Score function error for simple custom Gaussian node in RxInfer

Specific Domains Probabilistic programming
1

I’m testing a simple Gaussian custom node in RxInfer just to make sure I have the syntax correct, but I’m having trouble with the score function for Bethe free energy. My inference() function runs fine and returns BFE if I use a Gaussian node rather than my custom node. And it runs fine if I use my custom node but set free_energy = false. If I set free_energy = true it fails with the following message:

LoadError: MethodError: no method matching score(::ReactiveMP.AverageEnergy, ::Type{Main.RxTest2.MyNode}, ::Type{Val{(:out, :m, :v)}}, ::Tuple{ReactiveMP.Marginal{ReactiveMP.NormalWeightedMeanPrecision{Float64}, Nothing}, ReactiveMP.Marginal{ReactiveMP.PointMass{Float64}, Nothing}, ReactiveMP.Marginal{ReactiveMP.PointMass{Float64}, Nothing}}, ::Nothing)
Closest candidates are:

I’m not sure how the score function is supposed to be specified, or if I am missing something else. Perhaps it needs additional rules, but if so I’m not asked for any when the code runs. Any ideas, @bvdmitri or others? Code is below.

struct MyNode end  
@node MyNode Stochastic [out, m, v]

@rule MyNode(:out, Marginalisation) (
    q_m::PointMass, 
    q_v::PointMass, 
    ) = begin 
    
    m = mean(q_m) 
    W = 1. / mean(q_v)
    return NormalWeightedMeanPrecision(m* W, W)
end

# --------------------------------------------------------------------------------------------------
@model function ref_model()
    m_x = datavar(Float64)
    v_x = datavar(Float64)
    x = randomvar()  
    
    #x ~ NormalMeanVariance(m_x, v_x)  # this works and gives BFE
    x ~ MyNode(m_x, v_x)  # does not give BFE 
    x ~ NormalMeanVariance(0.0, 100000.) 
    return (x, )
end

# --------------------------------------------------------------------------------------------------    
result = inference(
    data = (
        m_x = 10.,
        v_x = 1.,
    ),
    model = ref_model(),
    initmarginals = (
        x = NormalMeanVariance(0.0, huge),
    ),
    free_energy = true,
    returnvars = (
        x = KeepLast(),
    ),
)
        
x_ = mean.(result.posteriors[:x])
println(x_)

As a suggestion, perhaps this simple test case might be a good addition to the documentation, as it might help users get stared with custom nodes. I’ve looked at the documentation and at the ReactiveMP code for a Gaussian node and associated rules, but it’s not clear to me what the bare minimum code is for creating a simple Gaussian custom node that returns BFE.

2

Hi @John_Boik! Thanks for the question. Indeed, it seems that this part is missing from our documentation. So please open an issue on GitHub and link this conversation. I think it’s important to explain how free energy is computed locally and what’s needed when implementing your node.

What’s missing is the implementation of the @average_energy function.
The average energy is computed as -\mathbb{E}_{m, v, y}\log{f} where f is the functional form of MyNode (see appendix from van de Laar thesis). You can also have a look at the implementation of average energy for a NormalMeanVariance node.

Given that MyNode is just a Gaussian node, then we can add the following average energy:

@average_energy MyNode (q_out::Any, q_m::Any, q_v::Any) = begin
    μ_mean, μ_var     = mean_var(q_m)
    out_mean, out_var = mean_var(q_out)
    return (ReactiveMP.log2π + mean(log, q_v) + mean(inv, q_v) * (μ_var + out_var + abs2(μ_mean - out_mean))) / 2
end

Then the computations of free energy should work.