Improve documentation and error handling mixture node

Initial Issue

The Mixture node's out marginalisation rule was failing with a cryptic error:

MethodError: no method matching getlogscale(::Nothing)

I was a bit surprised by that and I haven't found the tests for mixture node out rules, maybe I missing smt here. So I wrote them on my own. In the current ReactiveMP it fails. Could it be I am misunderstading how this node can be used correctly?

@testset "Marginalisation: (m_switch::Categorical, m_inputs::ManyOf)" begin
        @test_rules [check_type_promotion = false] Mixture(:out, Marginalisation) [
            # Test case 1: Equal weights
            (
                input = (
                    m_switch = Categorical([0.5, 0.5]),
                    m_inputs = ManyOf(
                        NormalMeanVariance(0.0, 1.0),
                        NormalMeanVariance(2.0, 1.0)
                    )
                ),
                output = MixtureDistribution([
                    NormalMeanVariance(0.0, 1.0),
                    NormalMeanVariance(2.0, 1.0)
                ], [0.5, 0.5])
            ),
            # Test case 2: Unequal weights
            (
                input = (
                    m_switch = Categorical([0.8, 0.2]),
                    m_inputs = ManyOf(
                        NormalMeanVariance(1.0, 1.0),
                        NormalMeanVariance(5.0, 2.0)
                    )
                ),
                output = MixtureDistribution([
                    NormalMeanVariance(1.0, 1.0),
                    NormalMeanVariance(5.0, 2.0)
                ], [0.8, 0.2])
            )
        ]
    end

So I decided to re-write this rule myself.

I come up with the following implementation

@rule Mixture(:out, Marginalisation) (m_switch::Any, m_inputs::ManyOf{N, Any}) where {N} = begin
    # Get logscales, defaulting to 0.0 if Nothing
    logscales_inputs = map(msg -> getlogscale(getdata(msg)) === nothing ? 0.0 : getlogscale(getdata(msg)), messages[2])
    logscale_switch = getlogscale(getdata(messages[1])) === nothing ? 0.0 : getlogscale(getdata(messages[1]))

    # compute logscales of individual components
    logscales = logscales_inputs .+ logscale_switch

    @logscale logsumexp(logscales)

    # Use probabilities directly from m_switch
    w = probvec(m_switch)
    T = promote_type(eltype(w), map(x -> eltype(mean(x)), m_inputs)...)

    # Convert inputs to the promoted type
    typed_inputs = map(x -> convert_paramfloattype(T, x), m_inputs)

    # return mixture with type-preserved components
    return MixtureDistribution(collect(typed_inputs), collect(w))
end

Also to make the rule work I need to write quite some helping methods, so to run test now, you need to use the following code

