mixle.inference.mcmc.parameter_bridge module

High-level parameter-posterior sampling for mixle.stats families.

Given a prototype distribution that fixes the model family, a dataset, and a prior over parameters, this samples p(theta | data) proportional to exp(sum_i log p(x_i | theta)) * prior(theta) by running Metropolis-Hastings or Hamiltonian Monte Carlo in an unconstrained reparameterization (log for positive scales, stick-breaking for probability simplices) and mapping the retained samples back to parameter space (or to rebuilt distribution objects).

Seven families (Gaussian, Gamma, Beta, Exponential, Poisson, Bernoulli, Categorical) get a hand-tuned bridge below. Every other family that declares a DistributionDeclaration – the same per-parameter constraint/differentiable metadata mixle.inference.gradient_fit already reads for autograd fitting – is bridged generically by _generic_declared_bridge(), so adding a new supported family here is a matter of the family declaring itself, not editing this module. See build_parameter_bridge()’s docstring for exactly which constraints have a generic reparameterization and why some declared families (a covariance matrix, a coupled bound, a natural/scoring parameterization that differs from the constructor) are still excluded.

class ParameterBridge(dim, to_unconstrained, from_unconstrained, log_abs_det_jacobian, build, param_names, initial_theta=None)

Bases: object

Map between a distribution’s parameters and an unconstrained vector.

Parameters:

• dim (int)

• to_unconstrained (Callable[[Any], ndarray])

• from_unconstrained (Callable[[ndarray], Any])

• log_abs_det_jacobian (Callable[[ndarray], float])

• build (Callable[[Any], Any])

• param_names (tuple[str, ...])

• initial_theta (Any)

dim

Length of the unconstrained vector phi.

Type:

int

to_unconstrained

theta -> phi for the prototype’s parameters.

Type:

collections.abc.Callable[[Any], numpy.ndarray]

from_unconstrained

phi -> theta (parameter-space value).

Type:

collections.abc.Callable[[numpy.ndarray], Any]

log_abs_det_jacobian

phi -> log|det dtheta/dphi| so a flat prior in theta-space maps to the right density in phi-space.

Type:

collections.abc.Callable[[numpy.ndarray], float]

build

theta -> distribution rebuilds a model from parameters.

Type:

collections.abc.Callable[[Any], Any]

param_names

Human-readable names for each theta block (diagnostics).

Type:

tuple[str, …]

initial_theta

The prototype’s parameters in theta-space (chain start).

Type:

Any

build_parameter_bridge(prototype)

Build a ParameterBridge for a prototype distribution.

Seven families get a hand-tuned bridge below:

• Gaussian (mu, sigma2) -> (mu, log sigma2)

• Gamma (k, theta) -> (log k, log theta)

• Exponential beta/lam (positive scalar) -> log

• Poisson lam -> log lam

• Bernoulli p -> logit p

• Beta (a, b) -> (log a, log b)

• Categorical probability map -> stick-breaking over the simplex

Every other family falls through to _generic_declared_bridge(), which reads the same DistributionDeclaration that mixle.inference.gradient_fit already uses for autograd fitting. A declared parameter is sampled – included in theta (a {name: value} dict here, unlike the tuple/scalar shapes above) – when it is differentiable and its constraint is one of real (identity), positive/positive_vector (log), unit_interval (logit), real_vector (identity), simplex_vector/simplex_map (stick-breaking), or unit_interval_map (independent logits); non-differentiable declared parameters (e.g. a Binomial’s n) are carried as fixed constructor keywords taken from prototype, the same nuisance-parameter treatment mixle.stats.bayes.conjugate uses. This covers 33 families as of this writing (the 7 above plus 26 more – see mcmc_test.py’s GenericParameterBridgeTestCase for the current list).

Still unsupported, deliberately: families with no declaration at all (compositional/ structured models – mixtures, HMMs, copulas, … – for which a single flat parameter vector isn’t the right shape); families whose declaration describes an exponential-family natural/scoring parameterization rather than the constructor’s (e.g. von Mises’ eta1/ eta2/log_const vs. its mu/kappa constructor – the declared names aren’t constructor keywords, so there’s nothing safe to dispatch on); and declared-but-exotic constraints with no generic reparameterization yet (a covariance/precision matrix’s positive_matrix, or a coupled greater_than:/less_than: bound such as Uniform’s high). These need a family-specific bridge (or a matrix/coupled-bound generalization) rather than a generic one and are out of scope here.

Raises:

NotImplementedError – if the family or a parameter shape is unsupported.

Parameters:

prototype (Any)

Return type:

ParameterBridge

sample_parameter_posterior(prototype_dist, data, prior=None, sampler='mh', steps=2000, burn_in=500, thin=1, seed=None, proposal=None, initial=None, step_size=0.05, num_steps=20, grad_log_target=None, return_distributions=False)

Sample the parameter posterior p(theta | data) of a distribution.

The model family is fixed by prototype_dist; its parameters define the sampled space. The unnormalized log target is the data log-likelihood (rebuilding the distribution per proposed theta and summing seq_log_density) plus a prior log-density and the reparameterization Jacobian. Sampling runs in an unconstrained space (see build_parameter_bridge()) so proposals stay in-domain, and retained samples are mapped back to parameter space.

Parameters:

• prototype_dist (Any) – A distribution instance fixing the model family.

• data (Any) – Observations accepted by the family’s encoder.

• prior (Any) – None (flat), a callable theta -> log p(theta), or a distribution with log_density over the parameter representation.

• sampler (str) – 'mh' (Metropolis-Hastings), 'hmc' (Hamiltonian Monte Carlo), or 'nuts' (No-U-Turn Sampler, self-tuning). HMC/NUTS use grad_log_target if given, else a finite-difference gradient.

• steps (int) – Number of retained posterior samples.

• burn_in (int) – Number of initial transitions to discard.

• thin (int) – Keep one sample every thin transitions.

• seed (int | None) – Seed for the RandomState.

• proposal (Proposal | None) – Optional MH proposal in the unconstrained space; defaults to a random walk (MH only).

• initial (Any) – Optional starting theta; defaults to the prototype’s parameters.

• step_size (float) – HMC leapfrog controls (NUTS self-tunes step size).

• num_steps (int) – HMC leapfrog controls (NUTS self-tunes step size).

• grad_log_target (Callable[[Any], Any] | None) – Optional exact gradient of log_target in the unconstrained space (e.g. from mixle.inference.mcmc.gradients.torch_gradient()). Replaces the finite-difference gradient for HMC/NUTS – one backward pass per step instead of O(dim) target evaluations.

• return_distributions (bool) – If True, samples are rebuilt distribution objects instead of parameter-space values.

Returns:

MCMCResult whose samples are parameter-space values (or rebuilt distributions) and whose diagnostics come from the underlying driver.

Return type:

MCMCResult