Fit a mixed model with automatic coordinate transforms

Build a Gaussian regression with a shared slope and a varying intercept for each site. We will choose proper priors, inspect the model’s parameters, sample with BlackJAX NUTS, check the chains, and predict for new rows.

FlatPPL compiles both the density and its coordinate transforms to StableHLO. Positive scales become free coordinates for NUTS, with their density correction included automatically. No hand-written transforms or Jacobians are needed.

Use the Installation instructions with BlackJAX and Polars installed. The Regression formulas reference lists formula terms, families and prior selectors; Coordinate transforms for inference covers the compiled coordinate interface.

Runtime budget: 60 s.

from typing import Any

import blackjax
import jax
import jax.numpy as jnp
import numpy as np
import polars as pl

from flatppl import flatppl, from_formula

Choose priors in model units

y ~ x + (1 | site) gives each site an intercept drawn from a shared population. The default response family is Gaussian. These eight observations make a small teaching example; two sites provide little information about the population variation between sites.

The explicit priors below regularize the intercept, slope, site variation and observation noise. FlatPPL’s brms-compatible default slope prior is improper Lebesgue(reals); choosing a proper Normal also makes this joint prior sampleable. Positive scale domains turn the two normal(0, 1) scale priors into normalized half-Normals.

data = pl.DataFrame(
    {
        "x": [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0],
        "y": [2.0, 3.0, 6.0, 7.0, 3.0, 6.0, 5.0, 8.0],
        "site": ["A"] * 4 + ["B"] * 4,
    }
)
model = from_formula(
    "y ~ x + (1 | site)",
    data,
    priors=[
        {"class": "Intercept", "prior": "normal(1, 3)"},
        {"class": "b", "prior": "normal(0, 2)"},
        {"class": "sd", "group": "site", "prior": "normal(0, 1)"},
        {"class": "sigma", "prior": "normal(0, 1)"},
    ],
)
print("Design columns:", model.formula.info["design_columns"])
groups = model.formula.recipe["predictors"]["mu"]["groups"]
print("Site order:", [level["value"] for level in groups[0]["levels"]])
sampled = [
    parameter
    for parameter in model.formula.info["parameters"]
    if parameter["role"] == "sampled"
]
print("Sampled parameters:", [parameter["name"] for parameter in sampled])

Design columns: ('x',)
Site order: ['A', 'B']
Sampled parameters: ['Intercept', 'b_x', 'sd_site__Intercept', 'z_1[1,1]', 'z_1[1,2]', 'sigma']

Identify the model coordinates

model.source contains the generated FlatPPL. Its Pi_post is the prior times the likelihood, without an evidence normalizer. NUTS only needs that unnormalized posterior density.

The metadata maps readable parameter names to FlatPPL binding names and one-based index or indices. A parameter’s role distinguishes sampled coordinates from fixed values and derived summaries. Here the sampled coordinates are a centred intercept, a length-one slope vector, two scales, and two standard Normal site coordinates in the printed site order.

temp_Intercept describes the mean at the training mean of x (2.5). The site effects use the non-centred parameterization r_site = sd_site * z_site. b_Intercept is a derived intercept at x = 0 and is not another sampling coordinate.

Registering the model makes it importable by an explicit density query. Its input domains describe model coordinates, including positive scales; Rust uses these same domains to choose the sampling transforms.

context = model.context
context.register("training.flatppl", model)
query = flatppl(
    """
    m = load_module("training.flatppl")
    p = elementof(cartprod(temp_Intercept = reals, b = cartpow(reals, 1),
                          sd_site = posreals, z_site = cartpow(reals, 2), sigma = posreals))
    inputs = p
    outputs = logdensityof(m.Pi_post, p)
    """,
    context=context,
)

Compile the density and its transforms together

compile_transformed(logdensity=True) returns four compiled callables: constrained evaluates the original query, unconstrained evaluates the density in free coordinates, unconstrain maps a model point to free coordinates, and constrain maps a free point back.

For this model, positive scales use log coordinates and the other coordinates stay unchanged. The unconstrained density includes the change-of-variables correction once. Both the transforms and density execute in StableHLO; Python performs no coordinate conversion or Jacobian calculation inside the sampler loop. Compilation defaults to float32; see Coordinate transforms for inference for float64 and the supported domains.

target = query.compile_transformed(logdensity=True)
point = {
    "temp_Intercept": 1.0,
    "b": jnp.array([2.0]),
    "sd_site": 0.5,
    "z_site": jnp.array([-0.2, 0.4]),
    "sigma": 1.5,
}
score, gradient = jax.jit(jax.value_and_grad(target.constrained))(point)
print("Model-coordinate log density:", float(score))
print("Gradient:", jax.tree.map(lambda value: np.asarray(value).tolist(), gradient))
initial = target.unconstrain(point)

Model-coordinate log density: -47.65321350097656
Gradient: {'b': [-2.2777774333953857], 'sd_site': 1.277777910232544, 'sigma': 32.21852111816406, 'temp_Intercept': 14.044445037841797, 'z_site': [3.3999998569488525, 3.422222137451172]}

Sample with BlackJAX