@testitem "rules:Mixture:out" begin
    using ReactiveMP, BayesBase, Random, ExponentialFamily, Distributions

    import ReactiveMP: @test_rules
    import ReactiveMP: getlogscale
    import BayesBase: paramfloattype
    import Base: isapprox
    using ExponentialFamily: NormalMeanVariance
    
    function getlogscale(d::NormalMeanVariance{T}) where {T}
        μ, τ = mean_precision(d)
        # Log of normalization constant for Normal(μ, 1/√τ)
        return -0.5 * (log(2π) - log(τ))
    end

    function getlogscale(d::Categorical{T}) where {T}
        # Categorical distribution is already normalized
        return 0.0
    end

    # Add paramfloattype for MixtureDistribution
    function paramfloattype(d::MixtureDistribution{D, T}) where {D, T}
        # The float type should be the promoted type of both the component distributions and weights
        return promote_type(T, paramfloattype(first(d.components)))
    end

    # Add isapprox for MixtureDistribution
    function isapprox(x::MixtureDistribution, y::MixtureDistribution; kwargs...)
        # Check if components and weights match approximately
        return length(x.components) == length(y.components) &&
               all(isapprox.(x.components, y.components; kwargs...)) &&
               isapprox(x.weights, y.weights; kwargs...)
    end

    @testset "Marginalisation: (m_switch::Categorical, m_inputs::ManyOf)" begin
        @test_rules [check_type_promotion = false] Mixture(:out, Marginalisation) [
            # Test case 1: Equal weights
            (
                input = (
                    m_switch = Categorical([0.5, 0.5]),
                    m_inputs = ManyOf(
                        NormalMeanVariance(0.0, 1.0),
                        NormalMeanVariance(2.0, 1.0)
                    )
                ),
                output = MixtureDistribution([
                    NormalMeanVariance(0.0, 1.0),
                    NormalMeanVariance(2.0, 1.0)
                ], [0.5, 0.5])
            ),
            # Test case 2: Unequal weights
            (
                input = (
                    m_switch = Categorical([0.8, 0.2]),
                    m_inputs = ManyOf(
                        NormalMeanVariance(1.0, 1.0),
                        NormalMeanVariance(5.0, 2.0)
                    )
                ),
                output = MixtureDistribution([
                    NormalMeanVariance(1.0, 1.0),
                    NormalMeanVariance(5.0, 2.0)
                ], [0.8, 0.2])
            )
        ]
    end

    @testset "Marginalisation: (m_inputs::ManyOf, q_switch::PointMass)" begin
        @test_rules [check_type_promotion = false] Mixture(:out, Marginalisation) [
            # Test case for one-hot encoded switch
            (
                input = (
                    m_inputs = ManyOf(
                        NormalMeanVariance(0.0, 1.0),
                        NormalMeanVariance(2.0, 1.0)
                    ),
                    q_switch = PointMass([1.0, 0.0])
                ),
                output = NormalMeanVariance(0.0, 1.0)
            ),
            (
                input = (
                    m_inputs = ManyOf(
                        NormalMeanVariance(0.0, 1.0),
                        NormalMeanVariance(2.0, 1.0)
                    ),
                    q_switch = PointMass([0.0, 1.0])
                ),
                output = NormalMeanVariance(2.0, 1.0)
            )
        ]
    end
end

Initially my goal was to run the following model

@model function mixture_model(y)
    # Use fixed mixing proportions
    s = [1/3, 2/3]
    
    # Parameters for the two Gaussian components
    m[1] ~ Normal(mean = -20.0, variance = 1e2)
    w[1] ~ InverseGamma(2.0, 1.0)  # Use InverseGamma for precision
    
    m[2] ~ Normal(mean = 20.0, variance = 1e2)
    w[2] ~ InverseGamma(2.0, 1.0)  # Use InverseGamma for precision
    
    obs_precision ~ InverseGamma(2.0, 1.0)
    
    # Generate mixture assignments and observations
    for i in eachindex(y)
        z[i] ~ Categorical(s)
        comp1[i] ~ NormalMeanVariance(m[1], w[1])  # Use variance parameterization
        comp2[i] ~ NormalMeanVariance(m[2], w[2])
        μ[i] ~ Mixture(switch = z[i], inputs = (comp1[i], comp2[i]))
        y[i] ~ NormalMeanVariance(μ[i], obs_precision)
    end
end

So the end-to-end script to run is the following one:

using RxInfer, Distributions
using Random
using Plots
using StatsPlots
using ReactiveMP: getlogscale

# Add getlogscale methods for both distribution types
import ReactiveMP: getlogscale
function getlogscale(d::NormalMeanPrecision{T}) where {T}
    μ, τ = mean_precision(d)
    # Log of normalization constant for Normal(μ, 1/√τ)
    return -0.5 * (log(2π) - log(τ))
end

function getlogscale(d::NormalWeightedMeanPrecision{T}) where {T}
    μ, τ = mean_precision(d)
    # Log of normalization constant for Normal(μ, 1/√τ)
    return -0.5 * (log(2π) - log(τ))
end



function getlogscale(d::NormalMeanVariance{T}) where {T}
    μ, τ = mean_precision(d)
    # Log of normalization constant for Normal(μ, 1/√τ)
    return -0.5 * (log(2π) - log(τ))
end

function getlogscale(d::Categorical{T}) where {T}
    # Categorical distribution is already normalized
    return 0.0
end

function getlogscale(_::Nothing)
    return 0.0
end

# include("mixture_entropy.jl") # include this to use free energy

