Add structured VMP rules for softdot node

src/nodes/predefined/softdot.jl

@@ -14,3 +14,21 @@ const softdot = SoftDot
  m_γ = mean(q_γ)
  return (-mean(log, q_γ) + log2π + m_γ * (V_y + m_y^2 - 2m_γ * m_y * m_θ'm_x + mul_trace(V_θ, V_x) + m_x'V_θ * m_x + m_θ' * (V_x + m_x * m_x') * m_θ)) / 2
 end
+
+@average_energy softdot (q_y_x::MultivariateNormalDistributionsFamily, q_θ::NormalDistributionsFamily, q_γ::GammaShapeRate) = begin
+ mθ, Vθ = mean_cov(q_θ)
+ myx, Vyx = mean_cov(q_y_x)
+ mγ = mean(q_γ)
+
+ order = length(mθ)
+ F = order == 1 ? Univariate : Multivariate
+
+ mx, Vx = ar_slice(F, myx, (order + 1):(2order)), ar_slice(F, Vyx, (order + 1):(2order), (order + 1):(2order))
+ my1, Vy1 = first(myx), first(Vyx)
+ Vy1x = ar_slice(F, Vyx, 1, (order + 1):(2order))
+
+ # Equivalent to AE = (-mean(log, q_γ) + log2π + mγ*(Vy1+my1^2 - 2*mθ'*(Vy1x + mx*my1) + tr(Vθ*Vx) + mx'*Vθ*mx + mθ'*(Vx + mx*mx')*mθ)) / 2
+ AE = (-mean(log, q_γ) + log2π + mγ * (Vy1 + my1^2 - 2 * mθ' * (Vy1x + mx * my1) + mul_trace(Vθ, Vx) + dot(mx, Vθ, mx) + dot(mθ, Vx, mθ) + abs2(dot(mθ, mx)))) / 2
+
+ return AE
+end

src/rules/dot_product/in1.jl

@@ -2,3 +2,7 @@
 @rule typeof(dot)(:in1, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_in2::PointMass, meta::Union{AbstractCorrectionStrategy, Nothing}) = begin
  return @call_rule typeof(dot)(:in2, Marginalisation) (m_out = m_out, m_in1 = m_in2, meta = meta)
 end
+
+@rule typeof(dot)(:in1, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_in2::NormalDistributionsFamily, meta::Union{AbstractCorrectionStrategy, Nothing}) = begin
+ return error("The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead.")
+end

src/rules/dot_product/in2.jl

@@ -8,3 +8,7 @@
 
  return convert(promote_variate_type(variate_form(typeof(m_in1)), NormalWeightedMeanPrecision), ξ, W)
 end
+
+@rule typeof(dot)(:in2, Marginalisation) (m_out::UnivariateNormalDistributionsFamily, m_in1::NormalDistributionsFamily, meta::Union{AbstractCorrectionStrategy, Nothing}) = begin
+ return error("The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead.")
+end

src/rules/dot_product/out.jl

@@ -8,3 +8,7 @@ end
  in2_mean, in2_cov = mean_cov(m_in2)
  return NormalMeanVariance(dot(A, in2_mean), dot(A, in2_cov, A))
 end
+
+@rule typeof(dot)(:out, Marginalisation) (m_in1::NormalDistributionsFamily, m_in2::NormalDistributionsFamily, meta::Union{AbstractCorrectionStrategy, Nothing}) = begin
+ return error("The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead.")
+end

src/rules/predefined.jl

@@ -101,11 +101,6 @@ include("dot_product/out.jl")
 include("dot_product/in1.jl")
 include("dot_product/in2.jl")
 
-include("softdot/y.jl")
-include("softdot/x.jl")
-include("softdot/theta.jl")
-include("softdot/gamma.jl")
-
 include("transition/marginals.jl")
 include("transition/out.jl")
 include("transition/in.jl")
@@ -123,6 +118,12 @@ include("autoregressive/theta.jl")
 include("autoregressive/gamma.jl")
 include("autoregressive/marginals.jl")
 
+include("softdot/y.jl")
+include("softdot/x.jl")
+include("softdot/theta.jl")
+include("softdot/gamma.jl")
+include("softdot/marginals.jl")
+
 include("probit/marginals.jl")
 include("probit/in.jl")
 include("probit/out.jl")

src/rules/softdot/gamma.jl

@@ -12,3 +12,25 @@
  β += (mul_trace(Vx, Vθ) + mθ'Vx * mθ + mx'Vθ * mx + mθ'mx * mx'mθ) / 2
  return GammaShapeRate(α, β)
 end
+
+# Variational MP: Structured
+@rule softdot(:γ, Marginalisation) (q_y_x::MultivariateNormalDistributionsFamily, q_θ::Any) = begin
+ # q_y is always Univariate
+ order = length(q_y_x) - 1
+ F = order == 1 ? Univariate : Multivariate
+
+ y_x_mean, y_x_cov = mean_cov(q_y_x)
+ mθ, Vθ = mean_cov(q_θ)
+
+ my, Vy = first(y_x_mean), first(y_x_cov)
+ mx, Vx = ar_slice(F, y_x_mean, 2:(order + 1)), ar_slice(F, y_x_cov, 2:(order + 1), 2:(order + 1))
+ Vyx = ar_slice(F, y_x_cov, 2:(order + 1))
+
+ C = rank1update(Vx, mx)
+ R = rank1update(Vy, my)
+ L = Vyx + mx * my
+
+ B = first(R) - 2 * first(mθ' * L) + first(mθ' * C * mθ) + mul_trace(Vθ, C)
+
+ return GammaShapeRate(convert(eltype(B), 3//2), B / 2)
+end

src/rules/softdot/marginals.jl

@@ -0,0 +1,27 @@
+
+# The following marginal rule is adaptation of the marginal rule for Autoregressive node
+@marginalrule SoftDot(:y_x) (m_y::NormalDistributionsFamily, m_x::NormalDistributionsFamily, q_θ::NormalDistributionsFamily, q_γ::Any) = begin
+ mθ, Vθ = mean_cov(q_θ)
+ mγ = mean(q_γ)
+
+ b_my, b_Vy = mean_cov(m_y)
+ f_mx, f_Vx = mean_cov(m_x)
+
+ inv_b_Vy = cholinv(b_Vy)
+ inv_f_Vx = cholinv(f_Vx)
+
+ D = inv_f_Vx + mγ * Vθ
+
+ W_11 = inv_b_Vy + mγ
+
+ W_12 = -mγ * mθ'
+
+ W_21 = -mθ * mγ
+
+ W_22 = D + mθ * mγ * mθ'
+
+ W = [W_11 W_12; W_21 W_22]
+ ξ = [inv_b_Vy * b_my; inv_f_Vx * f_mx]
+
+ return MvNormalWeightedMeanPrecision(ξ, W)
+end

src/rules/softdot/theta.jl

@@ -9,3 +9,23 @@
  zθ = mγ * mx * my
  return convert(promote_variate_type(variate_form(typeof(q_x)), NormalWeightedMeanPrecision), zθ, Dθ)
 end
+
+# Variational MP: Structured
+@rule softdot(:θ, Marginalisation) (q_y_x::MultivariateNormalDistributionsFamily, q_γ::Any) = begin
+ # q_y is always Univariate
+ order = length(q_y_x) - 1
+ F = order == 1 ? Univariate : Multivariate
+
+ myx, Vyx = mean_cov(q_y_x)
+ my, Vy = first(myx), first(Vyx)
+ mx, Vx = ar_slice(F, myx, 2:(order + 1)), ar_slice(F, Vyx, 2:(order + 1), 2:(order + 1))
+ Vyx = ar_slice(F, Vyx, 2:(order + 1))
+
+ mγ = mean(q_γ)
+
+ W = mγ * (Vx + mx * mx')
+
+ ξ = (Vyx + mx * my') * mγ
+
+ return convert(promote_variate_type(F, NormalWeightedMeanPrecision), ξ, W)
+end

src/rules/softdot/x.jl

@@ -9,3 +9,21 @@
  zx = mγ * mθ * my
  return convert(promote_variate_type(variate_form(typeof(q_θ)), NormalWeightedMeanPrecision), zx, Dx)
 end
+
+# Variational MP: Structured
+@rule softdot(:x, Marginalisation) (m_y::UnivariateNormalDistributionsFamily, q_θ::Any, q_γ::Any) = begin
+ # the naive call of AR rule is not possible, because the softdot rule expects m_y to be a UnivariateNormalDistributionsFamily
+ mθ, Vθ = mean_cov(q_θ)
+ my, Vy = mean_cov(m_y)
+
+ mγ = mean(q_γ)
+
+ mV = inv(mγ)
+
+ C = mθ * inv(add_transition(Vy, mV))
+
+ W = C * mθ' + mγ * Vθ
+ ξ = C * my
+
+ return convert(promote_variate_type(variate_form(typeof(q_θ)), NormalWeightedMeanPrecision), ξ, W)
+end

src/rules/softdot/y.jl

@@ -2,3 +2,7 @@
 
 # Variational MP: Mean-field
 @rule softdot(:y, Marginalisation) (q_θ::Any, q_x::Any, q_γ::Any) = NormalMeanPrecision(mean(q_θ)'mean(q_x), mean(q_γ))
+
+@rule softdot(:y, Marginalisation) (q_θ::Any, m_x::Any, q_γ::Any) = NormalMeanVariance(
+ first.(mean_cov((@call_rule AR(:y, Marginalisation) (m_x = m_x, q_θ = q_θ, q_γ = q_γ, meta = ARMeta(variate_form(typeof(m_x)), length(q_θ), ARsafe())))))...
+)

test/nodes/predefined/softdot_tests.jl

@@ -22,5 +22,14 @@
 
  @test score(AverageEnergy(), SoftDot, Val{(:y, :θ, :x, :γ)}(), marginals, nothing) ≈ 8.15193210352257
  end
+
+ begin
+ q_y_x = MvNormalMeanCovariance(zeros(2), diageye(2))
+ q_θ = NormalMeanVariance(0.0, 1.0)
+ q_γ = GammaShapeRate(2.0, 3.0)
+
+ marginals = (Marginal(q_y_x, false, false, nothing), Marginal(q_θ, false, false, nothing), Marginal(q_γ, false, false, nothing))
+ @test score(AverageEnergy(), SoftDot, Val{(:y_x, :θ, :γ)}(), marginals, nothing) ≈ 1.92351917665616
+ end
  end # testset: AverageEnergy
 end # testset

test/rules/dot_product/in1_tests.jl

@@ -84,4 +84,12 @@
  )
  ]
  end
+
+ @testset "Error Belief Propagation: (m_out::UnivariateNormalDistributionsFamily, m_in2::NormalDistributionsFamily)" begin
+ @test_throws r"The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead." @call_rule typeof(
+ dot
+ )(
+ :in1, Marginalisation
+ ) (m_out = NormalMeanVariance(2.0, 2.0), m_in2 = MvNormalMeanCovariance([-1.0, 1.0], [2.0 -1.0; -1.0 4.0]), meta = NoCorrection())
+ end
 end

test/rules/dot_product/in2_tests.jl

@@ -84,4 +84,12 @@
  )
  ]
  end
+
+ @testset "Error Belief Propagation: (m_out::UnivariateNormalDistributionsFamily, m_in1::NormalDistributionsFamily)" begin
+ @test_throws r"The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead." @call_rule typeof(
+ dot
+ )(
+ :in2, Marginalisation
+ ) (m_out = NormalMeanVariance(2.0, 2.0), m_in1 = MvNormalMeanCovariance([-1.0, 1.0], [2.0 -1.0; -1.0 4.0]), meta = NoCorrection())
+ end
 end

test/rules/dot_product/out_tests.jl

@@ -97,4 +97,12 @@
  )
  ]
  end
+
+ @testset "Error belief Propagation: (m_in1::NormalDistributionsFamily, m_in2::NormalDistributionsFamily)" begin
+ @test_throws r"The rule for the dot product node between two NormalDistributionsFamily instances is not available in closed form. Please use SoftDot instead." @call_rule typeof(
+ dot
+ )(
+ :out, Marginalisation
+ ) (m_in1 = NormalMeanVariance(2.0, 2.0), m_in2 = NormalMeanVariance(2.0, 2.0), meta = NoCorrection())
+ end
 end

test/rules/softdot/gamma_tests.jl

@@ -131,4 +131,25 @@
  )
  end
  end # testset: mean-field
