Rule for MvNormal and Dirichlets in a GMM.

[failin3]

Hey everyone.

I am building a Gaussian mixture model in which after having seen some data, it uses the Gaussians in the mixture to generate new (missing) data.

After having looked at N, 5 dimensional samples, new data is supplied where only 3 of the 5 dimensions are observed.
I am looking to extract the vector that describes the probvec for how much each Gaussian contributes to the distribution given the inputs, and then I'd like to multiply this probvec by the Gaussian components that have not been observed.
Here is my current code:

@model function GMM_strategy(n)
values_to_predict = 2
    
    # If we are only infering best option, skip this 
    if n > 0
        # Get all inputs as one single vector
        y = datavar(Vector{Float64}, n) 
        w = randomvar(2)
        w[1] ~ Wishart(5, 1e2*diagm(ones(5)))
        w[2] ~ Wishart(5, 1e2*diagm(ones(5)))
        precs = tuple(w...)
    end

    m = randomvar(2)
    
    z = randomvar(n+values_to_predict)

    m[1] ~ MvNormal(mean = [0.0, 0.0, 0.0, 10.0, 10.0], cov = diageye(5))
    m[2] ~ MvNormal(mean = [1.0, 0.0, 0.0, 0.0, 0.0], cov = diageye(5))

    s ~ Dirichlet(ones(2))
    

    means = tuple(m...)

    for k in 1:n
        z[k] ~ Categorical(s) 
        y[k] ~ NormalMixture(z[k], means, precs)
    end

    # Now lets calculate the options
    if values_to_predict > 0
        new_value = datavar(Vector{Float64}, values_to_predict) # [opp, posx, posy]

        a ~ dot(m[1], [1, 0, 0, 0, 0])*[1,0,0] + dot(m[1], [0, 1, 0, 0, 0])*[0,1,0] + dot(m[1], [0, 0, 1, 0, 0])*[0,0,1]
        b ~ dot(m[2], [1, 0, 0, 0, 0])*[1,0,0] + dot(m[2], [0, 1, 0, 0, 0])*[0,1,0] + dot(m[2], [0, 0, 1, 0, 0])*[0,0,1]
        for i = 1:values_to_predict
            z[n+i] ~ Categorical(s)
            new_value[i] ~ NormalMixture(z[n+i], (a,b), (diageye(3)/10, diageye(3)/10))
        end
    end
end

The top part is a typical GMM, and then at the bottom I create MvNormals a and b that have only 3 dimensions and use those to determine the mixture for my new input data.

What I then want to do is take the Gaussians of the data that was not observed (the other 2 dimensions) and multiply with s to mix the contributions. For a 2 Gaussian mixture that would be:

I have managed to create 4-vector for the MvNormals using
c ~ dot(m[1], [0, 0, 0, 1, 0])*[1, 0, 0, 0] + dot(m[1], [0, 0, 0, 1, 0])*[0, 0, 1, 0] + dot(m[2], [0, 0, 0, 1, 0])*[0, 1, 0, 0] + dot(m[2], [0, 0, 0, 1, 0])*[0, 0, 0, 1]
where m[1] and m[2] are the respective MvNormals.

I then think of using transition to reshape c into a 2x2 matrix and multiply it by s.

# Function to perform the state transition in the model.
# It reshapes vector `a` into a matrix and multiplies it with vector `x` to simulate the transition.
function transition(a, x)
    nm, n = length(a), length(x)
    m = nm ÷ n  # Calculate the number of rows for reshaping 'a' into a matrix
    A = reshape(a, (m, n))  
    return A * x
end

coords[i] ~ transition(c, s) where { meta = DeltaMeta(method = Linearization()) }
coords[i] ~ MvNormalMeanCovariance(zeros(2), diageye(2))

Unfortunately there is no rule to do this operation on a MvNormal and a Dirichlet.

ERROR: MarginalRuleMethodError: no method matching rule for the given arguments

Possible fix, define:

@marginalrule DeltaFn(:ins) (m_out::MvNormalMeanCovariance, m_ins::ManyOf{2, Union{MvNormalMeanCovariance,Dirichlet}}, meta::DeltaMeta{Linearization, Nothing}) = begin
    return ...
end

Is there anything I can do to get this to work?

Here's the full code including the synthetic data generator I use for testing
strategy_model_coordinates.zip

[failin3]

Having looked at the problem some more, it seems I can change my approach and use a second NormalMixture for the missing data.

If I split the 'trained' MvNormals of the original mixture into the 3 that I do know and the 2 that I don't like this

a ~ dot(m[1], [1, 0, 0, 0, 0])*[1,0,0] + dot(m[1], [0, 1, 0, 0, 0])*[0,1,0] + dot(m[1], [0, 0, 1, 0, 0])*[0,0,1]
b ~ dot(m[2], [1, 0, 0, 0, 0])*[1,0,0] + dot(m[2], [0, 1, 0, 0, 0])*[0,1,0] + dot(m[2], [0, 0, 1, 0, 0])*[0,0,1]

c ~ dot(m[1], [0, 0, 0, 1, 0])*[1, 0] + dot(m[1], [0, 0, 0, 0, 1])*[0, 1]
d ~ dot(m[2], [0, 0, 0, 1, 0])*[1, 0] + dot(m[2], [0, 0, 0, 0, 1])*[0, 1]

I can then, for each sample with missing data create two new NormalMixtures

new_value = datavar(Vector{Float64}, values_to_predict) # [opp, posx, posy]
coords = randomvar(values_to_predict)
for i = 1:values_to_predict
  z[n+i] ~ Categorical(s)
  # The 3D mixture with known outputs
  new_value[i] ~ NormalMixture(z[n+i], (a,b), ([0.5 0 0; 0 0.1 0; 0 0 0.1], [0.5 0 0; 0 0.1 0; 0 0 0.1]))
  # The 2D mixture with uknown outputs
  coords[i] ~ NormalMixture(z[n+i], (c,d), (diageye(2)*0.02, diageye(2)*0.02))
  coords[i] ~ MvNormalMeanCovariance(zeros(2), diageye(2)*huge)
end

Effectively splitting the mixture into the known part and unknown part, while using the same categorical z
Setting a large prior on coords should mean the prior doesn't influence the outcome much.