mixle.ppl.inference module

Bayesian inference for mixle.ppl: parameter MCMC, MAP, VI, and closed-form conjugate updates.

A model whose parameter slots hold distributions (priors) or free defines a joint log p(data | theta) + log p(theta). MAP (maximize), MCMC/HMC/NUTS (sample), and VI all run on the exact same target, scored with the existing vectorized seq_log_density and mixle’s mixle.inference.mcmc kernels — no new inference engine.

What this module supports:

  • Flat Sample models (Normal(Normal(0, 10), free)) and composite models (mixtures, sequences): composites collect their leaf free/prior parameters across the tree and rebuild a concrete model per evaluation (_collect_composite).

  • Hierarchical priors – a prior whose own hyperparameter is another random variable (Normal(0, tau) with tau estimated), scored differentiably via parent-slot substitution.

  • Non-centered reparameterization (.noncentered()) for funnel-prone location-scale priors.

  • Grouped / random-intercept (“plate”) models – Normal(Normal(mu, tau).each(), sigma) sampled jointly over the hyperparameters and every per-group latent.

Identifiability is handled automatically: mixture chains are relabeled (sorted by a leading parameter) before pooling, so label-switching no longer needs a user-supplied ordering constraint.

Section map (banners below): core slot/target construction; init + result assembly (finalize, convergence diagnostics, relabeling); parallel chains; inequality/region constraints; hierarchical & grouped models; sampler setup and the how= drivers; closed-form conjugate / conjugate-mixture / hierarchical-EM updates.

class Posterior(slots, value_samples, raw)[source]

Bases: object

Parameter-posterior result attached to a fitted RV’s .result.

Holds value-space draws and the raw MCMCResult. Look up a parameter by its RandomVariable handle, its name, or its slot index.

Parameters:

  • slots (list[_Slot])

  • value_samples (np.ndarray)

  • raw (Any)

pointwise_log_likelihood(data)[source]

Return the (n_draws, n_obs) log-likelihood of data under each posterior draw.

This is the input to the predictive model-comparison diagnostics (WAIC, PSIS-LOO).

Return type:

ndarray

samples(param=None)[source]
Return type:

ndarray

mean(param=None)[source]
summary()[source]
Return type:

dict

ensemble_fit(rv, data, *, draws=1500, burn=500, thin=1, walkers=None, constraints=None, penalty=None, potentials=None, rng=None, chains=1, parallel=False, missing='error')[source]

Affine-invariant ensemble MCMC (Goodman & Weare stretch move).

A population of walkers samples jointly with no per-dimension step tuning; it is invariant to affine rescalings, so it mixes well on correlated / poorly-scaled posteriors and gives very high ESS/sec on low/medium-dimensional models (no JIT-compile latency). Each draws sweep contributes all walkers states, so the pooled posterior has draws*walkers near-independent samples. Uses the fast NumPy scalar log-target (one eval per proposal). chains>1 runs independent ensembles for Gelman-Rubin R-hat / pooled ESS (parallel spreads them over a process pool).

Parameters:

  • rv (RandomVariable)

  • draws (int)

  • burn (int)

  • thin (int)

  • walkers (int | None)

  • chains (int)

  • parallel (bool)

  • missing (str)

Return type:

RandomVariable

mcmc_fit(rv, data, *, draws=2000, burn=1000, thin=1, scale=None, rng=None, chains=1, parallel=False, constraints=None, penalty=None, potentials=None, missing='error')[source]
Parameters:
Return type:

RandomVariable

hmc_fit(rv, data, *, draws=1000, burn=500, step_size=None, num_steps=15, thin=1, rng=None, chains=1, parallel=False, constraints=None, penalty=None, potentials=None, missing='error')[source]

Hamiltonian Monte Carlo over the parameter posterior.

Uses mixle’s hamiltonian_monte_carlo with a numerical gradient of the joint log-target and a diagonal mass matrix preconditioned to the data-informed posterior scale, so trajectories are well-conditioned without manual tuning. Inequality constraints truncate the posterior (trajectories leaving the region are rejected); for hard constraints how='ensemble' or 'mcmc' usually mixes better.

Parameters:

  • rv (RandomVariable)

  • draws (int)

  • burn (int)

  • step_size (float | None)

  • num_steps (int)

  • thin (int)

  • chains (int)

  • parallel (bool)

  • missing (str)

Return type:

RandomVariable

nuts_fit(rv, data, *, draws=1000, burn=1000, thin=1, target_accept=0.8, max_tree_depth=10, rng=None, chains=1, parallel=False, constraints=None, penalty=None, potentials=None, missing='error')[source]

No-U-Turn Sampler over the parameter posterior — auto-tuned HMC (trajectory length + dual-averaging step size), the default choice for correlated / higher-dimensional posteriors.

Uses the analytic Torch gradient when available (else a numeric gradient), preconditioned by a diagonal mass matrix at the data-informed scale. warmup (= burn) adapts the step size to target_accept. Inequality constraints truncate the posterior; penalty adds a smooth penalty (which enters the gradient via the numeric-gradient path).

Parameters:

  • rv (RandomVariable)

  • draws (int)

  • burn (int)

  • thin (int)

  • target_accept (float)

  • max_tree_depth (int)

  • chains (int)

  • parallel (bool)

  • missing (str)

Return type:

RandomVariable

sample_fit(rv, data, **kw)[source]