@model function mixture_model(y)
    # Use fixed mixing proportions
    s = [1/3, 2/3]
    
    # Parameters for the two Gaussian components
    m[1] ~ Normal(mean = -20.0, variance = 1e2)
    w[1] ~ InverseGamma(2.0, 1.0)  # Use InverseGamma for precision
    
    m[2] ~ Normal(mean = 20.0, variance = 1e2)
    w[2] ~ InverseGamma(2.0, 1.0)  # Use InverseGamma for precision
    
    obs_precision ~ InverseGamma(2.0, 1.0)
    
    # Generate mixture assignments and observations
    for i in eachindex(y)
        z[i] ~ Categorical(s)
        comp1[i] ~ NormalMeanVariance(m[1], w[1])  # Use variance parameterization
        comp2[i] ~ NormalMeanVariance(m[2], w[2])
        μ[i] ~ Mixture(switch = z[i], inputs = (comp1[i], comp2[i]))
        y[i] ~ NormalMeanVariance(μ[i], obs_precision)
    end
end

# Update constraints
@constraints function mixture_constraints()
    q(z, m, w, μ, comp1, comp2, obs_precision) = q(z)q(m)q(w)q(μ)q(comp1)q(comp2)q(obs_precision)
    q(m) = q(m[1])q(m[2])
    q(w) = q(w[1])q(w[2])
    q(z) = q(z[1]) .. q(z[end])
    q(μ) = q(μ[1]) .. q(μ[end])
    q(comp1) = q(comp1[1]) .. q(comp1[end])
    q(comp2) = q(comp2[1]) .. q(comp2[end])
end

# Generate synthetic data with unequal proportions
rng = MersenneTwister(42)
true_means = [-20.0, 20.0]
true_precisions = [1.0, 1.0]
N = 1000
switch = [1/3, 2/3]  # Unequal proportions as in test
z = rand(rng, Categorical(switch), N)
data = zeros(N)
for i in 1:N
    data[i] = randn(rng)/sqrt(true_precisions[z[i]]) + true_means[z[i]]
end

# Update initialization with better starting points
init = @initialization begin
    # Initialize means further apart
    q(m[1]) = NormalMeanVariance(-30.0, 10.0)  # More uncertainty in initial means
    q(m[2]) = NormalMeanVariance(30.0, 10.0)
    
    # Initialize variances with more informative priors
    q(w[1]) = InverseGamma(3.0, 2.0)  # Mode around 1.0
    q(w[2]) = InverseGamma(3.0, 2.0)
    q(obs_precision) = InverseGamma(3.0, 2.0)
    
    for i in 1:N
        # Initialize assignments closer to true proportions
        q(z[i]) = Categorical([0.4, 0.6])
        
        # Initialize components with wider separation
        q(comp1[i]) = NormalMeanVariance(-20.0, 5.0)
        q(comp2[i]) = NormalMeanVariance(20.0, 5.0)
        q(μ[i]) = NormalMeanVariance(0.0, 100.0)  # Very uncertain about mixture means
    end
end

result = infer(
    model = mixture_model(),
    constraints = mixture_constraints(),
    initialization = init,
    data = (y = data,),
    iterations = 10,
    allow_node_contraction = true,
    options = (limit_stack_depth = 100,),
    # free_energy = true
)

# Create a range for plotting
x_range = range(minimum(data) - 1, maximum(data) + 1, length=200)

# Get the final parameters
m1 = mean(result.posteriors[:m][1][end])
m2 = mean(result.posteriors[:m][2][end])
v1 = mean(result.posteriors[:w][1][end])  # This is variance now
v2 = mean(result.posteriors[:w][2][end])

# Create the plot
p = histogram(data, normalize=true, alpha=0.3, label="Data", bins=50)
plot!(x_range, 
    x -> pdf(Normal(m1, sqrt(v1)), x),
    label="Component 1", linestyle=:dash)
plot!(x_range, 
    x -> pdf(Normal(m2, sqrt(v2)), x),
    label="Component 2", linestyle=:dash)
title!("Gaussian Mixture Model Fit")
xlabel!("x")
ylabel!("Density")

# Save the plot
savefig(p, "mixture_fit.png")

# Print the fitted parameters
println("\nFitted Parameters:")
println("Mean 1: ", round(m1, digits=3))
println("Mean 2: ", round(m2, digits=3))
println("Variance 1: ", round(v1, digits=3))
println("Variance 2: ", round(v2, digits=3))
println("Observation variance: ", round(mean(result.posteriors[:obs_precision][end]), digits=3))

