Algorithm specification

Every rule of a node belongs to an algorithm. The algorithm selects which rules run and carries their parameters. Most nodes run under DefaultAlgorithm, and a model says nothing about them. Some nodes need a choice from you:

• the AR node, for autoregressive processes, needs the order of the process and whether it is univariate or multivariate: ARVMP;
• the GCV node, the Gaussian controlled variance, computes its messages with a cubature whose number of points you choose: GCVApproximation;
• a deterministic node, such as y := f(x), needs an approximation method: DeltaApproximation, see Deterministic nodes.

You give a node its algorithm in the model, with where { algorithm = ... }, or for the whole model with the @algorithm macro and the algorithm keyword of infer. The meta keyword, the @meta macro and where { meta = ... } still work, with a deprecation warning, as the old names of these three.

An algorithm in the model

The AR node lives in the AutoregressiveMessagePassingRules package, which RxInfer loads. This model observes an autoregressive process of order 2 through noise:

using RxInfer, StableRNGs

@model function ar_model(y, order)
    γ  ~ Gamma(shape = 1.0, rate = 1.0)
    θ  ~ MvNormal(mean = zeros(order), precision = diageye(order))
    x0 ~ MvNormal(mean = zeros(order), precision = diageye(order))
    c = [1.0; zeros(order - 1)]
    x_prev = x0
    for i in eachindex(y)
        x[i] ~ AR(x_prev, θ, γ) where { algorithm = ARVMP(Multivariate, order, ARsafe()) }
        y[i] ~ Normal(mean = dot(c, x[i]), precision = 10.0)
        x_prev = x[i]
    end
end

ARVMP(Multivariate, order, ARsafe()) is the algorithm of each AR node: a multivariate state of length order, with the numerically safe computation of the joint over the states. The node's rules are variational, so the model needs a factorization and initial marginals:

θ_real = [0.5, -0.3]
series = zeros(102)
rng    = StableRNG(42)
for t in 3:102
    series[t] = θ_real[1] * series[t - 1] + θ_real[2] * series[t - 2] + randn(rng)
end
dataset = series[3:end] .+ randn(rng, 100) ./ sqrt(10.0)

ar_constraints = @constraints begin
    q(x, x0, γ, θ) = q(x, x0)q(γ)q(θ)
end

ar_initialization = @initialization begin
    q(γ) = GammaShapeRate(1.0, 1.0)
    q(θ) = MvNormalMeanPrecision(zeros(2), diageye(2))
end

result = infer(
    model          = ar_model(order = 2),
    data           = (y = dataset,),
    constraints    = ar_constraints,
    initialization = ar_initialization,
    iterations     = 20,
)

mean(result.posteriors[:θ][end])

2-element Vector{Float64}:
  0.42590475897265156
 -0.29724154279417836

The posterior mean of the coefficients is close to θ_real. Without an algorithm, no rule of AR applies, and inference stops with an error that says so:

@model function ar_model_without_algorithm(y, order)
    γ  ~ Gamma(shape = 1.0, rate = 1.0)
    θ  ~ MvNormal(mean = zeros(order), precision = diageye(order))
    x0 ~ MvNormal(mean = zeros(order), precision = diageye(order))
    c = [1.0; zeros(order - 1)]
    x_prev = x0
    for i in eachindex(y)
        x[i] ~ AR(x_prev, θ, γ)
        y[i] ~ Normal(mean = dot(c, x[i]), precision = 10.0)
        x_prev = x[i]
    end
end

try
    infer(
        model          = ar_model_without_algorithm(order = 2),
        data           = (y = dataset,),
        constraints    = ar_constraints,
        initialization = ar_initialization,
        iterations     = 20,
    )
catch err
    showerror(stdout, err)
end

RuleNotFoundError: no message rule for AR towards :y under DefaultAlgorithm() takes the inputs (m[:x]::MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}}, q[:γ]::GammaShapeRate{Float64}, q[:θ]::MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}})
  rules of this shape exist under another algorithm (ARVMP)
  what to try: give the node the algorithm these rules are defined for
  near misses:
    rule at ~/.julia/packages/AutoregressiveMessagePassingRules/5n5Ch/src/autoregressive.jl:304
      ✗ algorithm ARVMP
      ✓ m[:x]::Union{NormalDistributionsFamily{T}, GaussianDistributionsFamily{T}} where T  got MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}}
      ✓ q[:θ]::Any  got MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}}
      ✓ q[:γ]::Any  got GammaShapeRate{Float64}
    rule at ~/.julia/packages/AutoregressiveMessagePassingRules/5n5Ch/src/autoregressive.jl:310
      ✗ algorithm ARVMP
      ✗ q[:x]::Any  not provided
      ✓ q[:θ]::Any  got MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}}
      ✓ q[:γ]::Any  got GammaShapeRate{Float64}
      ✗ m[:x]::MvNormalMeanPrecision{Float64, Vector{Float64}, Matrix{Float64}}  provided but not consumed

