Merge pull request #25 from mschauer/new2024

add marginal and conditional kernel for gauss

src/Mitosis.jl

@@ -35,9 +35,10 @@ import StatsBase.sample
 include("mt.jl")
 
 
-export Gaussian, Copy, fuse, weighted
+export Gaussian, conditional, marginal, Copy, fuse, weighted
 export Traced, BFFG, left′, right′, forward, backward, backwardfilter, forwardsampler
 export BF, density, logdensity, ⊕, kernel, correct, Kernel, WGaussian, Gaussian, ConstantMap, AffineMap, LinearMap, GaussKernel
+export likelihood
 
 function independent_sum
 end
@@ -108,6 +109,6 @@ include("wgaussian.jl")
 include("markov.jl")
 include("rules.jl")
 include("regression.jl")
-
+include("bayes.jl")
 
 end # module

src/bayes.jl

@@ -0,0 +1,18 @@
+
+function conditional(p::Gaussian{(:μ, :Σ)}, A, B)
+    Z = p.Σ[A,B]*inv(p.Σ[B,B])
+    β = p.μ[A] - Z*p.μ[B]
+    Σ = p.Σ[A,A] - Z*p.Σ[B,A]
+    kernel(Gaussian; μ=AffineMap(Z, β), Σ=ConstantMap(Σ))
+end
+
+function marginal(p::Gaussian{(:μ, :Σ)}, A)
+    Gaussian{(:μ, :Σ)}(p.μ[A], p.Σ[A,A])
+end
+
+function likelihood(k::AffineGaussianKernel, obs)
+    q = backward(BFFG(), k, obs; unfused=true)[2].y
+    F, Γ, c =  params(q)
+    c0 = Mitosis.logdensity0(Gaussian{(:F,:Γ)}(F, Γ))
+    WGaussian{(:F,:Γ, :c)}(F, Γ, c - c0)
+end

src/gauss.jl

@@ -119,9 +119,11 @@ Base.:*(M, p::Gaussian{P}) where {P} = Gaussian{P}(M * mean(p), Σ = M * cov(p)
 
 """
     conditional(p::Gaussian, A, B, xB)
+    conditional(p::Gaussian, A, B)
 
 Conditional distribution of `X[i for i in A]` given
-`X[i for i in B] == xB` if ``X ~ P``.
+`X[i for i in B] == xB` if ``X ~ P``. The version without
+the argument `xB` returns a kernel mapping `xB` to the conditional.
 """
 function conditional(p::Gaussian{(:μ, :Σ)}, A, B, xB)
     Z = p.Σ[A,B]*inv(p.Σ[B,B])

test/gauss.jl

@@ -24,4 +24,21 @@ p3 = convert(Gaussian{(:μ,:Σ)}, p2)
         @test 10/sqrt(K) > norm(m - mean(rand(G) for x in 1:K))
         @test 10/sqrt(K) > norm(Q - cov([rand(G) for x in 1:K]))
     end
+end
+
+@testset "conditional" begin
+    using Mitosis
+    d = 5
+    d1 = 2
+    μ = randn(d)
+    A = randn(d,d)
+    q = Gaussian{(:μ,:Σ)}(μ, A*A')
+    π = marginal(q, 1:d1)
+    k = conditional(q, d1+1:d, 1:d1)
+    @test k(π) ≈ marginal(q, d1+1:d)
+    x = rand(π)
+    obs = rand(k(x))
+
+    @test logdensity(k(x), obs) ≈ logdensity(likelihood(k, obs), x) 
+
 end