Add SparseFiniteGP type and associated functionality

src/Stheno.jl

@@ -60,6 +60,9 @@ module Stheno
     # include(joinpath("composite", "gradient.jl"))
     # include(joinpath("composite", "integrate.jl"))
 
+    # approximate inference
+    include("approximate_inference.jl")
+
     # Various stuff for convenience.
     include(joinpath("util", "model.jl"))
     include(joinpath("util", "plotting.jl"))

src/approximate_inference.jl

@@ -0,0 +1,50 @@
+import Base: rand, length
+import Distributions: logpdf, AbstractMvNormal
+
+export mean, std, cov, marginals, rand, logpdf, elbo, dtc
+export SparseFiniteGP
+
+struct SparseFiniteGP{T1<:FiniteGP, T2<:FiniteGP} <: AbstractMvNormal
+    fobs::T1
+    finducing::T2
+end
+
+function SparseFiniteGP(f::AbstractGP, x, xu, σ=1e-18, σu=σ)
+    return SparseFiniteGP(f(x, σ), f(xu, σu))
+end
+
+# The below convenience constructor conflicts with the constructor for a FiniteGP with
+# diagonal covariance matrix (in src/abstract_gp.jl, line 34)
+# """
+#     (f::AbstractGP)(x::AbstractVector, xu::AbstractVector, σ=1e-18, σu=σ)
+#
+# Construct a `SparseFiniteGP` representing the projection of `f` at `x` with inducing points at `xu`.
+# """
+# (f::AbstractGP)(x::AbstractVector, xu::AbstractVector, σ=1e-18, σu=σ) = SparseFiniteGP(f, x, xu, σ, σu)
+#
+
+
+length(f::SparseFiniteGP) = length(f.fobs)
+mean(f::SparseFiniteGP) = mean(f.fobs)
+# in Base, trying to instantiate a large dense matrix through e.g.
+# inv(A::SparseArrays.AbstractSparseMatrixCSC) raises an error, since the resulting dense
+# matrix will often be enormous and use up all available memory. Doing something similar here.
+covariance_error = "The covariance matrix of a sparse GP can often be dense and can cause the computer to run out of memory. If you are sure you have enough memory, you can use `cov(f.fobs)`."
+cov(f::SparseFiniteGP) = error(covariance_error) #cov(f.fobs)
+# Not sure how to implement the following...
+# cov(fx::FiniteGP, gx::FiniteGP) = cov(fx.f, gx.f, fx.x, gx.x)
+marginals(f::SparseFiniteGP) = Normal.(mean(f), sqrt.(cov_diag(f.fobs.f, f.fobs.x) .+ diag(f.fobs.Σy)))
+
+rand(rng::AbstractRNG, f::SparseFiniteGP, N::Int) = rand(rng, f.fobs, N)
+rand(f::SparseFiniteGP, N::Int) = rand(Random.GLOBAL_RNG, f, N)
+rand(rng::AbstractRNG, f::SparseFiniteGP) = vec(rand(rng, f, 1))
+rand(f::SparseFiniteGP) = vec(rand(f, 1))
+
+function elbo(f::SparseFiniteGP, y::AV{<:Real})
+    elbo(f.fobs, y, f.finducing)
+end
+logpdf(f::SparseFiniteGP, y::AV{<:Real}) = elbo(f, y)
+logpdf(f::SparseFiniteGP, Y::AbstractMatrix{<:Real}) = map(y -> logpdf(f, y), eachcol(Y))
+
+PseudoObs(fxu::SparseFiniteGP, y) = PseudoObs(Obs(fxu.fobs, y), fxu.finducing)
+←(fxu::SparseFiniteGP, y) = PseudoObs(fxu, y)

test/gp/sparse_gp.jl

@@ -0,0 +1,53 @@
+using Stheno
+using Random
+using Test
+
+x = 0:0.1:10
+xu = 0:10
+σ = 1.0
+σu = 1e-3
+f = GP(Matern32(), GPC())
+
+@timedtestset "sparse gp" begin
+    x = 0:0.1:10
+    xu = 0:10
+    σ = 1.0
+    σu = 1e-3
+    covariance_error = "The covariance matrix of a sparse GP can often be dense and can cause the computer to run out of memory. If you are sure you have enough memory, you can use `cov(f.fobs)`."
+
+    @timedtestset "SparseFiniteGP Constructors" begin
+        f = GP(Matern32(), GPC())
+        @test SparseFiniteGP(f(x), f(xu)) == SparseFiniteGP(f, x,  xu)
+        # @test SparseFiniteGP(f(x), f(xu)) == f(x, xu)
+        @test SparseFiniteGP(f(x, σ), f(xu, σ)) == SparseFiniteGP(f, x, xu, σ)
+        # @test SparseFiniteGP(f(x, σ), f(xu, σ)) == f(x, xu, σ)
+        @test SparseFiniteGP(f(x, σ), f(xu, σu)) == SparseFiniteGP(f, x,  xu, σ, σu)
+        # @test SparseFiniteGP(f(x, σ), f(xu, σu)) == f(x, xu, σ, σu)
+    end
+
+    @timedtestset "SparseFiniteGP methods" begin
+        f = GP(Matern32(), GPC())
+        fx = f(x)
+        fxu = SparseFiniteGP(f, x, xu)
+        @test mean(fxu) == mean(fx)
+        @test rand(MersenneTwister(12345), fxu) == rand(MersenneTwister(12345), fx)
+        @test_throws ErrorException(covariance_error) cov(fxu)
+        @test cov(fxu.fobs) == cov(fx)
+    end
+
+    @timedtestset "SparseFiniteGP inference" begin
+        f = GP(Matern32(), GPC())
+        fx = f(x, σ)
+        fxu = SparseFiniteGP(f, x, xu, σ, σu)
+        y = rand(MersenneTwister(12345), fxu)
+
+        fpost1 = f | Stheno.PseudoObs(Obs(fxu.fobs, y), fxu.finducing)
+        fpost2 = f | (fxu ← y)
+        @test marginals(fpost1(x)) == marginals(fpost2(x))
+        @test elbo(fxu, y) == logpdf(fxu, y)
+        @test logpdf(fxu, y) == elbo(fxu.fobs, y, fxu.finducing)
+        yy = rand(fxu, 10)
+        logpdf(fx, yy)
+        logpdf(fxu, yy)
+    end
+end

test/runtests.jl

@@ -24,6 +24,7 @@ include("test_util.jl")
 
 @testset "Stheno" begin
 
+    include(joinpath("gp", "sparse_gp.jl"))
     println("util:")
     @timedtestset "util" begin
         include(joinpath("util", "zygote_rules.jl"))