+
+ @testset "Structured: (q_y_x::MultivariateNormalDistributionsFamily, q_γ::Any)" begin
+ @test_rules [check_type_promotion = true] SoftDot(:γ, Marginalisation) [
+ (input = (q_y_x = MvNormalMeanCovariance(ones(2), diageye(2)), q_θ = NormalMeanPrecision(1.0, 1.0)), output = GammaShapeRate(3 / 2, 2.0)),
+ (input = (q_y_x = MvNormalMeanCovariance(2 * ones(2), diageye(2)), q_θ = NormalMeanPrecision(2.0, 1.0)), output = GammaShapeRate(3 / 2, 7.0))
+ ]
+ end
+
+ @testset "Structured : (q_y_x::MultivariateNormalDistributionsFamily, q_θ::Any)" begin
+ order = 2
+ @test_rules [check_type_promotion = true] SoftDot(:γ, Marginalisation) [
+ (
+ input = (q_y_x = MvNormalMeanCovariance(ones(order + 1), diageye(order + 1)), q_θ = MvNormalMeanPrecision(ones(order), diageye(order))),
+ output = GammaShapeRate(3 / 2, 4.0)
+ ),
+ (
+ input = (q_y_x = MvNormalMeanCovariance(ones(order + 1), diageye(order + 1)), q_θ = MvNormalMeanPrecision(zeros(order), diageye(order))),
+ output = GammaShapeRate(3 / 2, 3.0)
+ )
+ ]
+ end
 end # testset

