Additive GP affine transformation (#229)

* Add additive_gp transformation

* Bump patch

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "Stheno"
 uuid = "8188c328-b5d6-583d-959b-9690869a5511"
-version = "0.8.0"
+version = "0.8.1"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

docs/src/api.md

@@ -25,4 +25,5 @@ stretch
 periodic
 shift
 select
+additive_gp
 ```

src/Stheno.jl

@@ -36,12 +36,13 @@ module Stheno
     include(joinpath("affine_transformations", "addition.jl"))
     include(joinpath("affine_transformations", "compose.jl"))
     include(joinpath("affine_transformations", "product.jl"))
+    include(joinpath("affine_transformations", "additive_gp.jl"))
 
     # AbstractGP subtype which groups together other AbstractGP subtypes.
     include("gaussian_process_probabilistic_programme.jl")
 
     export atomic, BlockData, GPC, GPPPInput, @gppp
-    export ∘, select, stretch, periodic, shift
+    export ∘, select, stretch, periodic, shift, additive_gp
     export SparseFiniteGP
 
 end # module

src/affine_transformations/additive_gp.jl

@@ -0,0 +1,29 @@
+"""
+    additive_gp(fs)
+
+Produces the GP given by
+```julia
+sum(fs[1](x[1]) + fs[2](x[2]) + ... + fs[D](x[D]))
+```
+Requires that `length(fs)` is the same as the dimension of the inputs to be used.
+"""
+additive_gp(fs) = additive_gp(fs, 1:length(fs))
+
+"""
+    additive_gp(fs, indices)
+
+`fs` should be a collection of GPs, and `indices` a collection of collections of integer
+indices.
+For example, `indices` might be something like `[1:2, 3, 4:6]`, in which case `fs` would
+need to comprise exactly three elements. In general, this functions requires that
+`length(fs) == length(indices)`.
+
+Produces the GP given by
+```julia
+sum(fs[1](x[indices[1]]) + fs[2](x[indices[2]]) + ... + fs[D](x[indices[D]]))
+```
+"""
+function additive_gp(fs, indices)
+    fs_projected = map((f, idx) -> f ∘ Select(idx), fs, indices)
+    return sum(fs_projected)
+end

test/affine_transformations/additive_gp.jl

@@ -0,0 +1,43 @@
+@testset "additive_gp" begin
+    @testset "arbitrary indices" begin
+        f = @gppp let
+            f1 = GP(SEKernel())
+            f2 = GP(SEKernel())
+            f3 = additive_gp((f1, f2), [2:3, 1])
+        end
+        x_raw = ColVecs(randn(3, 7))
+        x1_raw = ColVecs(x_raw.X[2:3, :])
+        x2_raw = x_raw.X[1, :]
+        x = vcat(
+            GPPPInput(:f1, x1_raw),
+            GPPPInput(:f2, x2_raw),
+            GPPPInput(:f3, x_raw),
+        )
+        y = rand(f(x, 1e-9))
+
+        @test y[1:7] + y[8:14] ≈ y[15:21] rtol=1e-3
+
+        z = GPPPInput(:f3, ColVecs(3 * randn(3, 4)))
+        test_internal_abstractgps_interface(MersenneTwister(123456), f, x, z; jitter=1e-15)
+    end
+    @testset "regular indices" begin
+        f = @gppp let
+            f1 = GP(SEKernel())
+            f2 = GP(Matern52Kernel())
+            f3 = additive_gp((f1, f2))
+        end
+        x_raw = ColVecs(randn(2, 7))
+        x1_raw = x_raw.X[1, :]
+        x2_raw = x_raw.X[2, :]
+        x = vcat(
+            GPPPInput(:f1, x1_raw),
+            GPPPInput(:f2, x2_raw),
+            GPPPInput(:f3, x_raw),
+        )
+        y = rand(f(x, 1e-9))
+        @test y[1:7] + y[8:14] ≈ y[15:21] rtol=1e-3
+
+        z = GPPPInput(:f3, ColVecs(3 * randn(2, 4)))
+        test_internal_abstractgps_interface(MersenneTwister(123456), f, x, z; jitter=1e-15)
+    end
+end

test/runtests.jl

@@ -61,6 +61,7 @@ include("test_util.jl")
         include(joinpath("affine_transformations", "addition.jl"))
         include(joinpath("affine_transformations", "compose.jl"))
         include(joinpath("affine_transformations", "product.jl"))
+        include(joinpath("affine_transformations", "additive_gp.jl"))
     end
 
     include("gaussian_process_probabilistic_programme.jl")