# plot(1:10, result.free_energy)

Interestingly it shows a different behavior (and I would say more interesting one) comparing it with NormalMixture

Normal mixture model (at least the following one) showing collapsing behavior:

using RxInfer, Distributions
using Random
using Plots
using StatsPlots
using ReactiveMP: getlogscale


@model function mixture_model_normal_mixture(y)
    # Use fixed mixing proportions
    s = [1/3, 2/3]
    
    # Parameters for the two Gaussian components
    m[1] ~ Normal(mean = -20.0, variance = 1e2)
    w[1] ~ GammaShapeRate(2.0, 1.0)
    
    m[2] ~ Normal(mean = 20.0, variance = 1e2)
    w[2] ~ GammaShapeRate(2.0, 1.0)
    
    # Generate mixture assignments and observations
    for i in eachindex(y)
        z[i] ~ Categorical(s)
        # Using p for precision interface as shown in the tests
        y[i] ~ NormalMixture(
            switch = z[i],
            m = (m[1], m[2]),
            p = (w[1], w[2])  # Changed from v to p to match test code
        )
    end
end

@constraints function mixture_constraints()
    q(z, m, w) = q(z)q(m)q(w)
    q(m) = q(m[1])q(m[2])
    q(w) = q(w[1])q(w[2])
    q(z) = q(z[1]) .. q(z[end])
end

# Use same data generation
rng = MersenneTwister(42)
true_means = [-20.0, 20.0]
true_precisions = [1.0, 1.0]
N = 1000
switch = [1/3, 2/3]
z = rand(rng, Categorical(switch), N)
data = zeros(N)
for i in 1:N
    data[i] = randn(rng)/sqrt(true_precisions[z[i]]) + true_means[z[i]]
end

# Use same initialization strategy
init = @initialization begin
    q(m[1]) = NormalMeanVariance(-20.0, 5.0)
    q(m[2]) = NormalMeanVariance(20.0, 5.0)
    
    q(w[1]) = GammaShapeRate(3.0, 2.0)
    q(w[2]) = GammaShapeRate(3.0, 2.0)
    
    for i in 1:N
        q(z[i]) = Categorical([0.4, 0.6])
    end
end

result = infer(
    model = mixture_model_normal_mixture(),
    constraints = mixture_constraints(),
    initialization = init,
    data = (y = data,),
    iterations = 10,
    allow_node_contraction = true,
    options = (limit_stack_depth = 100,),
    free_energy = true
)

# Plotting
x_range = range(minimum(data) - 1, maximum(data) + 1, length=200)

m1 = mean(result.posteriors[:m][1][end])
m2 = mean(result.posteriors[:m][2][end])
v1 = mean(result.posteriors[:w][1][end])
v2 = mean(result.posteriors[:w][2][end])

p = histogram(data, normalize=true, alpha=0.3, label="Data", bins=50)
plot!(x_range, 
    x -> pdf(Normal(m1, sqrt(v1)), x),
    label="Component 1", linestyle=:dash)
plot!(x_range, 
    x -> pdf(Normal(m2, sqrt(v2)), x),
    label="Component 2", linestyle=:dash)
title!("Gaussian Mixture Model Fit (NormalMixture)")
xlabel!("x")
ylabel!("Density")

savefig(p, "normal_mixture_fit.png")

println("\nFitted Parameters (NormalMixture):")
println("Mean 1: ", round(m1, digits=3))
println("Mean 2: ", round(m2, digits=3))
println("Variance 1: ", round(v1, digits=3))
println("Variance 2: ", round(v2, digits=3))


@show result.free_energy[end]

plot(1:10, result.free_energy)

@bartvanerp

AFAIK the rules for Mixture are supposed to be used together with the log-scale addon. See the example here.

result_mary = infer(
    model = beta_model_mary(), 
    data  = (y = dataset, ),
    returnvars = (θ = KeepLast(), θ_john = KeepLast(), θ_jane = KeepLast(), john_is_right = KeepLast()),
    addons = AddonLogScale(),  # <----- here
)

Should be better documented @bartvanerp ?

As discussed earlier, the mixture node is specifically designed for sum-product message passing to mimic Bayesian model comparison. This is also where it's operations are limited to at the moment. I will draft a PR soon improving the error messages.

