Repository: biaslab/ReactiveMP.jl (julia)

Issue:
Relax constraints for Gaussian Mixture update rules
The current implementation of the update rules for Gaussian Mixture is very strict given the mean-field assumption around this node.

For example, the constraint on `q_m`  can be relaxed from `UnivariateNormalDistributionsFamily` to `Any` (same applies to `q_p`)

```julia
@rule NormalMixture{N}(:out, Marginalisation) (
    q_switch::Any,
    q_m::ManyOf{N, UnivariateNormalDistributionsFamily},
    q_p::ManyOf{N, GammaDistributionsFamily}
) where {N} = begin
    πs = probvec(q_switch)
    return NormalMeanPrecision(sum(πs .* mean.(q_m)), sum(πs .* mean.(q_p)))
end
```
A similar problem occurs for the update rules on other interfaces for GM node.

Pull request:
Updated Gaussian Mixture update and refined tests

Patch:
diff --git a/src/rules/normal_mixture/out.jl b/src/rules/normal_mixture/out.jl
index a5c028ff5..416f6172a 100644
--- a/src/rules/normal_mixture/out.jl
+++ b/src/rules/normal_mixture/out.jl
@@ -1,12 +1,6 @@
 
-@rule NormalMixture{N}(:out, Marginalisation) (q_switch::Any, q_m::ManyOf{N, UnivariateNormalDistributionsFamily}, q_p::ManyOf{N, GammaDistributionsFamily}) where {N} = begin
+@rule NormalMixture{N}(:out, Marginalisation) (q_switch::Any, q_m::ManyOf{N, Any}, q_p::ManyOf{N, Any}) where {N} = begin
     πs = probvec(q_switch)
-    return NormalMeanPrecision(sum(πs .* mean.(q_m)), sum(πs .* mean.(q_p)))
-end
-
-@rule NormalMixture{N}(:out, Marginalisation) (q_switch::Any, q_m::ManyOf{N, MultivariateNormalDistributionsFamily}, q_p::ManyOf{N, Wishart}) where {N} = begin
-    πs = probvec(q_switch)
-    d  = ndims(first(q_m))
 
     # Better to preinitialize
     q_p_m = mean.(q_p)
@@ -15,20 +9,7 @@ end
     W = mapreduce(x -> x[1] * x[2], +, zip(πs, q_p_m))
     ξ = mapreduce(x -> x[1] * x[2] * x[3], +, zip(πs, q_p_m, q_m_m))
 
-    return MvNormalWeightedMeanPrecision(ξ, W)
-end
-
-@rule NormalMixture{N}(:out, Marginalisation) (q_switch::Any, q_m::ManyOf{N, PointMass{T} where T <: Real}, q_p::ManyOf{N, PointMass{T} where T <: Real}) where {N} = begin
-    πs = probvec(q_switch)
-    return NormalMeanPrecision(sum(πs .* mean.(q_m)), sum(πs .* mean.(q_p)))
-end
-
-@rule NormalMixture{N}(:out, Marginalisation) (q_switch::Any, q_m::ManyOf{N, PointMass{<:AbstractVector}}, q_p::ManyOf{N, PointMass{<:AbstractMatrix}}) where {N} = begin
-    πs = probvec(q_switch)
-    d  = ndims(first(q_m))
-
-    w  = mapreduce(x -> x[1] * mean(x[2]), +, zip(πs, q_p))
-    xi = mapreduce(x -> x[1] * mean(x[2]) * mean(x[3]), +, zip(πs, q_p, q_m))
+    F = variate_form(ξ)
 
-    return MvNormalWeightedMeanPrecision(xi, w)
+    return convert(promote_variate_type(F, NormalWeightedMeanPrecision), ξ, W)
 end


