Form constraints

A form constraint gives a variable's marginal a chosen functional form. The variable multiplies the messages that arrive at it, and the constraint approximates their normalised product by a distribution of that form. With the constraint as $f$, the marginal is

\[q(x) = f\left(\frac{\overrightarrow{\mu}(x)\overleftarrow{\mu}(x)}{\int \overrightarrow{\mu}(x)\overleftarrow{\mu}(x) \mathrm{d}x}\right).\]

A form constraint keeps a posterior tractable where the product has no closed form, or keeps it in a parameterisation the model needs. The product itself is the prod function of BayesBase.

The examples use the Gaussian node of The example node.

The interface

A variable receives its form constraint through its ReactiveMP.MessageProductContexts, as the keywords form_constraint, form_constraint_check_strategy and prod_constraint. RxInfer builds them from a model's constraints. A constraint implements the functions below, and the example implements all of them.

The supertype

A form constraint subtypes AbstractFormConstraint. UnspecifiedFormConstraint, the default, leaves the product as it is, and CompositeFormConstraint applies several constraints in turn.

ReactiveMP.AbstractFormConstraint — Type

AbstractFormConstraint

The supertype of form constraints, which give a product of messages a chosen functional form (see constrain_form). A constraint of another type works too, wrapped by ReactiveMP.preprocess_form_constraints in a ReactiveMP.WrappedFormConstraint; subtyping it lets constraints be combined with + into a CompositeFormConstraint.

A constraint implements constrain_form, default_form_check_strategy and default_prod_constraint.

source

ReactiveMP.UnspecifiedFormConstraint — Type

UnspecifiedFormConstraint()

The form constraint that constrains nothing, returning the distribution as it is: the default of ReactiveMP.MessageProductContext. Its check strategy is FormConstraintCheckLast and its product strategy BayesBase.GenericProd().

source

ReactiveMP.CompositeFormConstraint — Type

CompositeFormConstraint(constraints::Tuple)

Several form constraints applied in order, the output of each the input of the next. Combine constraints with +, a + b, or give a tuple to ReactiveMP.preprocess_form_constraints. Its product strategy resolves its parts'.

Throws

ErrorException from default_form_check_strategy when its parts' default check strategies differ.

source

ReactiveMP.preprocess_form_constraints — Function

ReactiveMP.preprocess_form_constraints(constraints)

Form constraints as the engine takes them: an AbstractFormConstraint as it is, a tuple as a CompositeFormConstraint of its elements, each preprocessed, and any other value wrapped in a ReactiveMP.WrappedFormConstraint with its ReactiveMP.prepare_context.

source

Form check strategy

A form constraint implements default_form_check_strategy, which says when it applies:

FormConstraintCheckLast applies it once, to the whole product: $q(x) = f(\mu_1(x) \mu_2(x) \mu_3(x))$;
FormConstraintCheckEach applies it after each pairwise product: $q(x) = f(f(\mu_1(x) \mu_2(x)) \mu_3(x))$.

ReactiveMP.default_form_check_strategy — Function

default_form_check_strategy(form_constraint)

The check strategy form_constraint applies with by default, FormConstraintCheckEach or FormConstraintCheckLast. A constraint implements it; RxInfer calls it to build a variable's ReactiveMP.MessageProductContext, and a CompositeFormConstraint requires its parts to agree.

source

ReactiveMP.FormConstraintCheckEach — Type

FormConstraintCheckEach()

The check strategy that applies the form constraint after each pairwise product, inside ReactiveMP.compute_product_of_two_messages: f(f(μ₁ μ₂) μ₃). It keeps every intermediate product in the constrained form, where the product would otherwise grow or lose its closed form.

source

ReactiveMP.FormConstraintCheckLast — Type

FormConstraintCheckLast()

The check strategy that applies the form constraint once, to the whole product, at the end of ReactiveMP.compute_product_of_messages: f(μ₁ μ₂ μ₃). It is the default of ReactiveMP.MessageProductContext, and the cheaper where the intermediate form does not matter.

The product strategy

A form constraint implements default_prod_constraint, the strategy it needs BayesBase.prod to multiply with. GenericProd() computes a closed-form product where one exists, and otherwise keeps the two sides as a BayesBase.ProductOf.

ReactiveMP.default_prod_constraint — Function

default_prod_constraint(form_constraint)

The product strategy for BayesBase.prod that form_constraint needs, such as BayesBase.GenericProd(), which keeps a product with no closed form as a ProductOf for the constraint to approximate. A constraint implements it; RxInfer calls it to build a variable's ReactiveMP.MessageProductContext, and a CompositeFormConstraint resolves its parts' with BayesBase.resolve_prod_strategy.

source

The constraint

constrain_form is the $f$ above: it takes the product and returns a distribution of the chosen form. A constraint that changes the product, returning something other than it was given, leaves the product's log scale undefined.

ReactiveMP.constrain_form — Function

constrain_form(constraint, distribution)
constrain_form(constraint, context, distribution)

distribution in the form constraint imposes: the f of q(x) = f(μ₁(x) μ₂(x)). A constraint implements it for the distributions it supports; the second form is called for a ReactiveMP.WrappedFormConstraint with the context its ReactiveMP.prepare_context prepared. A constraint that returns its input unchanged keeps the product's log scale; one that returns something else leaves it undefined.