The starting point is inside every domain. Transforms do not clip invalid initial values or repair points on a boundary.

Each chain adapts its step size and mass matrix for 1,000 warmup steps, then keeps 2,000 posterior draws. Warmup draws are discarded. Four independent random keys drive four chains, vectorized with jax.vmap; jax.lax.scan keeps the sampling loop compiled. BlackJAX owns the sampler and adaptation.

warmup = blackjax.window_adaptation(blackjax.nuts, target.unconstrained)


@jax.jit
def sample_chain(key: jax.Array):
    """Adapt one NUTS chain, then collect its draws and divergence flags."""
    warmup_key, sample_key = jax.random.split(key)
    (state, parameters), _ = warmup.run(warmup_key, initial, num_steps=1000)
    sampler = blackjax.nuts(target.unconstrained, **parameters)

    def step(current: Any, draw_key: jax.Array):
        updated, info = sampler.step(draw_key, current)
        return updated, (updated.position, info.is_divergent)

    return jax.lax.scan(step, state, jax.random.split(sample_key, 2000))[1]


keys = jax.random.split(jax.random.key(827), 4)
free_draws, divergences = jax.vmap(sample_chain)(keys)
draws = jax.vmap(jax.vmap(target.constrain))(free_draws)

Check the chains in model units

free_draws has leading (chain, draw) axes. The two nested vmap calls above apply the compiled constrain function across both axes, so scale summaries below are in the original units.

R-hat compares variation within and between chains; values near one indicate agreement. Effective sample size (ESS) measures how much information the correlated draws provide. Check both together with divergences before using the summaries. For poor mixing or divergences, check the parameterization, priors and sampler tuning. Longer runs reduce Monte Carlo uncertainty once sampling behaves well.

parameters = {
    "Centred intercept": draws["temp_Intercept"],
    "b_x": draws["b"][..., 0],
    "sd_site": draws["sd_site"],
    "sigma": draws["sigma"],
}
for name, chains in parameters.items():
    interval = np.quantile(chains, [0.05, 0.95])
    print(
        f"{name}: mean={float(jnp.mean(chains)):.2f}, "
        f"90% interval=({interval[0]:.2f}, {interval[1]:.2f}), "
        f"ESS={float(blackjax.ess(chains)):.0f}, "
        f"R-hat={float(blackjax.rhat(chains)):.3f}"
    )
print("Divergences:", int(jnp.sum(divergences)))

Centred intercept: mean=4.75, 90% interval=(3.46, 5.80), ESS=1156, R-hat=1.004
b_x: mean=1.56, 90% interval=(1.01, 2.09), ESS=3291, R-hat=1.000
sd_site: mean=0.76, 90% interval=(0.08, 1.81), ESS=1973, R-hat=1.003
sigma: mean=1.04, 90% interval=(0.65, 1.59), ESS=2957, R-hat=1.001
Divergences: 68

This small example can produce divergences even when R-hat is near one. Treat its summaries as a workflow illustration until those divergences are resolved. Inspect the scale geometry and prior sensitivity; more draws alone do not fix a biased exploration of the posterior.

Predict with the training recipe

New rows need predictors but no response column. prepare returns a new module and retains the training mean, site order, contrasts and prior defaults. It does not refit preprocessing on the prediction data or change the training model. Unseen sites are currently refused.

new_data = pl.DataFrame({"x": [4.0, 5.0], "site": ["A", "B"]})
prediction = model.formula.prepare(new_data)
design = prediction.query("X").compile(autodiff=False)()
print("New design rows:", design)  # x minus its training mean, 2.5

New design rows: [[1.5]
 [2.5]]

The prediction module’s K_model is the observation law conditional on model parameters. This query scores the pair of proposed responses together. The first call uses our explicit parameter point; then we evaluate across posterior draws to account for parameter uncertainty.

context.register("prediction.flatppl", prediction)
predictive_density = flatppl(
    """
    m = load_module("prediction.flatppl")
    p = elementof(cartprod(temp_Intercept = reals, b = cartpow(reals, 1),
                          sd_site = posreals, z_site = cartpow(reals, 2), sigma = posreals))
    observed = elementof(cartpow(reals, 2))
    inputs = (p, observed)
    outputs = logdensityof(m.K_model(p), record(y = observed))
    """,
    context=context,
).compile()
observed = jnp.array([5.0, 7.0])
prediction_score = predictive_density(point, observed)
print("New-row log density:", float(prediction_score))


def score_responses(parameters: Any):
    """Score both proposed responses at one posterior parameter draw."""
    return predictive_density(parameters, observed)


log_scores = jax.jit(jax.vmap(jax.vmap(score_responses)))(draws)
predictive_score = jax.scipy.special.logsumexp(log_scores) - jnp.log(log_scores.size)
print("Posterior predictive log density:", float(predictive_score))

New-row log density: -3.0599184036254883
Posterior predictive log density: -4.491494178771973

logsumexp computes the logarithm of the average conditional density stably. This gives a Monte Carlo estimate of the joint posterior predictive log density for [5, 7] at the two new rows, with the sampled site effects included. The retained recipe lets us repeat the calculation for other rows; a changed input shape may require another compilation.