Test patch:
diff --git a/test/rules/normal_mixture/test_out.jl b/test/rules/normal_mixture/test_out.jl
index 1b4d1bca4..ff6f39056 100644
--- a/test/rules/normal_mixture/test_out.jl
+++ b/test/rules/normal_mixture/test_out.jl
@@ -12,36 +12,44 @@ import ReactiveMP: @test_rules
         @test_rules [with_float_conversions = true] NormalMixture{2}(:out, Marginalisation) [
             (
                 input = (q_switch = Categorical([0.5, 0.5]), q_m = ManyOf(PointMass(1.0), PointMass(1.0)), q_p = ManyOf(PointMass(1.0), PointMass(1.0))),
-                output = NormalMeanPrecision(1.0, 1.0)
+                output = NormalWeightedMeanPrecision(1.0, 1.0)
             ),
             (
                 input = (q_switch = Categorical([1.0, 0.0]), q_m = ManyOf(PointMass(1.0), PointMass(2.0)), q_p = ManyOf(PointMass(2.0), PointMass(1.0))),
-                output = NormalMeanPrecision(1.0, 2.0)
+                output = NormalWeightedMeanPrecision(2.0, 2.0)
             ),
             (
                 input = (q_switch = Categorical([0.5, 0.5]), q_m = ManyOf(PointMass(1.0), PointMass(1.0)), q_p = ManyOf(PointMass(1.0), PointMass(1.0))),
-                output = NormalMeanPrecision(1.0, 1.0)
+                output = NormalWeightedMeanPrecision(1.0, 1.0)
             ),
             (
                 input = (q_switch = Categorical([1.0, 0.0]), q_m = ManyOf(PointMass(1.0), PointMass(2.0)), q_p = ManyOf(PointMass(2.0), PointMass(1.0))),
-                output = NormalMeanPrecision(1.0, 2.0)
+                output = NormalWeightedMeanPrecision(2.0, 2.0)
             ),
             (
                 input = (q_switch = Categorical([0.0, 1.0]), q_m = ManyOf(PointMass(2.0), PointMass(-3.0)), q_p = ManyOf(PointMass(4.0), PointMass(3.0))),
-                output = NormalMeanPrecision(-3.0, 3.0)
+                output = NormalWeightedMeanPrecision(-9.0, 3.0)
             )
         ]
     end
 
     @testset "Variational : (m_μ::UnivariateNormalDistributionsFamily..., m_p::GammaDistributionsFamily...)" begin
         @test_rules [with_float_conversions = true] NormalMixture{2}(:out, Marginalisation) [
+            (
+                input = (
+                    q_switch = Bernoulli(0.2),
+                    q_m = ManyOf(NormalMeanVariance(5.0, 2.0), NormalMeanVariance(10.0, 3.0)),
+                    q_p = ManyOf(GammaShapeRate(1.0, 2.0), GammaShapeRate(2.0, 1.0))
+                ),
+                output = NormalWeightedMeanPrecision(33 / 2, 17 / 10)
+            ),
             (
                 input = (
                     q_switch = Categorical([0.5, 0.5]),
                     q_m = ManyOf(NormalMeanVariance(1.0, 2.0), NormalMeanPrecision(-2.0, 3.0)),
                     q_p = ManyOf(GammaShapeRate(1.0, 1.0), GammaShapeScale(2.0, 0.1))
                 ),
-                output = NormalMeanPrecision(-1 / 2, 6 / 10)
+                output = NormalWeightedMeanPrecision(3 / 10, 6 / 10)
             ),
             (
                 input = (
@@ -49,7 +57,7 @@ import ReactiveMP: @test_rules
                     q_m = ManyOf(NormalWeightedMeanPrecision(-1.0, 2.0), NormalMeanPrecision(2.0, 3.0)),
                     q_p = ManyOf(GammaShapeScale(1.0, 1.0), GammaShapeRate(2.0, 0.1))
                 ),
-                output = NormalMeanPrecision(1 / 8, 5.75)
+                output = NormalWeightedMeanPrecision(77 / 8, 23 / 4)
             ),
             (
                 input = (
@@ -57,7 +65,7 @@ import ReactiveMP: @test_rules
                     q_m = ManyOf(NormalMeanVariance(1.0, 2.0), NormalMeanPrecision(-2.0, 3.0)),
                     q_p = ManyOf(GammaShapeRate(1.0, 1.0), GammaShapeScale(2.0, 0.1))
                 ),
-                output = NormalMeanPrecision(1, 1)
+                output = NormalWeightedMeanPrecision(1, 1)
             )
         ]
     end