The @algorithm macro

@algorithm gives algorithms to nodes from outside the model. Each line selects nodes by their function and, optionally, by the variables they connect, and gives the algorithm after ->:

ar_algorithm = @algorithm begin
    AR() -> ARVMP(Multivariate, 2, ARsafe())
end

result = infer(
    model          = ar_model_without_algorithm(order = 2),
    data           = (y = dataset,),
    constraints    = ar_constraints,
    initialization = ar_initialization,
    algorithm      = ar_algorithm,
    iterations     = 20,
)

mean(result.posteriors[:θ][end])

2-element Vector{Float64}:
  0.42590475897265156
 -0.29724154279417836

AR() selects every AR node of the model. The result is the same as with where { algorithm = ... }. An algorithm given in the model with where takes precedence over one from @algorithm.

@algorithm also defines a function that returns a specification, so that the algorithm can depend on arguments:

@algorithm function ar_algorithm_of_order(order)
    AR() -> ARVMP(Multivariate, order, ARsafe())
end

ar_algorithm = ar_algorithm_of_order(2)

@algorithm begin
    f(x, y) -> DeltaApproximation(method = Unscented())
    NormalMixture() -> NormalMixtureVMP()
end

Selecting nodes by their variables

The GCV node lives in the GCVMessagePassingRules package, which RxInfer loads too. It models an observation whose log-variance is linear in a latent variable, $y \sim \mathcal{N}(x, \exp(\kappa z + \omega))$:

@model function volatility_model(y, κ, ω)
    x ~ Normal(mean = 0.0, variance = 10.0)
    z ~ Normal(mean = 0.0, variance = 1.0)
    for i in eachindex(y)
        y[i] ~ GCV(x, z, κ, ω)
    end
end

Its default algorithm, GCVApproximation(), uses a Gauss–Hermite cubature with 20 points. The entry GCV(x, z) -> ... selects the GCV nodes connected to both x and z, and gives them a cubature with 32 points:

volatility_algorithm = @algorithm begin
    GCV(x, z) -> GCVApproximation(method = GaussHermiteCubature(32))
end

volatility_initialization = @initialization begin
    q(x) = NormalMeanVariance(0.0, 10.0)
    q(z) = NormalMeanVariance(0.0, 1.0)
end

volatility_data = 1.0 .+ sqrt(exp(0.5)) .* randn(StableRNG(1), 200)

result = infer(
    model          = volatility_model(κ = 1.0, ω = 0.0),
    data           = (y = volatility_data,),
    constraints    = MeanField(),
    initialization = volatility_initialization,
    algorithm      = volatility_algorithm,
    iterations     = 20,
)

(x = mean(result.posteriors[:x][end]), z = mean(result.posteriors[:z][end]))

(x = 0.9370954784155182, z = 0.5006426759837177)

The data have mean 1 and log-variance 0.5, which the posteriors of x and z recover. A specification holds as many entries as you need, one per line:

combined_algorithm = @algorithm begin
    GCV(x, z) -> GCVApproximation(method = GaussHermiteCubature(32))
    AR() -> ARVMP(Multivariate, 2, ARsafe())
end

Algorithms for nodes in submodels

As in @constraints, a for meta in submodel block gives algorithms to the nodes of a submodel:

@model function noisy_ar_step(y, x_next, x_prev, θ, γ, order)
    x_next ~ AR(x_prev, θ, γ)
    y ~ Normal(mean = dot([1.0; zeros(order - 1)], x_next), precision = 10.0)
end

@model function ar_model_with_submodel(y, order)
    γ  ~ Gamma(shape = 1.0, rate = 1.0)
    θ  ~ MvNormal(mean = zeros(order), precision = diageye(order))
    x0 ~ MvNormal(mean = zeros(order), precision = diageye(order))
    x_prev = x0
    for i in eachindex(y)
        x[i] ~ noisy_ar_step(y = y[i], x_prev = x_prev, θ = θ, γ = γ, order = order)
        x_prev = x[i]
    end
end

submodel_algorithm = @algorithm begin
    for meta in noisy_ar_step
        AR() -> ARVMP(Multivariate, 2, ARsafe())
    end
end

result = infer(
    model          = ar_model_with_submodel(order = 2),
    data           = (y = dataset,),
    constraints    = ar_constraints,
    initialization = ar_initialization,
    algorithm      = submodel_algorithm,
    iterations     = 20,
)

mean(result.posteriors[:θ][end])

2-element Vector{Float64}:
  0.42590475897265156
 -0.29724154279417836

The block is written for meta in … in both macros. The entries inside it apply only to the nodes of noisy_ar_step, so the same node function can run under different algorithms in different submodels.

A custom algorithm