I am not sure why this issue was closed. The documentation is probably handled fine but it also points to a bug or strange behavior with the NormalMixture. We should discuss why we observe collapsing behavior.

@ismailsenoz

@ismailsenoz NormalMixture and Mixture are different nodes with quite different implementations. And what is the "collapsing behavior"?

@Nimrais

Yes I understand that the implementations are different and aimed at solving different problems. However, the example @Nimrais posted shows that the NormalMixture node can not properly recover means/precisions/regimes. Basically, it estimates one component and the other component is not estimated or estimated as the same (Collapsing. Sorry for vague explanations in the earlier comment). Although, the issue was about documentation I think this behavior of NormalMixture is not desirable. As far as I know, the collapsing behavior was independently observed by me ,@albertpod and @bartvanerp as well. Perhaps, we can open a new issue if need be.

Usually, by mixture collapse, people mean the following: a phenomenon where mixture components k and j in model p(x) converge to identical parameters θₖ = θⱼ during fitting, reducing the mixture from K to K-1 effective components:

p(x) = ∑ᵢ πᵢ f(x|θᵢ) → ∑ᵢ≠ₖ,ⱼ πᵢ f(x|θᵢ) + (πₖ + πⱼ)f(x|θₖ)

where πᵢ are mixing weights and f(x|θᵢ) are component distributions with parameters θᵢ. And this is precisely what happens in the example above.

@Nimrais

I see, you're right, something is wrong. We don't have this behavior in the example and the there is no collapsing but there it uses Beta/Bernoulli for the z[i] variable (only 2 components). Later it uses Dirichlet/Categorical for more than two components setup which also doesn't have the collapsing and works oke. In the example from @Nimrais he used Categorical with fixed parameters, smth is broken with this setup then. Or Categorical setup is broken specifically for 2 components

Hi! We are moving to a cleaner Epic -> Feature -> Task issue hierarchy to better organize our backlog.
This issue is currently either underspecified or not tagged appropriately.

To keep this issue open, please do the following within the next 7 days (by 25-11-2025):

Update/Replace: Ensure the description is clear and actionable.
Tag Correctly:For Tasks/Features, add the correct label (e.g., feature, task) AND include a link to the Parent Epic or Feature it belongs to.
For Bugs, add the Bug label. (Bugs do not require a parent link.)

Issues not updated, linked, or tagged correctly by the deadline will be closed and purged.

Thank you for helping us clean up and organize our backlog!

@bvdmitri

Following up here from issue #600 and PR #614. That fixed a related but different bug in InputArgumentsAnnotations. Two Point I think are worth looking at :

First, PR #446's helpful getlogscale(::Nothing) error message seems to have been lost during the addons to annotations rewrite in PR #596. I checked the diff and it deletes those lines with no equivalent added. Running Mixture today with no annotations option gives a raw KeyError again, not that guided message.

Second separate thing, I tested the actually documented setup, LogScaleAnnotations() enabled like bvdmitri described in the first comment above. It still crashes with the same KeyError, specifically when a Mixture input or its switch is passed in as a pre built Distribution rather than built inline. Traced it to a few direct getlogscale reads in the Mixture rules that assume every message went through a rule, which clamped messages never do.

A fix for this could be, defaulting a missing logscale to 0, I validated it numerically against a few reference models, it does stop crashing and the existing test suite still passes.

Given what I have understood is that Mixture support is intentionally narrow right now, should a clamped constant input actually be made to work, or should it fail with a clear intentional error instead?

Happy to cover this If you guys want.

cc @bvdmitri tagging you directly since this is basically picking up where your comment above left off

Navigation Menu

Repository navigation

Improve documentation and error handling mixture node #444

Initial Issue

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Improve documentation and error handling mixture node #444

Description

Initial Issue

Activity

bvdmitri commented on Feb 2, 2025

bartvanerp commented on Feb 3, 2025

ismailsenoz commented on Feb 10, 2025

bvdmitri commented on Feb 11, 2025

ismailsenoz commented on Feb 13, 2025

Nimrais commented on Feb 13, 2025

bvdmitri commented on Feb 13, 2025

wouterwln commented on Nov 18, 2025

gauraaansh commented on Jul 7, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions