Skip to content

Commit 50cfa22

Browse files
bvdmitribvdmitri
authored
Merge pull request #497 from ReactiveBayes/mv_normal_mean_scale_matrix_precision
Mv normal mean scale matrix precision
2 parents a67b643 + 9b6b53b commit 50cfa22

14 files changed

Lines changed: 532 additions & 1 deletion

File tree

‎src/nodes/predefined.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ include("predefined/mv_normal_mean_covariance.jl")
88
include("predefined/mv_normal_mean_precision.jl")
99
include("predefined/mv_normal_mean_scale_precision.jl")
1010
include("predefined/mv_normal_weighted_mean_precision.jl")
11+
include("predefined/mv_normal_mean_scale_matrix_precision.jl")
1112
include("predefined/gamma.jl")
1213
include("predefined/gamma_inverse.jl")
1314
include("predefined/gamma_shape_rate.jl")

‎src/nodes/predefined/mv_normal_mean_scale_matrix_precision.jl‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
export MvNormalMeanScaleMatrixPrecision
2+
3+
import StatsFuns: log2π
4+
5+
struct MvNormalMeanScaleMatrixPrecision{T <: Real, M <: AbstractVector{T}, H <: AbstractMatrix{T}} <: AbstractMvNormal
6+
μ::M
7+
γ::T
8+
G::H
9+
end
10+
11+
@node MvNormalMeanScaleMatrixPrecision Stochastic [out, (μ, aliases = [mean]), (γ, aliases = [scale]), (G, aliases = [matrix])]
12+
13+
# default method for mean-field assumption
14+
@average_energy MvNormalMeanScaleMatrixPrecision (q_out::Any, q_μ::Any, q_γ::Any, q_G::Any) = begin
15+
dim = ndims(q_out)
16+
17+
m_mean, v_mean = mean_cov(q_μ)
18+
m_out, v_out = mean_cov(q_out)
19+
m_Λ = mean(q_γ) * mean(q_G)
20+
21+
result = zero(promote_samplefloattype(q_out, q_μ, q_γ, q_G))
22+
result += dim * log2π
23+
result -= dim * mean(log, q_γ)
24+
result -= mean(logdet, q_G)
25+
@inbounds for k1 in 1:dim, k2 in 1:dim
26+
# optimize trace operation (indices can be interchanges because of symmetry)
27+
result += m_Λ[k1, k2] * (v_out[k1, k2] + v_mean[k1, k2] + (m_out[k2] - m_mean[k2]) * (m_out[k1] - m_mean[k1]))
28+
end
29+
result /= 2
30+
31+
return result
32+
end
33+
34+
# default method for structured mean-field assumption
35+
@average_energy MvNormalMeanScaleMatrixPrecision (q_out_μ::Any, q_γ::Any, q_G::Any) = begin
36+
dim = div(ndims(q_out_μ), 2)
37+
38+
m, V = mean_cov(q_out_μ)
39+
m_out = m[1:dim]
40+
m_mean = m[(dim + 1):end]
41+
# slice V into blocks according to dim-lengths of out and μ
42+
v_out = V[1:dim, 1:dim]
43+
v_mean = V[(dim + 1):end, (dim + 1):end]
44+
v_out_mean = V[1:dim, (dim + 1):end]
45+
v_mean_out = V[(dim + 1):end, 1:dim]
46+
m_Λ = mean(q_γ) * mean(q_G)
47+
48+
result = zero(promote_samplefloattype(q_out_μ, q_γ, q_G))
49+
result += dim * log2π
50+
result -= dim * mean(log, q_γ)
51+
result -= mean(logdet, q_G)
52+
@inbounds for k1 in 1:dim, k2 in 1:dim
53+
# optimize trace operation (indices can be interchanges because of symmetry)
54+
result += m_Λ[k1, k2] * (v_out[k1, k2] - v_out_mean[k1, k2] - v_mean_out[k1, k2] + v_mean[k1, k2] + (m_out[k2] - m_mean[k2]) * (m_out[k1] - m_mean[k1]))
55+
end
56+
result /= 2
57+
58+
return result
59+
end