source

An example

The constraint below keeps a marginal a normal distribution in the mean-precision parameterisation:

struct MeanPrecisionFormConstraint <: AbstractFormConstraint end

ReactiveMP.default_form_check_strategy(::MeanPrecisionFormConstraint) = FormConstraintCheckLast()
ReactiveMP.default_prod_constraint(::MeanPrecisionFormConstraint) = GenericProd()

ReactiveMP.constrain_form(::MeanPrecisionFormConstraint, distribution) =
    NormalMeanPrecision(mean(distribution), precision(distribution))

constraint = ReactiveMP.preprocess_form_constraints(MeanPrecisionFormConstraint())
constrain_form(constraint, NormalMeanVariance(0.0, 2.0))

ExponentialFamily.NormalMeanPrecision{Float64}(μ=0.0, w=0.5)

The first method needs a distribution with a mean and a precision. A product with no closed form is a ProductOf, which has neither. The second method matches its moments on a grid, using the product's unnormalised log-density:

function ReactiveMP.constrain_form(::MeanPrecisionFormConstraint, product::ProductOf)
    left = product.left
    xs = range(mean(left) - 10 * std(left), mean(left) + 10 * std(left); length = 2001)
    weights = exp.(logpdf.(Ref(product), xs))
    weights ./= sum(weights)
    m = sum(weights .* xs)
    v = sum(weights .* abs2.(xs .- m))
    return NormalMeanPrecision(m, inv(v))
end

constrain_form(constraint, prod(GenericProd(), NormalMeanVariance(0.0, 1.0), Laplace(1.0, 0.5)))

ExponentialFamily.NormalMeanPrecision{Float64}(μ=0.7312397164755419, w=3.3355747105101963)

The grid spans ten standard deviations of the left side, so the method assumes the left side covers the product's mass.

A variable applies the constraint to its marginal when its marginal's product context carries it:

c = MeanPrecisionFormConstraint()
marginal_context = MessageProductContext(;
    form_constraint = c,
    prod_constraint = default_prod_constraint(c),
    form_constraint_check_strategy = default_form_check_strategy(c),
)

x, y = randomvar(label = :x), datavar(label = :y)
prior = factornode(Gaussian, [(:out, x), (:μ, constvar(0.0)), (:v, constvar(10.0))])
likelihood = factornode(Gaussian, [(:out, y), (:μ, x), (:v, constvar(1.0))])
activate!(x, RandomVariableActivationOptions(nothing, MessageProductContext(), marginal_context))
activate!(y, DataVariableActivationOptions())
foreach(n -> activate!(n, FactorNodeActivationOptions(; logscales = true)), (prior, likelihood))

subscription = subscribe!(get_stream_of_marginals(x), (q) -> println(q))
new_observation!(y, 2.0)
unsubscribe!(subscription)

Marginal(ExponentialFamily.NormalMeanPrecision{Float64}(μ=1.8181818181818181, w=1.1)) with logscale = UndefinedLogScale(:form_constraint)

The posterior is a NormalMeanPrecision. The constraint changed the product's type, so its log scale is undefined, and says it was the form constraint.

Wrapped form constraints

A constraint need not subtype AbstractFormConstraint, as when another package defines it or it subtypes another abstract type. ReactiveMP.preprocess_form_constraints wraps such an object in a ReactiveMP.WrappedFormConstraint, which makes it a form constraint.

A wrapped constraint may implement ReactiveMP.prepare_context. The engine computes its result once and stores it with the constraint. constrain_form then receives three arguments: the constraint, the context and the distribution.

ReactiveMP.WrappedFormConstraint — Type

ReactiveMP.WrappedFormConstraint(constraint, context)

A form constraint of a type that does not subtype AbstractFormConstraint, with the context its ReactiveMP.prepare_context returned. ReactiveMP.preprocess_form_constraints builds it. constrain_form, default_form_check_strategy and default_prod_constraint forward to constraint, the first with the context when there is one.

source

ReactiveMP.prepare_context — Function

ReactiveMP.prepare_context(constraint)

A context for constraint, computed once when it is wrapped and passed to every constrain_form(constraint, context, distribution) call: a random number generator, a precomputed quadrature, and so on. By default there is none, and constrain_form(constraint, distribution) is called. A constraint that needs one adds a method.

source

The constraint below replaces a distribution by ten samples from it, drawn with the random number generator its context holds:

struct SampleListFormConstraint end

ReactiveMP.default_form_check_strategy(::SampleListFormConstraint) = FormConstraintCheckLast()
ReactiveMP.default_prod_constraint(::SampleListFormConstraint) = GenericProd()
ReactiveMP.prepare_context(::SampleListFormConstraint) = Xoshiro(42)
ReactiveMP.constrain_form(::SampleListFormConstraint, rng, distribution) = rand(rng, distribution, 10)

sampler = ReactiveMP.preprocess_form_constraints(SampleListFormConstraint())
constrain_form(sampler, Normal(0.0, 10.0))

10-element Vector{Float64}:
  -3.6335748145177753
   2.517372155742292
  -3.1498797116895605
  -3.1125240132442067
   8.163067649323274
   4.76738379831878
  -8.595553820616212
 -14.692882055065464
  -2.0661311448119264
  -3.1074387308373415