This example has been auto-generated from the examples/ folder at GitHub repository.

Simple Nonlinear Node

# Activate local environment, see `Project.toml`
import Pkg; Pkg.activate(".."); Pkg.instantiate();

using RxInfer, Random, StableRNGs

Here is an example of creating custom node with nonlinear function approximation with samplelist.

Custom node creation

struct NonlinearNode end # Dummy structure just to make Julia happy

struct NonlinearMeta{R, F}
    rng      :: R
    fn       :: F   # Nonlinear function, we assume 1 float input - 1 float output
    nsamples :: Int # Number of samples used in approximation
end

@node NonlinearNode Deterministic [ out, in ]

We need to define two Sum-product message computation rules for our new custom node

• Rule for outbound message on out edge given inbound message on in edge
• Rule for outbound message on in edge given inbound message on out edge
• Both rules accept optional meta object

# Rule for outbound message on `out` edge given inbound message on `in` edge
@rule NonlinearNode(:out, Marginalisation) (m_in::NormalMeanVariance, meta::NonlinearMeta) = begin 
    samples = rand(meta.rng, m_in, meta.nsamples)
    return SampleList(map(meta.fn, samples))
end

# Rule for outbound message on `in` edge given inbound message on `out` edge
@rule NonlinearNode(:in, Marginalisation) (m_out::Gamma, meta::NonlinearMeta) = begin     
    return ContinuousUnivariateLogPdf((x) -> logpdf(m_out, meta.fn(x)))
end

Model specification

After we have defined our custom node with custom rules we may proceed with a model specification:

\[\begin{aligned} p(\theta) &= \mathcal{N}(\theta|\mu_{\theta}, \sigma_{\theta}),\\ p(m) &= \mathcal{N}(\theta|\mu_{m}, \sigma_{m}),\\ p(w) &= f(\theta),\\ p(y_i|m, w) &= \mathcal{N}(y_i|m, w), \end{aligned}\]

Given this IID model, we aim to estimate the precision of a Gaussian distribution. We pass a random variable $\theta$ through a non-linear transformation $f$ to make it positive and suitable for a precision parameter of a Gaussian distribution. We, later on, will estimate the posterior of $\theta$.

@model function nonlinear_estimation(n)
    
    θ ~ Normal(mean = 0.0, variance = 100.0)
    m ~ Normal(mean = 0.0, variance = 1.0)
    
    w ~ NonlinearNode(θ)
    
    y = datavar(Float64, n)
    
    for i in 1:n
        y[i] ~ Normal(mean = m, precision = w)
    end
    
end

@constraints function nconstsraints(nsamples)
    q(θ) :: SampleList(nsamples, LeftProposal())
    q(w) :: SampleList(nsamples, RightProposal())
    
    q(θ, w, m) = q(θ)q(m)q(w)
end

nconstsraints (generic function with 1 method)

@meta function nmeta(fn, nsamples)
    NonlinearNode(θ, w) -> NonlinearMeta(StableRNG(123), fn, nsamples)
end

nmeta (generic function with 1 method)

Here we generate some data

nonlinear_fn(x) = abs(exp(x) * sin(x))

nonlinear_fn (generic function with 1 method)

seed = 123
rng  = StableRNG(seed)

niters   = 15 # Number of VMP iterations
nsamples = 5_000 # Number of samples in approximation

n = 500 # Number of IID samples
μ = -10.0
θ = -1.0
w = nonlinear_fn(θ)

data = rand(rng, NormalMeanPrecision(μ, w), n);

result = inference(
    model = nonlinear_estimation(n),
    meta =  nmeta(nonlinear_fn, nsamples),
    constraints = nconstsraints(nsamples),
    data = (y = data, ), 
    initmarginals = (m = vague(NormalMeanPrecision), w = vague(Gamma)),
    returnvars = (θ = KeepLast(), ),
    iterations = niters,  
    showprogress = true
)

Inference results:
  Posteriors       | available for (θ)

θposterior = result.posteriors[:θ]

SampleList(Univariate, 5000)

using Plots, StatsPlots

estimated = Normal(mean_std(θposterior)...)

plot(estimated, title="Posterior for θ", label = "Estimated", legend = :bottomright, fill = true, fillopacity = 0.2, xlim = (-3, 3), ylim = (0, 2))
vline!([ θ ], label = "Real value of θ")