Drawing function samples (#24)

* Implement drawing function samples

* ColVecs/RowVecs test

* Apply suggestions from code review

Co-authored-by: willtebbutt <[email protected]>

* Efficient multisample generation

* Fix test

* Non-mutating rand and autodiff tests

* Remove `Random.eltype` definition

* Add back in-place and clean up tests

* test Random.Sampler

* Patch bump

Co-authored-by: willtebbutt <[email protected]>

Author: rossviljoen

Project.toml

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

src/bayesian_linear_regression.jl

@@ -73,3 +73,34 @@ function __compute_inference_quantities(fx::FiniteBLR, y::AbstractVector{<:Real}
 
     return Uw, Bt, δy, logpdf_δy, Λεy
 end
+
+# Random function sample generation
+# Following the Random API: https://docs.julialang.org/en/v1/stdlib/Random/#Hooking-into-the-Random-API
+struct BLRFunctionSample{Tw<:AbstractVector}
+    w::Tw
+end
+
+(s::BLRFunctionSample)(X::AbstractMatrix{<:Real}) = X's.w
+(s::BLRFunctionSample)(X::ColVecs) = X.X's.w
+(s::BLRFunctionSample)(X::RowVecs) = X.X * s.w
+
+Random.Sampler(::Type{<:AbstractRNG}, blr::BayesianLinearRegressor, ::Random.Repetition) = blr
+
+function Random.rand(rng::AbstractRNG, blr::BayesianLinearRegressor)
+    w = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, size(blr.mw))
+    return BLRFunctionSample(w)
+end
+
+function Random.rand(rng::AbstractRNG, blr::BayesianLinearRegressor, dims::Dims)
+    ws = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, (only(size(blr.mw)), prod(dims)))
+    bs = [BLRFunctionSample(w) for w in eachcol(ws)]
+    return reshape(bs, dims)
+end
+
+function Random.rand!(rng::AbstractRNG, A::AbstractArray{<:BLRFunctionSample}, blr::BayesianLinearRegressor)
+    ws = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, (only(size(blr.mw)), prod(size(A))))
+    for i in LinearIndices(A)
+        @inbounds A[i] = BLRFunctionSample(ws[:,i])
+    end
+    return A
+end

test/bayesian_linear_regression.jl

@@ -121,4 +121,89 @@ end
             @test cov(f′(X′, Σy)) ≈ cov(f′2(X′, Σy))
         end
     end
+
+    @testset "sampling functions" begin
+        rng, N, D = MersenneTwister(123456), 11, 5
+        X, f, Σy = generate_toy_problem(rng, N, D)
+
+        g = rand(rng, f)
+        @test g(X) == g(X)  # check the sample doesn't change between evaluations
+
+        Xc = ColVecs(X)
+        Xr = RowVecs(X')
+        @test g(X) == g(Xc)
+        @test g(X) == g(Xr)
+
+        # test the Random interface
+        @test rand(rng, Random.Sampler(rng, f, Val(Inf))) isa BayesianLinearRegressors.BLRFunctionSample
+
+        samples1, samples2 = 10_000, 1000
+        samples = samples1 * samples2
+        gs = rand(rng, f, samples1, samples2)
+        @test size(gs) == (samples1, samples2)
+
+        # test statistical properties of the sampled functions
+        let
+            Y = reduce(hcat, map(h -> h(X), reshape(gs, :)))
+            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 + Σy atol = 1e-3 rtol = 1e-3
+        end
+
+        # test statistical properties of in-place rand
+        let
+            A = Array{BayesianLinearRegressors.BLRFunctionSample,2}(
+                undef,
+                samples1,
+                samples2,
+            )
+            A = rand!(rng, A, f)
+            Y = reduce(hcat, map(h -> h(X), reshape(gs, :)))
+            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 + Σy atol = 1e-3 rtol = 1e-3
+        end
+
+        @testset "Zygote (everything dense)" begin
+            function test_rand_funcs_adjoints(sample_function)
+                rng, N, D = MersenneTwister(123456), 11, 5
+                X, f, _ = generate_toy_problem(rng, N, D)
+                mw, A_Λw = f.mw, 0.1 .* randn(rng, D, D)
+
+                # Run the model forwards and check that output agrees with non-Zygote output.
+                z, back = Zygote.pullback(sample_function, X, mw, A_Λw)
+                @test z == sample_function(X, mw, A_Λw)
+
+                # Compute adjoints using Zygote.
+                z̄ = randn(rng, size(z))
+                dX, dmw, dA_Λw = back(z̄)
+
+                # Verify adjoints via finite differencing.
+                fdm = central_fdm(5, 1)
+                @test dX ≈ first(j′vp(fdm, X -> sample_function(X, mw, A_Λw), z̄, X))
+                @test dmw ≈ first(j′vp(fdm, mw -> sample_function(X, mw, A_Λw), z̄, mw))
+                @test dA_Λw ≈
+                      first(j′vp(fdm, A_Λw -> sample_function(X, mw, A_Λw), z̄, A_Λw))
+            end
+
+            function rand_funcs_single(X, mw, A_Λw)
+                Λw = Symmetric(A_Λw * A_Λw' + I)
+                f = BayesianLinearRegressor(mw, Λw)
+                g = rand(MersenneTwister(123456), f)
+                return g(X)
+            end
+
+            function rand_funcs_multi(X, mw, A_Λw)
+                Λw = Symmetric(A_Λw * A_Λw' + I)
+                f = BayesianLinearRegressor(mw, Λw)
+                gs = rand(MersenneTwister(123456), f, 1, 1)
+                return reduce(hcat, map(h -> h(X), reshape(gs, :)))
+            end
+
+            test_rand_funcs_adjoints(rand_funcs_single)
+            test_rand_funcs_adjoints(rand_funcs_multi)
+        end
+    end
 end