Allow AbstractGPs in WrappedGP (#217)

* Bump patch

* Allow AbstractGP in WrappedGP

* Test nested GPPP

* Check WrappedGP with fresh AbstractGP

* Remove redundant code

* Remove redudant code

Author: willtebbutt

Project.toml

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

src/Stheno.jl

@@ -37,7 +37,6 @@ module Stheno
     include(joinpath("util", "zygote_rules.jl"))
     include(joinpath("util", "covariance_matrices.jl"))
     include(joinpath("util", "block_arrays", "dense.jl"))
-    include(joinpath("util", "block_arrays", "diagonal.jl"))
     include(joinpath("util", "abstract_data_set.jl"))
     include(joinpath("util", "proper_type_piracy.jl"))
 

src/gp/gp.jl

@@ -12,14 +12,14 @@ struct WrappedGP{Tgp<:AbstractGP} <: SthenoAbstractGP
     gp::Tgp
     n::Int
     gpc::GPC
-    function WrappedGP{Tgp}(gp::Tgp, gpc::GPC) where {Tgp<:GP}
+    function WrappedGP{Tgp}(gp::Tgp, gpc::GPC) where {Tgp<:AbstractGP}
         wgp = new{Tgp}(gp, next_index(gpc), gpc)
         gpc.n += 1
         return wgp
     end
 end
 
-wrap(gp::Tgp, gpc::GPC) where {Tgp<:GP} = WrappedGP{Tgp}(gp, gpc)
+wrap(gp::Tgp, gpc::GPC) where {Tgp<:AbstractGP} = WrappedGP{Tgp}(gp, gpc)
 
 mean(f::WrappedGP, x::AbstractVector) = mean(f.gp, x)
 

src/util/block_arrays/diagonal.jl

@@ -82,8 +82,7 @@
             GPPPInput(:f1, randn(4)),
         ),
     ]
-        rng = MersenneTwister(123456)
-        AbstractGPs.TestUtils.test_internal_abstractgps_interface(rng, f, x0, x1)
+        test_internal_abstractgps_interface(MersenneTwister(123456), f, x0, x1)
     end
 
     @timedtestset "gppp macro" begin
@@ -103,4 +102,20 @@
         y = rand(f(x, s))
         Zygote.gradient((x, y, f, s) -> logpdf(f(x, s), y), x, y, f, s)
     end
+
+    # Check that we can use one GPPP inside another.
+    @timedtestset "nested gppp" begin
+
+        gpc_outer = GPC()
+        f1_outer = Stheno.wrap(f, gpc_outer)
+        f2_outer = 5 * f1_outer
+        f_outer = Stheno.GPPP((f1=f1_outer, f2=f2_outer), gpc_outer)
+
+        x0 = GPPPInput(:f1, randn(5))
+        x1 = GPPPInput(:f2, randn(4))
+        x0_outer = GPPPInput(:f1, x0)
+        x1_outer = GPPPInput(:f2, x1)
+        rng = MersenneTwister(123456)
+        test_internal_abstractgps_interface(rng, f_outer, x0_outer, x1_outer)
+    end
 end

test/gaussian_process_probabilistic_programme.jl

@@ -1,3 +1,5 @@
+struct ToyAbstractGP <: AbstractGP end
+
 @timedtestset "gp" begin
 
     # Ensure that basic functionality works as expected.
@@ -33,4 +35,8 @@
 
         @test cov(f1, f1, x′, x) ≈ cov(f1, f1, x, x′)'
     end
+
+    @timedtestset "wrapped AbstractGP" begin
+        wrap(ToyAbstractGP(), GPC())
+    end
 end

test/gp/gp.jl

@@ -24,7 +24,6 @@ using Stheno:
     GPC,
     AV,
     FiniteGP,
-    block_diagonal,
     AbstractGP,
     BlockData,
     blocks,
@@ -36,10 +35,9 @@ using Stheno:
     diag_At_B,
     diag_Xt_invA_X,
     diag_Xt_invA_Y,
-    block_diagonal,
-    BlockDiagonal,
     blocksizes
 
+using Stheno.AbstractGPs.TestUtils: test_internal_abstractgps_interface
 using Stheno.AbstractGPs.Distributions: MvNormal
 using FiniteDifferences: j′vp
 
@@ -60,7 +58,6 @@ include("test_util.jl")
         @testset "block_arrays" begin
             include(joinpath("util", "block_arrays", "test_util.jl"))
             include(joinpath("util", "block_arrays", "dense.jl"))
-            include(joinpath("util", "block_arrays", "diagonal.jl"))
         end
         include(joinpath("util", "abstract_data_set.jl"))
     end