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: _DroppedClosurePickling

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 posterior samples for all parameters or one selected parameter.

Return type:

ndarray

mean(param=None)[source]

Return posterior means for all parameters or one selected parameter.

summary()[source]

Return parameter summaries and any available convergence diagnostics.

Every draw-based mean ships with mcse – its own Monte Carlo standard error, std / sqrt(ESS) (STAT-RR19-11: the primary summary is the claim-bearing surface, and publishing a posterior mean without its noise floor invited reading MC noise as signal; a one-chain fit’s mean 0.338 carried a 0.017 MCSE that appeared nowhere). The ESS is the fit’s own bulk ESS when finite; otherwise it is computed from the pooled draws as one chain – the result’s deliberate single-chain NaN markers are about MIXING being unassessable, not about the ESS being incomputable.

Return type:

dict

class MultiChainResult(chain_results, chain_samples)[source]

Bases: object

Complete raw result for a multi-chain fit without pretending the first chain is the whole run.

validated_potentials(potentials)[source]

Return potentials as a list of potentials, refusing anything else by name.

A string is iterable, so potentials='x' walked straight into the resolution loop and failed with “‘str’ object has no attribute ‘vars’” (P04-F11).

Return type:

list

validated_sampler_rng(rng)[source]

Coerce rng to a RandomState, accepting the spellings predict() already accepts.

The samplers took only a RandomState while predict() took ints and Generator``s, so the same ``rng=0 worked on one call and raised on the next (P04-F11).

Return type:

RandomState

ensemble_fit(rv, data, *, draws=1500, burn=None, 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 to the pooled posterior, but those draws*walkers states are NOT independent: a walker is autocorrelated across sweeps and the walkers move together, so the effective sample size is typically a small fraction of the pooled count (read ess_bulk, not the count). Uses the fast NumPy scalar log-target (one eval per proposal).

The other walkers start around an anchor: the log-target’s maximum over the real location parameters, found from the data-informed start before sampling. Starting every location at one shared data mean left coordinates whose posterior sat far from it unreachable – an 8-parameter location model came back 141 posterior sds off at any burn-in (V01-F01).

Half the walkers start from prior draws, so burn has to pay for contracting that cloud. Left at None it scales with the number of parameters d as max(1500, 300*d); an explicit smaller value is honoured as given, and a burn-in too short to contract the cloud returns a posterior that is too wide and off-centre with a single-chain summary that cannot see it – on an 18-group model burn=500 left rates 8.8 posterior sds from the exact answer (Q04-F11, Q07-F06). chains>1 runs independent ensembles for Gelman-Rubin R-hat / pooled ESS (parallel spreads them over a process pool). That between-chain R-hat is the diagnostic that does flag an uncontracted ensemble – on the same 18-group model, burn=500 with chains=2 reports R-hat 11.2, against 1.03 at the default burn – so use it whenever the budget is not the default.

Parameters:

  • rv (RandomVariable)

  • draws (int)

  • burn (int | None)

  • 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]

Fit a PPL model with adaptive random-walk Metropolis-Hastings.

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' selects among MCMC, HMC, NUTS, and ensemble samplers. Low- and medium-dimensional models use the affine-invariant ensemble sampler; larger models use NUTS for correlated, higher-dimensional posteriors. Common controls include draws, burn, thin, rng, chains, parallel, constraints, and penalty.

Parameters:

rv (RandomVariable)

Return type:

RandomVariable

map_fit(rv, data, *, rng=None, constraints=None, penalty=None, potentials=None, missing='error', max_iter=5000, tol=1e-8)[source]

Fit a PPL model by maximum a posteriori optimization.

The point estimate maximizes log p(data | theta) + log p(theta) in the constrained parameter space – no support-transform log-Jacobian – so with flat priors MAP coincides with the MLE (e.g. the sqrt(S/n) sd for a Normal). The samplers and VI/Laplace keep the Jacobian: they need the unconstrained-space density.

Parameters:

  • rv (RandomVariable)

  • max_iter (int)

  • tol (float)

Return type:

RandomVariable

laplace_fit(rv, data, *, rng=None, draws=2000, missing='error', constraints=None, penalty=None, potentials=None, max_iter=5000, tol=1e-8)[source]

Gaussian (Laplace) posterior approximation at the 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)

  • max_iter (int)

  • tol (float)

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, seed=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. seed makes the fit deterministic when rng is not given (fit(how='vi', seed=...), matching the samplers’ seeding).

Parameters:
Return type:

RandomVariable

class ConjugatePosterior(post)[source]

Bases: _DroppedClosurePickling

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]

Draw samples from the exact conjugate posterior.

Parameters:

n (int)

mean(param=None)[source]

Return posterior means for all conjugate parameters or one parameter.

summary()[source]

Return conjugate posterior family, hyperparameters, and posterior means.

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]

Fit a registered conjugate PPL model and attach its exact posterior.

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]

Return the posterior mixture mean for the conjugate parameter.

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

Draw posterior samples by first sampling a conjugate prior component.

Parameters:

n (int)

summary()[source]

Return posterior mixture weights, component summaries, and mean.

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]

Fit a likelihood with a mixture of conjugate priors exactly.

Parameters:

rv (RandomVariable)

Return type:

RandomVariable

class HierarchicalPosterior(group_means, group_vars, hyper, posterior_family)[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, n=4000, rng=None)[source]

Draw group effects from the declared conditional posterior.

Parameters:

n (int)

summary()[source]

Return fitted hyperparameters and per-group posterior means.

Return type:

dict

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

Empirical-Bayes random-effects fitting with conditional conjugate group posteriors.

Supports Normal-Normal (exact), Gamma-Poisson, and Beta-Bernoulli. data is a list of groups. Declared population parameters initialize the hyperparameter optimization; returned group posterior draws condition on the fitted plug-in hyperparameters and do not include their uncertainty. The result receipt states those semantics explicitly.

Parameters:

  • rv (RandomVariable)

  • max_its (int)

  • tol (float)

Return type:

RandomVariable

class IndexedPosterior(latent_entries, scalars, field_layouts, optimizer)[source]

Bases: object

Point-estimate result for an indexed-latent MAP fit.

latents preserves the historical name-to-vector mapping. estimate() additionally accepts the original free(K) handle, avoiding ambiguity when a caller keeps structural parameter objects. Original field labels and their exact vector positions are retained in index_labels and index_map. group_means remains the single-vector compatibility alias.

Parameters:
estimate(param=None)[source]

Return one fitted point estimate by original handle, vector name, or scalar name.

samples(param=None)[source]

Reject a posterior-sample request on a MAP point estimate.

summary()[source]

Return fitted parameters, label identity, and optimizer termination evidence.

Return type:

dict

indexed_fit(rv, data, *, given=None, how='map', rng=None, draws=2000, burn=1000, thin=1, constraints=None, penalty=None, potentials=None, max_iter=2000, tol=1e-8)[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