test/rules/softdot/test_marginals.jl

@@ -0,0 +1,26 @@
+
+@testitem "marginalrules:SoftDot" begin
+ using ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions
+
+ import ReactiveMP: @test_marginalrules
+
+ @testset "y_x: (m_y::UnivariateNormalDistributionsFamily, m_x::UnivariateNormalDistributionsFamily, q_θ::UnivariateNormalDistributionsFamily, q_γ::Any)" begin
+ @test_marginalrules [check_type_promotion = true] SoftDot(:y_x) [(
+ input = (m_y = NormalMeanPrecision(0.0, 1.0), m_x = NormalMeanPrecision(0.0, 1.0), q_θ = NormalMeanPrecision(1.0, 1.0), q_γ = GammaShapeRate(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision(zeros(2), [2.0 -1.0; -1.0 3.0])
+ )]
+ end
+
+ @testset "y_x: (m_y::UnivariateNormalDistributionsFamily), m_x::MultivariateNormalDistributionsFamily, q_θ::MultivariateNormalDistributionsFamily, q_γ::Any)" begin
+ order = 2
+ @test_marginalrules [check_type_promotion = true] SoftDot(:y_x) [(
+ input = (
+ m_y = NormalMeanPrecision(1.0, 1.0),
+ m_x = MvNormalMeanCovariance(ones(order), diageye(order)),
+ q_θ = MvNormalMeanCovariance(ones(order), diageye(order)),
+ q_γ = GammaShapeRate(1.0, 1.0)
+ ),
+ output = MvNormalWeightedMeanPrecision(ones(3), [2.0 -1.0 -1.0; -1.0 3.0 1.0; -1.0 1.0 3.0])
+ )]
+ end
+end

test/rules/softdot/theta_tests.jl

@@ -129,4 +129,25 @@
  end
  # NOTE: γ can theoretically be Any, so also NormalMeanVariance
  end
+
+ @testset "Structured: (q_y_x::MultivariateNormalDistributionsFamily, q_γ::Any)" begin
+ @test_rules [check_type_promotion = true] SoftDot(:θ, Marginalisation) [
+ (input = (q_y_x = MvNormalMeanCovariance(ones(2), diageye(2)), q_γ = GammaShapeRate(1.0, 1.0)), output = NormalWeightedMeanPrecision(1.0, 2.0)),
+ (input = (q_y_x = MvNormalMeanCovariance(2 * ones(2), diageye(2)), q_γ = GammaShapeScale(2.0, 1.0)), output = NormalWeightedMeanPrecision(8.0, 10.0))
+ ]
+ end
+
+ @testset "Structured : (q_y_x::MultivariateNormalDistributionsFamily, q_γ::Any)" begin
+ order = 2
+ @test_rules [check_type_promotion = true] SoftDot(:θ, Marginalisation) [
+ (
+ input = (q_y_x = MvNormalMeanCovariance(ones(order + 1), diageye(order + 1)), q_γ = GammaShapeRate(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision(ones(order), [2.0 1.0; 1.0 2.0])
+ ),
+ (
+ input = (q_y_x = MvNormalMeanCovariance(zeros(order + 1), diageye(order + 1)), q_γ = GammaShapeRate(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision(zeros(order), [1.0 0.0; 0.0 1.0])
+ )
+ ]
+ end
 end # testset

test/rules/softdot/x_tests.jl

@@ -129,4 +129,34 @@
  end
  # NOTE: γ can theoretically be Any, so also NormalMeanVariance
  end
+
+ @testset "VMP: structured rules" begin
+ @testset "(m_y::NormalMeanVariance, q_θ::NormalMeanVariance, q_γ::Any)" begin
+ @test_rules [check_type_promotion = true] SoftDot(:x, Marginalisation) [
+ (input = (m_y = NormalMeanVariance(1.0, 1.0), q_θ = NormalMeanVariance(1.0, 1.0), q_γ = GammaShapeRate(1.0, 1.0)), output = NormalWeightedMeanPrecision(0.5, 1.5)),
+ (
+ input = (m_y = NormalWeightedMeanPrecision(1.0, 1.0), q_θ = NormalMeanPrecision(1.0, 2.0), q_γ = GammaShapeScale(1.0, 1.0)),
+ output = NormalWeightedMeanPrecision(0.5, 1.0)
+ )
+ ]
+ end
+
+ @testset "(m_y::UnivariateNormalDistributionsFamily, q_θ::MultivariateNormalDistributionsFamily, q_γ::Any)" begin
+ order = 2
+ @test_rules [check_type_promotion = false] SoftDot(:x, Marginalisation) [
+ (
+ input = (m_y = NormalMeanVariance(0.0, 1.0), q_θ = MvNormalMeanCovariance(ones(order), diageye(order)), q_γ = GammaShapeRate(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision([0.0, 0.0], [1.5 0.5; 0.5 1.5])
+ ),
+ (
+ input = (m_y = NormalMeanVariance(1.0, 1.0), q_θ = MvNormalMeanCovariance(zeros(order), diageye(order)), q_γ = GammaShapeScale(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision([0.0, 0.0], [1.0 0.0; 0.0 1.0])
+ ),
+ (
+ input = (m_y = NormalMeanVariance(1.0, 1.0), q_θ = MvNormalMeanCovariance(ones(order), diageye(order)), q_γ = Gamma(1.0, 1.0)),
+ output = MvNormalWeightedMeanPrecision([0.5, 0.5], [1.5 0.5; 0.5 1.5])
+ )
+ ]
+ end
+ end
 end # testset

test/rules/softdot/y_tests.jl

@@ -97,4 +97,37 @@
  )
  end
  end
+
+ @testset "VMP: structured rules" begin
+ @testset "(q_θ::NormalMeanVariance, m_x::NormalMeanVariance, q_γ::Any" begin
+ @test_rules [check_type_promotion = true] SoftDot(:y, Marginalisation) [
+ (input = (q_θ = PointMass(3.0), q_x = PointMass(11.0), q_γ = GammaShapeRate(7.0, 5.0)), output = NormalMeanPrecision(33.0, 1.4)),
+ (input = (q_θ = PointMass(3.0), q_x = PointMass(11.0), q_γ = GammaShapeScale(7.0, 5.0)), output = NormalMeanPrecision(33.0, 35.0))
+ ]
+
+ @test_rules [check_type_promotion = true] SoftDot(:y, Marginalisation) [
+ (input = (m_x = NormalMeanVariance(1.0, 1.0), q_θ = NormalMeanVariance(1.0, 1.0), q_γ = GammaShapeRate(1.0, 1.0)), output = NormalMeanVariance(0.5, 1.5)),
+ (
+ input = (m_x = NormalWeightedMeanPrecision(1.0, 1.0), q_θ = NormalMeanPrecision(1.0, 2.0), q_γ = GammaShapeScale(2.0, 1.0)),
+ output = NormalMeanVariance(0.5, 1.0)
+ )
+ ]
+ end
+
+ @testset "(q_θ::MvNormalMeanCovariance, m_x::MvNormalMeanCovariance, q_γ::Any" begin
+ order = 2
+ @test_rules [check_type_promotion = true] SoftDot(:y, Marginalisation) [
+ (
+ input = (
+ m_x = MvNormalMeanCovariance(ones(order), diageye(order)), q_θ = MvNormalMeanCovariance(zeros(order), diageye(order)), q_γ = GammaShapeScale(1.0, 1.0)
+ ),
+ output = NormalMeanVariance(0.0, 1.0)
+ ),
+ (
+ input = (m_x = MvNormalMeanCovariance(ones(order), diageye(order)), q_θ = MvNormalMeanCovariance(ones(order), diageye(order)), q_γ = Gamma(1.0, 1.0)),
+ output = NormalMeanVariance(1.0, 2.0)
+ )
+ ]
+ end
+ end
 end # testset