Bayesian linear Gaussian models

lingauss.jl

using LinearAlgebra, Test, Random, Statistics
outer(x) = x*x'
lchol(x) = cholesky(Symmetric(x)).L
Random.seed!(1)

# dimensions
n = 10 # observed
p = 5 # latent

# parameters
A0 = 0.2randn(p, p) + I
B = randn(n, p)
A = 0.2randn(n, n) + I
Σ0 = outer(A0)
Γ0 = inv(Σ0)
Σ = outer(A)
Γ = inv(Σ)

# Sample
x = A0*randn(p)
y = B*x + A*randn(n)

# Many samples
N = 10000
X = A0*randn(p, N)
Y = B*X + A*randn(n, N)
@test norm(cov(X') - Σ0) < 2sum(eigvals(Σ0))/sqrt(N) # Wishart(N-1, A0*A0')
@test norm(cov(Y' - (B*X)') - Σ) < 2sum(eigvals(Σ))/sqrt(N) # Wishart(N-1, A0*A0')
ΣY = outer(B*A0) + outer(A)
@test norm(cov(Y') - ΣY) < 2sum(eigvals(ΣY))/sqrt(N) # Wishart(N-1, A0*A0')

# cholesky of Σ is qr of A
L0 = cholesky(Σ0).L 
L = cholesky(Σ).L 
LA = qr(A').R'
@test L ≈ LA*Diagonal(sign.(diag(LA)))

# Joint distribution and cholesky
LL = [L0 zeros(p,n); B*L0 L]
ΣΣ = LL*LL'
@test ΣΣ[p+1:end, 1:p] ≈ B*Σ0
@test ΣΣ[p+1:end, p+1:end] ≈ ΣY

# A block-triangular factor
AA = [A0 zeros(p,n); B*A0 A]
@test ΣΣ ≈ AA*AA'

# empirical marginal Y
@test norm(cov(Y' - (B*X)') - Σ) < 2sum(eigvals(Σ))/sqrt(N) # Wishart(N-1, A0*A0')

# posterior with and without woodbury
Σpost = Σ0 - Σ0*B'*inv(outer(B*A0) + outer(A))*B*Σ0
@test Σpost ≈ inv(Γ0 + B'*Γ*B)

# block cholesky step of a 2x2 matrix
L11 = lchol(outer(B*A0) + outer(A))
U = Σ0*B'
L21 = U/L11'   
L22 = lchol(Σ0 - L21*L21')
@test Σ0 - L21*L21' ≈ Σpost

# qr (or cholesky) on product
LAA = qr(AA').R'
@test LL ≈ LAA*Diagonal(sign.(diag(LAA)))

# using a qr/cholesky to "reorder" and compute conditional variance
AAb = [A B*A0; zeros(p,n) A0]
LAAb = qr(AAb').R'
@test outer(LAAb[n+1:end,n+1:end]) ≈ Σpost