An algorithm is a type, and its fields are the parameters of its rules. You can define one for a node you did not write, and give it rules of your own. This section tempers the likelihood of a Gaussian observation: it raises $\mathcal{N}(y \mid x, v)$ to a power $\beta \in (0, 1]$, which widens its message towards the mean,

\[\mathcal{N}(y \mid x, v)^{\beta} \propto \mathcal{N}(x \mid y, v / \beta).\]

A subtype of DefaultAlgorithmExtension overrides some rules of the default algorithm and inherits the others:

using RxInfer

struct Tempered{T} <: DefaultAlgorithmExtension
    β::T
end

@define_message_update_rule(
    node = NormalMeanVariance,
    target = :μ,
    algorithm = Tempered,
    args = (q[:out]::PointMass, q[:v]::PointMass),
    body = (algo, args) -> NormalMeanVariance(mean(args.q[:out]), mean(args.q[:v]) / algo.β),
)

@call_message_update_rule(
    node = NormalMeanVariance, target = :μ,
    q = (out = PointMass(4.0), v = PointMass(2.0)),
    algorithm = Tempered(0.5),
)

(algo, args)->NormalMeanVariance(mean(args.q[:out]), mean(args.q[:v]) / algo.β)

The body names the algo slot before args, and reads $\beta$ from it. The rule takes marginals because y is observed and v is a constant: each is a cluster of its own. The tutorial on algorithms of MessagePassingRulesBase covers algorithms that stand alone and declare what their rules take.

@algorithm gives the new algorithm to the node connected to y:

@model function gaussian_model(y)
    x ~ NormalMeanVariance(2.5, 0.5)
    y ~ NormalMeanVariance(x, 2.0)
end

tempered = @algorithm begin
    NormalMeanVariance(y) -> Tempered(0.5)
end

result          = infer(model = gaussian_model(), data = (y = 4.0,))
result_tempered = infer(model = gaussian_model(), data = (y = 4.0,), algorithm = tempered)

(exact = mean_var(result.posteriors[:x]), tempered = mean_var(result_tempered.posteriors[:x]))

(exact = (2.8000000000000003, 0.4), tempered = (2.6666666666666665, 0.4444444444444444))

The tempered posterior stays closer to the prior mean, 2.5, and has a larger variance. Its closed form is the product of the prior with $\mathcal{N}(x \mid 4, 2 / 0.5)$:

mean_var(prod(PreserveTypeProd(Distribution), NormalMeanVariance(2.5, 0.5), NormalMeanVariance(4.0, 4.0)))

(2.6666666666666665, 0.4444444444444444)

The prior's node, x ~ NormalMeanVariance(2.5, 0.5), keeps the default algorithm, because the entry selects only the node connected to y.

Algorithms that name variables

Most algorithms only carry parameters. Some also give the variables of a node roles that its rules must know. A planning algorithm for the transition x[t + 1] ~ DiscreteTransition(x[t], A, u[t], θ) must know that u[t] is the control and that θ is a parameter it holds fixed. The rules know the node's interfaces, not the model's variables: u[t] is the first member of the group T, (:T, 1), and θ its second, (:T, 2). Written as keys, the roles depend on the order of the node's inputs, and a model that writes DiscreteTransition(x[t], A, θ, u[t]) silently swaps them. An algorithm can name the variables instead, and RxInfer resolves them to the keys of the interfaces they connect to before the node runs. The algorithm opts in with a method of RxInfer.resolve_variables:

struct Planning{C, P} <: DefaultAlgorithmExtension
    control::C
    parameter::P
end
Planning(; control, parameter) = Planning(control, parameter)

RxInfer.resolve_variables(algorithm::Planning, interface) =
    Planning(interface(algorithm.control), interface(algorithm.parameter))

@model function rollout(A, x0, horizon)
    θ ~ Categorical([0.5, 0.5])
    x[1] ~ Categorical(x0)
    for t in 1:horizon
        u[t] ~ Categorical([0.5, 0.5])
        x[t + 1] ~ DiscreteTransition(x[t], A, u[t], θ) where {
            algorithm = Planning(control = u[t], parameter = θ)
        }
    end
end

The rules of each transition run under Planning((:T, 1), (:T, 2)) and read the roles from their algo slot. With the inputs in the other order, they run under Planning((:T, 2), (:T, 1)), and the rules need not change.

Only an algorithm given in the model, with where { algorithm = … }, can name variables: the variables exist there. @algorithm is evaluated outside the model, so an algorithm from it gives the keys themselves, DiscreteTransition() -> Planning((:T, 1), (:T, 2)).

interface resolves the variables of the model only. It returns any other value as it is: a key given explicitly, as above, or a value that is not a variable of the model, whatever its name. After u[t] = 1.0 in a model, Planning(control = u[t], …) gets 1.0, as an algorithm that names no variables does. A variable the node is not connected to is an error.

RxInfer.resolve_variables(algorithm, interface)

RxInfer.resolve_variables(algorithm::Planning, interface) = Planning(interface(algorithm.control))