biaslab / ForneyLab.jl

Julia package for automatically generating Bayesian inference algorithms through message passing on Forney-style factor graphs.
MIT License
149 stars 33 forks source link

Algorithm construction bug (Nonlinear node) #166

Closed albertpod closed 3 years ago

albertpod commented 3 years ago

Given the following model

using ForneyLab
using LinearAlgebra
import ForneyLab: unsafeMean, unsafeCov

graph = FactorGraph()

T = 2
dimensionality = 2
x = Vector{Variable}(undef, T)
z = Vector{Variable}(undef, T)
y = Vector{Variable}(undef, T)

n_samples = 10
inputs = randn(n_samples, dimensionality)

@RV θ  ~ GaussianMeanPrecision(placeholder(:m_θ, dims=(dimensionality,)), placeholder(:W_θ, dims=(dimensionality, dimensionality)))
f(w, x) = 1/(1+exp(-w'x))
for i in 1:T
    @RV z[i] ~ GaussianMeanPrecision(inputs[i, :], 1e4*diageye(dimensionality))
    @RV x[i] ~ Nonlinear{Sampling}(θ, z[i], g=f, in_variates=[Multivariate, Multivariate], out_variate=Univariate)
    @RV y[i] ~ Bernoulli(x[i])
    placeholder(y[i], :y, index=i)
end

Following algorithm constuction

# Compile algorithm
algo = messagePassingAlgorithm(free_energy=true)
# Generate source code
src_code = algorithmSourceCode(algo, free_energy=true);
# Load algorithm
eval(Meta.parse(src_code))

After running

outputs = [rand([0,1]) for _ in 1:T]
data = Dict(:y => outputs, :m_θ => zeros(dimensionality), :W_θ => 0.1*diageye(dimensionality))
messages = init()
marginals = Dict()
step!(data, marginals, messages)

We get the following error: DimensionMismatch("first array has length 1 which does not match the length of the second, 2.") This happens because ForneyLab returns improper init() function:

function init()

  messages = Array{Message}(undef, 14)

  messages[1] = Message(vague(GaussianMeanPrecision))
  messages[2] = Message(vague(GaussianMeanPrecision))
  messages[6] = Message(vague(GaussianWeightedMeanPrecision))
  messages[9] = Message(vague(GaussianWeightedMeanPrecision))

  return messages

end

The messages are supposed to carry Multivariate distributions. We can circumvent this issue by creating a custom init() function, i.e.

function init()
    messages = Array{Message}(undef, 14)
    for i in 1:length(messages)
        messages[i] = Message(vague(GaussianMeanPrecision, 2))
    end
    messages
end

However, when fixing this issue, we encounter a different problem, i.e.

outputs = [rand([0,1]) for _ in 1:T]
data = Dict(:y => outputs, :m_θ => zeros(dimensionality), :W_θ => 0.1*diageye(dimensionality))
messages = init()
marginals = Dict()
step!(data, marginals, messages)

MethodError: no method matching ruleSPNonlinearSOutNGX(::typeof(f), ::Nothing, ::Message{GaussianWeightedMeanPrecision, Multivariate}, ::Message{GaussianMeanPrecision, Multivariate}; variate=Univariate) Closest candidates are: ruleSPNonlinearSOutNGX(::Function, ::Nothing, ::Message{var"#s158", V} where var"#s158"<:Gaussian...; n_samples) where V<:ForneyLab.VariateType at /Users/apodusenko/.julia/dev/ForneyLab/src/engines/julia/update_rules/nonlinear_sampling.jl:73 got unsupported keyword argument "variate" The error happens when ForneyLab tries to compute the following message:

messages[13] = ruleSPNonlinearSOutNGX(f, nothing, messages[6], messages[1], variate=Univariate)

For some reason, when I omit the argument variate=Univariate the message will be computed. It's strange because nonlinear_sampling.jl exports the method with the variate argument, i.e. ruleSPNonlinearSOutNGX(g::Function, msg_out::Nothing, msgs_in::Vararg{Message{<:Gaussian, <:VariateType}}; n_samples=default_n_samples, variate) I can't figure out why it's happening. Any ideas?

albertpod commented 3 years ago

To make ForneyLab.jl generate correct init(), we need to provide optional parameters that specify the dimension of Nonlinear node's inputs. The model specification should be changed as follows:

@RV θ  ~ GaussianMeanPrecision(placeholder(:m_θ, dims=(dimensionality,)), placeholder(:W_θ, dims=(dimensionality, dimensionality)))
f(w, x) = 1/(1+exp(-w'x))
for i in 1:T
    @RV z[i] ~ GaussianMeanPrecision(inputs[i, :], 1e4*diageye(dimensionality))
    @RV x[i] ~ Nonlinear{Sampling}(θ, z[i], g=f, in_variates=[Multivariate, Multivariate], out_variate=Univariate)
    @RV y[i] ~ Bernoulli(x[i])
    placeholder(y[i], :y, index=i)
end

Thanks to @ThijsvdLaar.

The rest of the issue will be addressed in future PR.