‎src/rules/mv_normal_mean_scale_matrix_precision/marginals.jl‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
export marginalrule
2+
3+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ_γ_G) (m_out::MultivariateNormalDistributionsFamily, m_μ::PointMass, m_γ::PointMass, m_G::PointMass) = begin
4+
return (out = prod(ClosedProd(), MvNormalMeanPrecision(mean(m_μ), mean(m_γ) * mean(m_G)), m_out), μ = m_μ, γ = m_γ, G = m_G)
5+
end
6+
7+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ_γ_G) (m_out::PointMass, m_μ::MultivariateNormalDistributionsFamily, m_γ::PointMass, m_G::PointMass) = begin
8+
return (out = m_out, μ = prod(ClosedProd(), m_μ, MvNormalMeanPrecision(mean(m_out), mean(m_γ) * mean(m_G))), γ = m_γ, G = m_G)
9+
end
10+
11+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ_γ_G) (
12+
m_out::MultivariateNormalDistributionsFamily, m_μ::MultivariateNormalDistributionsFamily, m_γ::PointMass, m_G::PointMass
13+
) = begin
14+
xi_y, W_y = weightedmean_precision(m_out)
15+
xi_m, W_m = weightedmean_precision(m_μ)
16+
17+
W_bar = mean(m_γ) * mean(m_G)
18+
19+
T = promote_samplefloattype(m_out, m_μ, m_γ, m_G)
20+
d = length(xi_y)
21+
Λ = Matrix{T}(undef, (2 * d, 2 * d))
22+
@inbounds for k2 in 1:d
23+
@inbounds for k1 in 1:d
24+
tmp1 = W_bar[k1, k2]
25+
tmp2 = -tmp1
26+
k1d = k1 + d
27+
k2d = k2 + d
28+
Λ[k1, k2] = tmp1 + W_y[k1, k2]
29+
Λ[k1d, k2] = tmp2
30+
Λ[k1, k2d] = tmp2
31+
Λ[k1d, k2d] = tmp1 + W_m[k1, k2]
32+
end
33+
end
34+
35+
ξ = [xi_y; xi_m]
36+
37+
return (out_μ = MvNormalWeightedMeanPrecision(ξ, Λ), γ = m_γ, G = m_G)
38+
end
39+
40+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ) (m_out::MultivariateNormalDistributionsFamily, m_μ::MultivariateNormalDistributionsFamily, q_γ::Any, q_G::Any) = begin
41+
xi_y, W_y = weightedmean_precision(m_out)
42+
xi_m, W_m = weightedmean_precision(m_μ)
43+
44+
W_bar = mean(q_γ) * mean(q_G)
45+
46+
T = promote_samplefloattype(m_out, m_μ, q_γ, q_G)
47+
d = length(xi_y)
48+
Λ = Matrix{T}(undef, (2 * d, 2 * d))
49+
@inbounds for k2 in 1:d
50+
@inbounds for k1 in 1:d
51+
tmp1 = W_bar[k1, k2]
52+
tmp2 = -tmp1
53+
k1d = k1 + d
54+
k2d = k2 + d
55+
Λ[k1, k2] = tmp1 + W_y[k1, k2]
56+
Λ[k1d, k2] = tmp2
57+
Λ[k1, k2d] = tmp2
58+
Λ[k1d, k2d] = tmp1 + W_m[k1, k2]
59+
end
60+
end
61+
62+
ξ = [xi_y; xi_m]
63+
64+
return MvNormalWeightedMeanPrecision(ξ, Λ)
65+
end
66+
67+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ) (m_out::PointMass, m_μ::MultivariateNormalDistributionsFamily, q_γ::Any, q_G::Any) = begin
68+
return (out = m_out, μ = prod(ClosedProd(), MvNormalMeanPrecision(mean(m_out), mean(q_γ) * mean(q_G)), m_μ))
69+
end
70+
71+
@marginalrule MvNormalMeanScaleMatrixPrecision(:out_μ) (m_out::MultivariateNormalDistributionsFamily, m_μ::PointMass, q_γ::Any, q_G::Any) = begin
72+
return (out = prod(ClosedProd(), MvNormalMeanPrecision(mean(m_μ), mean(q_γ) * mean(q_G)), m_out), μ = m_μ)
73+
end

