Implement cov_diag and mean_and_cov_diag for FiniteGPs (#116)

Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractGPs"
 uuid = "99985d1d-32ba-4be9-9821-2ec096f28918"
 authors = ["JuliaGaussianProcesses Team"]
-version = "0.2.20"
+version = "0.2.21"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/AbstractGPs.jl

@@ -11,7 +11,7 @@ module AbstractGPs
 
     using KernelFunctions: ColVecs, RowVecs
 
-    export GP, mean, cov, std, cov_diag, mean_and_cov, marginals, rand,
+    export GP, mean, cov, std, cov_diag, mean_and_cov, mean_and_cov_diag, marginals,
         logpdf, elbo, dtc, posterior, approx_posterior, VFE, DTC, sampleplot,
         update_approx_posterior, LatentGP, ColVecs, RowVecs
 

src/abstract_gp/abstract_gp.jl

@@ -30,7 +30,7 @@ Statistics.cov(::AbstractGP, x::AbstractVector)
 
 Compute only the diagonal elements of `cov(f(x))`.
 """
-cov_diag(::AbstractGP, x::AbstractVector)
+cov_diag(::AbstractGP, ::AbstractVector)
 
 """
     cov(f::AbstractGP, x::AbstractVector, y::AbstractVector)
@@ -45,12 +45,12 @@ Statistics.cov(::AbstractGP, x::AbstractVector, y::AbstractVector)
 Compute both `mean(f(x))` and `cov(f(x))`. Sometimes more efficient than separately
 computation, particularly for posteriors.
 """
-mean_and_cov(::AbstractGP, ::AbstractVector)
+mean_and_cov(f::AbstractGP, x::AbstractVector) = (mean(f, x), cov(f, x))
 
 """
     mean_and_cov_diag(f::AbstractGP, x::AbstractVector)
 
 Compute both `mean(f(x))` and the diagonal elements of `cov(f(x))`. Sometimes more efficient
 than separately computation, particularly for posteriors.
 """
-mean_and_cov_diag(f::AbstractGP, x::AbstractVector)
+mean_and_cov_diag(f::AbstractGP, x::AbstractVector) = (mean(f, x), cov_diag(f, x))

src/abstract_gp/finite_gp.jl

@@ -92,6 +92,25 @@ true
 """
 Statistics.cov(f::FiniteGP) = cov(f.f, f.x) + f.Σy
 
+"""
+    cov_diag(f::FiniteGP)
+
+Compute only the diagonal elements of [`cov(f)`](@ref).
+
+# Examples
+
+```jldoctest
+julia> fx = GP(Matern52Kernel())(randn(10), 0.1);
+
+julia> cov_diag(fx) == diag(cov(fx))
+true
+```
+"""
+function cov_diag(f::FiniteGP)
+    Σy = f.Σy
+    return cov_diag(f.f, f.x) + view(Σy, diagind(Σy))
+end
+
 """
     mean_and_cov(f::FiniteGP)
 
@@ -111,6 +130,28 @@ function mean_and_cov(f::FiniteGP)
     return m, C + f.Σy
 end
 
+"""
+    mean_and_cov_diag(f::FiniteGP)
+
+Compute both `mean(f)` and the diagonal elements of `cov(f)`.
+
+Sometimes more efficient than computing them separately, particularly for posteriors.
+
+# Examples
+
+```jldoctest
+julia> fx = GP(SqExponentialKernel())(range(-3.0, 3.0; length=10), 0.1);
+
+julia> mean_and_cov_diag(fx) == (mean(fx), cov_diag(fx))
+true
+```
+"""
+function mean_and_cov_diag(f::FiniteGP)
+    m, c = mean_and_cov_diag(f.f, f.x)
+    Σy = f.Σy
+    return m, c + view(Σy, diagind(Σy))
+end
+
 """
     cov(fx::FiniteGP, gx::FiniteGP)
 
@@ -154,7 +195,10 @@ julia> std.(fs) == sqrt.(diag(cov(f(x))))
 true
 ```
 """
-marginals(f::FiniteGP) = Normal.(mean(f), sqrt.(cov_diag(f.f, f.x) .+ diag(f.Σy)))
+function marginals(f::FiniteGP)
+    m, c = mean_and_cov_diag(f)
+    return Normal.(m, sqrt.(c))
+end
 
 """
     rand(rng::AbstractRNG, f::FiniteGP, N::Int=1)

src/gp/gp.jl

@@ -62,8 +62,6 @@ GP(mean, kernel::Kernel) = GP(CustomMean(mean), kernel)
 GP(mean::Real, kernel::Kernel) = GP(ConstMean(mean), kernel)
 GP(kernel::Kernel) = GP(ZeroMean(), kernel)
 
-
-
 # AbstractGP interface implementation.
 
 Statistics.mean(f::GP, x::AbstractVector) = _map(f.mean, x)
@@ -73,7 +71,3 @@ Statistics.cov(f::GP, x::AbstractVector) = kernelmatrix(f.kernel, x)
 cov_diag(f::GP, x::AbstractVector) = kernelmatrix_diag(f.kernel, x)
 
 Statistics.cov(f::GP, x::AbstractVector, x′::AbstractVector) = kernelmatrix(f.kernel, x, x′)
-
-mean_and_cov(f::GP, x::AbstractVector) = (mean(f, x), cov(f, x))
-
-mean_and_cov_diag(f::GP, x::AbstractVector) = (mean(f, x), cov_diag(f, x))

test/abstract_gp/finite_gp.jl

@@ -20,6 +20,7 @@ end
         @test FiniteGP(f, Xmat, σ², obsdim=2) == FiniteGP(f, ColVecs(Xmat), σ²) 
         @test mean(fx) == mean(f, x)
         @test cov(fx) == cov(f, x)
+        @test cov_diag(fx) == diag(cov(fx))
         @test cov(fx, fx′) == cov(f, x, x′)
         @test mean.(marginals(fx)) == mean(f(x))
         @test var.(marginals(fx)) == cov_diag(f, x)
@@ -29,6 +30,11 @@ end
             @test m == mean(fx)
             @test C == cov(fx)
         end
+        let
+            m, c = mean_and_cov_diag(fx)
+            @test m == mean(fx)
+            @test c == cov_diag(fx)
+        end
     end
     @testset "rand (deterministic)" begin
         rng = MersenneTwister(123456)