Wct/fix bfr (#38)

* Fix BFR posterior

* Add comment for future readers

* Bump patch version

* Make tests more lightweight

* Fix typo

Author: Will Tebbutt

Project.toml

@@ -1,7 +1,7 @@
 name = "BayesianLinearRegressors"
 uuid = "f579363c-4606-5e5c-a623-c4549f609c4b"
 authors = ["Will Tebbutt <[email protected]>"]
-version = "0.3.7"
+version = "0.3.8"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

src/basis_function_regression.jl

@@ -36,6 +36,30 @@ struct BasisFunctionRegressor{Tblr<:BayesianLinearRegressor,Tϕ} <: AbstractGP
     ϕ::Tϕ
 end
 
-function (bfblr::BasisFunctionRegressor)(x::AbstractVector, args...)
-    return bfblr.blr(bfblr.ϕ(x), args...)
+const FiniteBFR = FiniteGP{<:BasisFunctionRegressor}
+
+_to_finite_blr(fx::FiniteBFR) = fx.f.blr(fx.f.ϕ(fx.x), fx.Σy)
+
+# All functionality below just implements the primary and secondary AbstractGPs APIs.
+# See AbstractGPs.jl's documentation for information regarding their semantics.
+
+AbstractGPs.mean(fx::FiniteBFR) = mean(_to_finite_blr(fx))
+
+AbstractGPs.cov(fx::FiniteBFR) = cov(_to_finite_blr(fx))
+
+AbstractGPs.var(fx::FiniteBFR) = var(_to_finite_blr(fx))
+
+AbstractGPs.mean_and_cov(fx::FiniteBFR) = mean_and_cov(_to_finite_blr(fx))
+
+AbstractGPs.mean_and_var(fx::FiniteBFR) = mean_and_var(_to_finite_blr(fx))
+
+function AbstractGPs.rand(rng::AbstractRNG, fx::FiniteBFR, samples::Int)
+    return rand(rng, _to_finite_blr(fx), samples)
+end
+
+AbstractGPs.logpdf(fx::FiniteBFR, y::AbstractVector{<:Real}) = logpdf(_to_finite_blr(fx), y)
+
+function AbstractGPs.posterior(fx::FiniteBFR, y::AbstractVector{<:Real})
+    f_post = posterior(_to_finite_blr(fx), y)
+    return BasisFunctionRegressor(f_post, fx.f.ϕ)
 end

test/basis_function_regression.jl

@@ -1,6 +1,6 @@
 @testset "basis_function_regression" begin
     @testset "basis_func_blr $Tx" for Tx in [Matrix, ColVecs, RowVecs]
-        @testset "consistency" begin
+        @testset "self-consistency" begin
             rng, N, D, samples = MersenneTwister(123456), 11, 2, 1_000_000
             X, f, Σy = generate_toy_problem(rng, N, D, Tx)
 
@@ -10,8 +10,24 @@
                 rng, f_bf(X, Σy)
             )
         end
+        @testset "consistency with BLR" begin
+            rng = MersenneTwister(123456)
+            N = 11
+            D = 2
+            X, f, Σy = generate_toy_problem(rng, N, D, Tx)
+            f_bf = BasisFunctionRegressor(f, ϕ)
+
+            # Compute logpdf using both the BLR and BFR. Should agree.
+            y = rand(rng, f_bf(X, Σy))
+            @test logpdf(f(ϕ(X), Σy), y) ≈ logpdf(f_bf(X, Σy), y)
+
+            # Check that posteriors agree.
+            f_bf_post = posterior(f_bf(X, Σy), y)
+            f_post = posterior(f(ϕ(X), Σy), y)
+            @test mean(f_bf_post(X)) ≈ mean(f_post(ϕ(X)))
+        end
         @testset "rand" begin
-            rng, N, D, samples = MersenneTwister(123456), 11, 2, 10_000_000
+            rng, N, D, samples = MersenneTwister(123456), 11, 2, 1_000_000
             X, f, Σy = generate_toy_problem(rng, N, D, Tx)
 
             f_bf = BasisFunctionRegressor(f, ϕ)
@@ -20,8 +36,8 @@
             Y = rand(rng, f_bf(X, Σy), samples)
             m_empirical = mean(Y; dims=2)
             Σ_empirical = (Y .- mean(Y; dims=2)) * (Y .- mean(Y; dims=2))' ./ samples
-            @test mean(f_bf(X, Σy)) ≈ m_empirical atol = 1e-3 rtol = 1e-3
-            @test cov(f_bf(X, Σy)) ≈ Σ_empirical atol = 1e-3 rtol = 1e-3
+            @test mean(f_bf(X, Σy)) ≈ m_empirical atol = 1e-2 rtol = 1e-2
+            @test cov(f_bf(X, Σy)) ≈ Σ_empirical atol = 1e-2 rtol = 1e-2
         end
     end
 end

test/bayesian_linear_regression.jl

@@ -9,15 +9,15 @@
             )
         end
         @testset "rand" begin
-            rng, N, D, samples = MersenneTwister(123456), 11, 3, 10_000_000
+            rng, N, D, samples = MersenneTwister(123456), 11, 3, 1_000_000
             X, f, Σy = generate_toy_problem(rng, N, D, Tx)
 
             # Roughly test the statistical properties of rand.
             Y = rand(rng, f(X, Σy), samples)
             m_empirical = mean(Y; dims=2)
             Σ_empirical = (Y .- mean(Y; dims=2)) * (Y .- mean(Y; dims=2))' ./ samples
-            @test mean(f(X, Σy)) ≈ m_empirical atol = 1e-3 rtol = 1e-3
-            @test cov(f(X, Σy)) ≈ Σ_empirical atol = 1e-3 rtol = 1e-3
+            @test mean(f(X, Σy)) ≈ m_empirical atol = 1e-2 rtol = 1e-2
+            @test cov(f(X, Σy)) ≈ Σ_empirical atol = 1e-2 rtol = 1e-2
 
             @testset "Zygote (everything dense)" begin
                 function rand_blr(X, A_Σy, mw, A_Λw)