‎src/rules/mv_normal_mean_scale_matrix_precision/matrix.jl‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Variational #
2+
# --------------------------------- #
3+
@rule MvNormalMeanScaleMatrixPrecision(:G, Marginalisation) (q_out::Any, q_μ::Any, q_γ::Any) = begin
4+
m_out, v_out = mean_cov(q_out)
5+
m_mean, v_mean = mean_cov(q_μ)
6+
γ_bar = mean(q_γ)
7+
8+
n_G = ndims(q_μ) + 2
9+
V_G = inv(γ_bar * (v_out + v_mean + (m_out - m_mean) * (m_out - m_mean)'))
10+
11+
return Wishart(convert(eltype(V_G), n_G), V_G)
12+
end
13+
14+
@rule MvNormalMeanScaleMatrixPrecision(:G, Marginalisation) (q_out_μ::Any, q_γ::Any) = begin
15+
m_out_μ, v_out_μ = mean_cov(q_out_μ)
16+
γ_bar = mean(q_γ)
17+
18+
d = div(ndims(q_out_μ), 2)
19+
20+
n_G = d + 2
21+
22+
mdiff = @views m_out_μ[1:d] - m_out_μ[(d + 1):end]
23+
vdiff = @views v_out_μ[1:d, 1:d] - v_out_μ[1:d, (d + 1):end] - v_out_μ[(d + 1):end, 1:d] + v_out_μ[(d + 1):end, (d + 1):end]
24+
V_G = inv((vdiff + mdiff * mdiff') * γ_bar)
25+
26+
return Wishart(convert(eltype(V_G), n_G), V_G)
27+
end

‎src/rules/mv_normal_mean_scale_matrix_precision/mean.jl‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Variational #
2+
# --------------------------------- #
3+
@rule MvNormalMeanScaleMatrixPrecision(:μ, Marginalisation) (q_out::Any, q_γ::Any, q_G::Any) = MvNormalMeanPrecision(mean(q_out), mean(q_γ) * mean(q_G))
4+
5+
@rule MvNormalMeanScaleMatrixPrecision(:μ, Marginalisation) (m_out::MultivariateNormalDistributionsFamily, q_γ::Any, q_G::Any) = begin
6+
out_bar, Λ_out = mean_precision(m_out)
7+
Λ_f = mean(q_γ) * mean(q_G)
8+
9+
# Step 1: form M = Λ_out + Λ_f
10+
M = Λ_out + Λ_f
11+
12+
# Step 2: factorize M (Cholesky since M is PD)
13+
F = fastcholesky(M)
14+
15+
# Step 3: solve M X = Λ_out (multiple RHS solve)
16+
X = F \ Λ_out
17+
18+
# Step 4: form Λ_μ
19+
Λ_μ = Λ_out - Λ_out * X
20+
21+
return MvNormalMeanPrecision(out_bar, Λ_μ)
22+
end

‎src/rules/mv_normal_mean_scale_matrix_precision/out.jl‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Variational #
2+
# --------------------------------- #
3+
@rule MvNormalMeanScaleMatrixPrecision(:out, Marginalisation) (q_μ::Any, q_γ::Any, q_G::Any) = MvNormalMeanPrecision(mean(q_μ), mean(q_γ) * mean(q_G))
4+
5+
@rule MvNormalMeanScaleMatrixPrecision(:out, Marginalisation) (m_μ::MultivariateNormalDistributionsFamily, q_γ::Any, q_G::Any) = begin
6+
μ_bar, Λ_μ = mean_precision(m_μ)
7+
Λ_f = mean(q_γ) * mean(q_G)
8+
9+
# Step 1: form M = Λ_μ + Λ_f
10+
M = Λ_μ + Λ_f
11+
12+
# Step 2: factorize M (Cholesky since M is PD)
13+
F = fastcholesky(M)
14+
15+
# Step 3: solve M X = Λ_μ (multiple RHS solve)
16+
X = F \ Λ_μ
17+
18+
# Step 4: form Λ_out
19+
Λ_out = Λ_μ - Λ_μ * X
20+
21+
return MvNormalMeanPrecision(μ_bar, Λ_out)
22+
end

‎src/rules/mv_normal_mean_scale_matrix_precision/precision.jl‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Variational #
2+
# --------------------------------- #
3+
@rule MvNormalMeanScaleMatrixPrecision(:γ, Marginalisation) (q_out::Any, q_μ::Any, q_G::Any) = begin
4+
m_out, v_out = mean_cov(q_out)
5+
m_mean, v_mean = mean_cov(q_μ)
6+
G_bar = mean(q_G)
7+
8+
α = ndims(q_μ) / 2 + 1
9+
β = tr(G_bar * (v_mean + v_out + (m_out - m_mean) * (m_out - m_mean)')) / 2
10+
11+
return GammaShapeRate(convert(eltype(β), α), β)
12+
end
13+
14+
@rule MvNormalMeanScaleMatrixPrecision(:γ, Marginalisation) (q_out_μ::Any, q_G::Any) = begin
15+
m_out_μ, v_out_μ = mean_cov(q_out_μ)
16+
G_bar = mean(q_G)
17+
18+
d = div(ndims(q_out_μ), 2)
19+
20+
α = d / 2 + 1
21+
22+
mdiff = @views m_out_μ[1:d] - m_out_μ[(d + 1):end]
23+
vdiff = @views v_out_μ[1:d, 1:d] - v_out_μ[1:d, (d + 1):end] - v_out_μ[(d + 1):end, 1:d] + v_out_μ[(d + 1):end, (d + 1):end]
24+
β = tr(G_bar * (vdiff + mdiff * mdiff')) / 2
25+
26+
return GammaShapeRate(convert(eltype(β), α), β)
27+
end

‎src/rules/mv_normal_mean_scale_precision/marginals.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ end
3232

3333
ξ = [xi_y; xi_m]
3434

35-
return (out_μ = MvNormalWeightedMeanPrecision(ξ, Λ), Λ = m_Λ)
35+
return (out_μ = MvNormalWeightedMeanPrecision(ξ, Λ), γ = m_γ)
3636
end
3737

3838
@marginalrule MvNormalMeanScalePrecision(:out_μ) (m_out::MultivariateNormalDistributionsFamily, m_μ::MultivariateNormalDistributionsFamily, q_γ::Any) = begin

‎src/rules/predefined.jl‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ include("mv_normal_mean_precision/mean.jl")
6666
include("mv_normal_mean_precision/precision.jl")
6767
include("mv_normal_mean_precision/marginals.jl")
6868

69+
include("mv_normal_mean_scale_matrix_precision/out.jl")
70+
include("mv_normal_mean_scale_matrix_precision/mean.jl")
71+
include("mv_normal_mean_scale_matrix_precision/matrix.jl")
72+
include("mv_normal_mean_scale_matrix_precision/precision.jl")
73+
include("mv_normal_mean_scale_matrix_precision/marginals.jl")
74+
6975
include("mv_normal_mean_scale_precision/out.jl")
7076
include("mv_normal_mean_scale_precision/mean.jl")
7177
include("mv_normal_mean_scale_precision/precision.jl")

‎test/nodes/predefined/mv_normal_mean_scale_matrix_precision_tests.jl‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
2+
@testitem "MvNormalMeanScaleMatrixPrecision" begin
3+
using ReactiveMP, Random, BayesBase, ExponentialFamily, LinearAlgebra
4+
5+
@testset "AverageEnergy" begin
6+
begin
7+
q_out = PointMass([1.0, 1.0])
8+
q_μ = MvNormalMeanPrecision([1.0, 1.0], [1.0 0.0; 0.0 1.0])
9+
q_γ = GammaShapeRate(1.0, 1.0)
10+
q_G = Wishart(ndims(q_out) + 2, [1.0 0.0; 0.0 1.0])
11+
# m_μ, Cov_μ = mean_cov(q_μ)
12+
# m_out, Cov_out = mean_cov(q_out)
13+
# AE = - div(ndims(q_μ),2) * mean(log,q_γ) - 0.5*mean(logdet, q_G) + div(ndims(q_μ),2)*log(2pi) + 0.5*tr(mean(q_γ)*mean(q_G)*( Cov_out + Cov_μ + (m_out - m_μ)*(m_out - m_μ)'))
14+
15+
for N in (MvNormalMeanPrecision, MvNormalMeanCovariance, MvNormalWeightedMeanPrecision), g in (Gamma,), M in (Wishart,)
16+
marginals = (
17+
Marginal(q_out, false, false, nothing),
18+
Marginal(convert(N, q_μ), false, false, nothing),
19+
Marginal(convert(g, q_γ), false, false, nothing),
20+
Marginal(convert(M, q_G), false, false, nothing)
21+
)
22+
@test score(AverageEnergy(), MvNormalMeanScaleMatrixPrecision, Val{(:out, :μ, :γ, :G)}(), marginals, nothing) ≈ 5.49230839621241
23+
end
24+
end
25+
26+
begin
27+
q_out = MvNormalMeanPrecision([0.8625589157256603, 0.6694783342639599], [1.0014322413749484 0.7989099036521625; 0.7989099036521625 1.0976639268696966])
28+
q_μ = MvNormalMeanPrecision([0.9334416739853251, 0.38318522701093105], [0.21867945696266933 0.5704895781120056; 0.5704895781120056 1.5321190185800933])
29+
q_γ = Gamma(2.0, 1.0)
30+
q_G = Wishart(ndims(q_out) + 2, 0.25*[1.0 0.0; 0.0 1.0])
31+
# m_μ, Cov_μ = mean_cov(q_μ)
32+
# m_out, Cov_out = mean_cov(q_out)
33+
# AE = - div(ndims(q_μ),2) * mean(log,q_γ) - 0.5*mean(logdet, q_G) + div(ndims(q_μ),2)*log(2pi) + 0.5*tr(mean(q_γ)*mean(q_G)*( Cov_out + Cov_μ + (m_out - m_μ)*(m_out - m_μ)'))
34+
35+
for N1 in (MvNormalMeanPrecision, MvNormalMeanCovariance, MvNormalWeightedMeanPrecision),
36+
N2 in (MvNormalMeanPrecision, MvNormalMeanCovariance, MvNormalWeightedMeanPrecision), g in (Gamma,),
37+
M in (Wishart,)
38+
39+
marginals = (
40+
Marginal(convert(N1, q_out), false, false, nothing),
41+
Marginal(convert(N2, q_μ), false, false, nothing),
42+
Marginal(convert(g, q_γ), false, false, nothing),
43+
Marginal(convert(M, q_G), false, false, nothing)
44+
)
45+
@test score(AverageEnergy(), MvNormalMeanScaleMatrixPrecision, Val{(:out, :μ, :γ, :G)}(), marginals, nothing) ≈ 189.18709448153223
46+
end
47+
end
48+
49+
begin
50+
q_out_μ = MvNormalMeanPrecision(
51+
[0.2932046487282065, 0.7716085147100042, 0.03978072440454361, 0.2814883836121471],
52+
[
53+
1.0684808331872628 0.6721958601342372 0.7164104160110533 0.2869444930570181
54+
0.6721958601342372 1.4472786772965438 1.3516272546828385 0.5533932426057602
55+
0.7164104160110533 1.3516272546828385 1.2958919150214063 0.5171784879755998
56+
0.2869444930570181 0.5533932426057602 0.5171784879755998 0.30898006058959576
57+
]
58+
)
59+
d = div(ndims(q_out_μ), 2)
60+
q_γ = GammaShapeRate(3.0, 2.0)
61+
q_G = Wishart(d + 2, 0.25*[0.349811 0.318591; 0.318591 0.401713])
62+
# m_out_μ, Cov_out_μ = mean_cov(q_out_μ)
63+
# m_out, m_μ = @views m_out_μ[1:d], m_out_μ[(d + 1):end]
64+
# Cov_out, Cov_μ = @views Cov_out_μ[1:d, 1:d], Cov_out_μ[(d + 1):end, (d + 1):end]
65+
# Cov_out_out, Cov_μ_μ = @views Cov_out_μ[1:d, (d + 1):end], Cov_out_μ[(d + 1):end, 1:d]
66+
# AE = - div(ndims(q_μ),2) * mean(log,q_γ) - 0.5*mean(logdet, q_G) + div(ndims(q_μ),2)*log(2pi) + 0.5*tr(mean(q_γ)*mean(q_G)*( Cov_out + Cov_μ - Cov_out_out - Cov_μ_μ + (m_out - m_μ)*(m_out - m_μ)'))
67+
68+
for N in (MvNormalMeanPrecision, MvNormalMeanCovariance, MvNormalWeightedMeanPrecision)
69+
marginals = (Marginal(convert(N, q_out_μ), false, false, nothing), Marginal(q_γ, false, false, nothing), Marginal(q_G, false, false, nothing))
70+
@test score(AverageEnergy(), MvNormalMeanScaleMatrixPrecision, Val{(:out_μ, :γ, :G)}(), marginals, nothing) ≈ 57.01394241406646
71+
end
72+
end
73+
74+
begin
75+
q_out_μ = MvNormalMeanPrecision(
76+
[0.35156676223859784, 0.6798203100143094, 0.954485919235333, 0.9236981452828203],
77+
[
78+
1.3182839156957349 0.9159049032047119 1.170482409249098 1.132202025059748
79+
0.9159049032047119 1.4737964254194567 1.4024254322343757 0.7350293025705011
80+
1.170482409249098 1.4024254322343757 2.0577570913647105 1.3137472032115916
81+
1.132202025059748 0.7350293025705011 1.3137472032115916 1.2083880803032556
82+
]
83+
)
84+
q_γ = GammaShapeRate(4.0, 3.0)
85+
q_G = Wishart(ndims(q_out) + 2, 0.25*[5.60439 4.34489; 4.34489 3.69273])
86+
# m_out_μ, Cov_out_μ = mean_cov(q_out_μ)
87+
# m_out, m_μ = @views m_out_μ[1:d], m_out_μ[(d + 1):end]
88+
# Cov_out, Cov_μ = @views Cov_out_μ[1:d, 1:d], Cov_out_μ[(d + 1):end, (d + 1):end]
89+
# Cov_out_out, Cov_μ_μ = @views Cov_out_μ[1:d, (d + 1):end], Cov_out_μ[(d + 1):end, 1:d]
90+
# AE = - div(ndims(q_μ),2) * mean(log,q_γ) - 0.5*mean(logdet, q_G) + div(ndims(q_μ),2)*log(2pi) + 0.5*tr(mean(q_γ)*mean(q_G)*( Cov_out + Cov_μ - Cov_out_out - Cov_μ_μ + (m_out - m_μ)*(m_out - m_μ)'))
91+
92+
for N in (MvNormalMeanPrecision, MvNormalMeanCovariance, MvNormalWeightedMeanPrecision)
93+
marginals = (Marginal(convert(N, q_out_μ), false, false, nothing), Marginal(q_γ, false, false, nothing), Marginal(q_G, false, false, nothing))
94+
@test score(AverageEnergy(), MvNormalMeanScaleMatrixPrecision, Val{(:out_μ, :γ, :G)}(), marginals, nothing) ≈ 60.58871007449749
95+
end
96+
end
97+
end
98+
end

0 commit comments

Comments
 (0)