Draw the parameter posterior with an automatically chosen sampler (how='sample') — so the user doesn’t pick among mcmc/hmc/nuts/ensemble. Low/medium-dimensional models use the affine-invariant ensemble (highest ESS/sec, handles constraints by rejection); larger models use nuts (auto-tuned HMC for correlated, higher-dimensional posteriors). Takes the common knobs (draws, burn, thin, rng, chains, parallel, constraints, penalty).

Parameters:

rv (RandomVariable)

Return type:

RandomVariable

map_fit(rv, data, *, rng=None, constraints=None, penalty=None, potentials=None, missing='error')[source]
Parameters:

rv (RandomVariable)

Return type:

RandomVariable

laplace_fit(rv, data, *, rng=None, draws=2000, missing='error', **_)[source]

Gaussian (Laplace) posterior approximation at the MAP – the cheap uncertainty rung above map.

Finds the posterior mode, takes the Gaussian whose precision is the Hessian of the negative joint log-density there, and draws from it in the unconstrained space (so positivity/simplex supports are respected after the link transform). Returns a full Posterior (summary / intervals / predictive), so how='laplace' gives credible intervals where how='map' gives only a point.

Parameters:

  • rv (RandomVariable)

  • draws (int)

Return type:

RandomVariable

vi_fit(rv, data, *, samples=4000, mc=16, max_iter=4000, steps=600, lr=0.05, batch_size=None, family='meanfield', alpha=1.0, rng=None, missing='error')[source]

Variational Bayes (ADVI) — a Gaussian variational posterior fit by reparameterized-MC Adam.

family='meanfield' (diagonal q, default) or 'fullrank' (full covariance via a Cholesky factor, capturing posterior correlations). alpha selects the tilted Renyi objective: alpha=1 is the KL-ELBO (default), alpha=0 the importance-weighted (IWAE) bound, and alpha<1 is mass-covering (widens the often-too-narrow KL fit). batch_size subsamples the data per step (SGVB). Without Torch it falls back to derivative-free mean-field ELBO. Works for non-conjugate priors; returns a variational Posterior with draws and posterior-predictive.

Parameters:
Return type:

RandomVariable

class ConjugatePosterior(post)[source]

Bases: object

Exact closed-form posterior over a conjugate parameter.

post maps parameter name -> {mean, sample(n, rng), name, hyper}. This is the ideal case VB approximates: exact, instant, no iteration.

Parameters:

post (dict)

samples(param=None, n=4000, rng=None)[source]
Parameters:

n (int)

mean(param=None)[source]
summary()[source]
Return type:

dict

conjugate_spec(rv)[source]

Return (builder, prior_slot_index, prior_rv) if exactly one slot is a conjugate prior and every other slot is a fixed constant; else None.

Parameters:

rv (RandomVariable)

stats_conjugate_supported(rv)[source]

True if the general exp-family conjugate bridge can fit rv (used by how=’auto’ routing).

Return type:

bool

conjugate_fit(rv, data, *, prior=None, **_)[source]
Parameters:

rv (RandomVariable)

Return type:

RandomVariable

class ConjugateMixturePosterior(entries, weights, param_name)[source]

Bases: object

Exact posterior for a mixture-of-conjugate-priors model.

The posterior is again a mixture: the per-component conjugate posteriors with weights reweighted by each component’s marginal likelihood, w'_k ∝ w_k · m_k. Sampling draws a component by w' then samples that component’s conjugate posterior.

mean(param=None)[source]
samples(param=None, n=4000, rng=None)[source]
Parameters:

n (int)

summary()[source]
Return type:

dict

conjugate_mixture_spec(rv)[source]

Return (builder, logm, slot_index, component_rvs, prior_weights) when exactly one slot is a Mix of conjugate priors (all forming the same registered conjugate pair) and every other slot is a fixed constant; else None.

Parameters:

rv (RandomVariable)

conjugate_mixture_fit(rv, data)[source]
Parameters:

rv (RandomVariable)

Return type:

RandomVariable

class HierarchicalPosterior(group_means, group_vars, hyper)[source]

Bases: object

Per-group posteriors q(mu_i) = Normal(group_means[i], group_vars[i]) plus the fitted hyperparameters of a Normal-Normal random-effects model.

samples(param=None)[source]
summary()[source]
Return type:

dict

hierarchical_fit(rv, data, *, max_its=300, tol=1e-8)[source]

Conjugate hierarchical (random-effects) EM, dispatched by conjugate pair.

Supports Normal-Normal (exact), Gamma-Poisson, and Beta-Bernoulli. data is a list of groups; returns the fitted population distribution plus per-group posteriors.

Parameters:

  • rv (RandomVariable)

  • max_its (int)

  • tol (float)

Return type:

RandomVariable

class IndexedPosterior(latents, scalars)[source]

Bases: object

Result of an indexed-latent fit: the fitted latent vector(s) and scalar parameters.

latents maps each free(K) vector’s name to its fitted K-vector; group_means aliases the single-vector case (mirrors HierarchicalPosterior).

Parameters:
samples(param=None)[source]
summary()[source]
Return type:

dict

indexed_fit(rv, data, *, given=None, how='map', rng=None, draws=2000, burn=1000, thin=1, **_)[source]

Fit a flat model with a data-indexed latent vector (theta[Field("g")]).

how='map' (the default) maximizes the per-observation joint and returns the fitted vectors on .result (latents / group_means). how='mcmc' samples the joint by adaptive random-walk Metropolis and returns a full Posterior (per-latent summary / credible intervals). Other how values (hmc/nuts) are not yet wired for the indexed target.

Parameters:

rv (RandomVariable)

Return type:

RandomVariable