gensbi.inference#

Inference wrappers: NLE posterior sampling over trained density flows.

Submodules#

Classes#

`MCLMC`	Microcanonical Langevin Monte Carlo sampler.
`MclmcInfo`	Tuning parameters and diagnostics from an MCLMC run.
`NLEPosterior`	Amortized NLE posterior over a trained likelihood flow.
`NestedSampler`	Nested slice sampling (blackjax `nss`) posterior sampler.
`NestedSamplerInfo`	Diagnostics from a nested sampling run.
`PosteriorTarget`	Log-densities for one observation `x_o`.
`Sampler`	Abstract base class for posterior samplers.
`SmcInfo`	Diagnostics from an adaptive tempered SMC run.
`TemperedSMC`	Adaptive tempered Sequential Monte Carlo for (possibly multimodal) posteriors.

Package Contents#

class gensbi.inference.MCLMC(*, adjusted=True, num_samples=1000, num_tuning_steps=5000, num_chains=1, target_acceptance=0.9, diagonal_preconditioning=True)[source]#

Bases: Sampler

Microcanonical Langevin Monte Carlo sampler.

adjusted=True (the default) applies an MH correction for asymptotically exact sampling. adjusted=False uses the faster unadjusted variant, which is biased by the discretization error.

Parameters:

adjusted (bool, optional) – Whether to apply an MH correction. Default is True.
num_samples (int, optional) – Number of posterior samples to collect per chain. Default is 1000.
num_tuning_steps (int, optional) – Number of warmup steps for the L / step-size tuning loop. Default is 5000.
num_chains (int, optional) – Number of independent MCLMC chains. Default is 1.
target_acceptance (float, optional) – Target Metropolis acceptance rate for the adjusted variant. Default is 0.9.
diagonal_preconditioning (bool, optional) – Whether to use diagonal preconditioning during tuning. Default is True.

_run_adjusted(key, target)[source]#

_run_single(key, target)[source]#

_run_unadjusted(key, target)[source]#

run(key, target)[source]#

Draw posterior samples using MCLMC.

Parameters:

key (jax.random.PRNGKey) – Random key.
target (PosteriorTarget) – Posterior target produced by build_target().

Returns:

samples (Array) – Posterior samples of shape (num_chains * num_samples, dim).
info (MclmcInfo) – Tuning parameters and diagnostic information.

adjusted = True#

diagonal_preconditioning = True#

num_chains = 1#

num_samples = 1000#

num_tuning_steps = 5000#

target_acceptance = 0.9#

class gensbi.inference.MclmcInfo[source]#

Tuning parameters and diagnostics from an MCLMC run.

Parameters:

L (float) – Tuned trajectory length.
step_size (float) – Tuned integrator step size.
acceptance_rate (float) – Mean Metropolis acceptance rate over the sampling run. float('nan') for the unadjusted variant.
num_samples (int) – Number of samples drawn per chain.
num_chains (int) – Number of independent chains.

L: float#

acceptance_rate: float#

num_chains: int#

num_samples: int#

step_size: float#

class gensbi.inference.NLEPosterior(flow, prior, *, structured_obs=False)[source]#

Amortized NLE posterior over a trained likelihood flow.

Parameters:

flow (object) – Exposes log_prob(x, cond) -> (B,) (an NLE-trained MAFlow/TarFlow).
prior (numpyro.distributions.Distribution) – Prior over theta; prior.log_prob(theta) is a scalar and prior.sample(key, ()) returns (dim,).
structured_obs (bool, optional) – If True, x_o keeps its (image/field) shape instead of being flattened. Default is False.

build_target(x_o)[source]#

Build a posterior target for a single observation.

Parameters:: x_o (Array) – Observed data. For non-structured: squeezed then promoted to (dim_x, 1) (channel-carrying). For structured_obs=True: kept as-is (image/field shape).
Returns:: Frozen log-density container for x_o.
Return type:: PosteriorTarget

sample(key, x_o, sampler=None, *, return_info=False)[source]#

Draw posterior samples for a single observation.

Parameters:

key (jax.random.PRNGKey) – Random key.
x_o (Array) – Observed data passed to build_target().
sampler (Sampler or None, optional) – Sampler instance to use. If None, defaults to MCLMC. Default is None.
return_info (bool, optional) – If True, return a (samples, info) tuple instead of just samples. Default is False.

Returns:

samples (Array) – Posterior samples of shape (n, dim, 1). When return_info=False (the default), this is the only return value.
info (object) – Sampler-specific info object (MclmcInfo, SmcInfo, or NestedSamplerInfo). Only present when return_info=True.

flow#

prior#

structured_obs = False#

class gensbi.inference.NestedSampler(*, num_live=500, num_delete=None, num_inner_steps=None, num_samples=1000, dlogz=-3.0, max_iterations=100000, num_rejuvenation_steps=0)[source]#

Bases: Sampler

Nested slice sampling (blackjax nss) posterior sampler.

Runs blackjax’s Nested Slice Sampling from prior-drawn live points, accumulating dead points until the live set’s evidence share is negligible, then resamples the dead-point history into equal-weight posterior draws. Unlike the MCMC samplers this also estimates the log evidence, and handles multimodal posteriors without tempering.

Parameters:

num_live (int, optional) – Number of live points. Default is 500.
num_delete (int or None, optional) – Number of lowest-likelihood points replaced per step (device batching). If None, defaults to max(1, num_live // 10).
num_inner_steps (int or None, optional) – Constrained slice moves per replacement. If None, resolved at run time to max(5, 2 * target.dim) (blackjax’s rule of thumb for reliable mixing). Default is None.
num_samples (int, optional) – Number of equal-weight posterior draws returned. Default is 1000.
dlogz (float, optional) – Termination threshold (blackjax convention): stop once logZ_live - logZ < dlogz. Default is -3.0; use e.g. -10.0 near phase transitions.
max_iterations (int, optional) – Safety cap on the number of outer NS steps. Default is 100_000.
num_rejuvenation_steps (int, optional) – Posterior-invariant hit-and-run slice moves applied to each equal-weight draw after resampling. The with-replacement resampling duplicates draws whenever num_samples is comparable to the run’s ESS; a few slice moves break the duplicated atoms without changing the sampled distribution. Default is 0 (no rejuvenation).

_rejuvenate(key, positions, target)[source]#

Break duplicated equal-weight draws with posterior-invariant moves.

Runs num_rejuvenation_steps hit-and-run slice moves on every resampled draw (one vmapped chain per draw), targeting the unconstrained log-posterior. Directions are shaped by the empirical covariance of the resampled cloud and scaled to Mahalanobis norm 2, the same proposal the NS run’s inner kernel uses – so the moves are local decorrelation only; mode coverage and weights stay as the NS run left them.

_resolve_num_inner_steps(dim)[source]#: num_inner_steps if set, else blackjax’s max(5, 2 * dim).

run(key, target)[source]#

Draw posterior samples using nested slice sampling.

Parameters:

key (jax.random.PRNGKey) – Random key.
target (PosteriorTarget) – Posterior target produced by build_target().

Returns:

samples (Array) – Equal-weight posterior samples of shape (num_samples, dim).
info (NestedSamplerInfo) – Evidence estimate, ESS and the raw finalised run.

dlogz = -3.0#

max_iterations = 100000#

num_delete = None#

num_inner_steps = None#

num_live = 500#

num_rejuvenation_steps = 0#

num_samples = 1000#

class gensbi.inference.NestedSamplerInfo[source]#

Diagnostics from a nested sampling run.

Parameters:

log_evidence (float) – Log marginal likelihood estimate (mean over stochastic prior-volume draws).
log_evidence_err (float) – Standard deviation of the log-evidence estimate over the stochastic prior-volume draws.
ess (float) – Effective sample size of the weighted dead-point set.
num_dead (int) – Total number of points in the finalised run (dead points plus the final live set).
dead (object) – Raw finalised blackjax.ns.base.NSInfo carrying the full point history (positions, log-likelihoods, birth contours). Kept for downstream re-weighting or anesthetic-style analysis.

dead: object#

ess: float#

log_evidence: float#

log_evidence_err: float#

num_dead: int#

class gensbi.inference.PosteriorTarget[source]#

Log-densities for one observation x_o.

Frozen dataclass produced by NLEPosterior.build_target(). All callables accept a flat parameter vector theta of shape (dim,).

Parameters:

log_prior (Callable) – Log-prior density. Signature: log_prior(theta) -> float.
log_likelihood (Callable) – Log-likelihood log q(x_o | theta) from the NLE-trained flow. Signature: log_likelihood(theta) -> float.
log_posterior (Callable) – Unnormalised log-posterior log_likelihood(theta) + log_prior(theta). Signature: log_posterior(theta) -> float.
prior (object) – Prior distribution; must expose sample(key, shape) and log_prob(theta).
dim (int) – Dimensionality of the parameter space.

dim: int#

log_likelihood: Callable#

log_posterior: Callable#

log_prior: Callable#

prior: object#

class gensbi.inference.Sampler[source]#

Bases: abc.ABC

Abstract base class for posterior samplers.

Subclasses consume a PosteriorTarget and return an array of posterior samples together with a sampler-specific info object.

abstractmethod run(key, target)[source]#

Draw posterior samples from a log-density target.

Parameters:

key (jax.random.PRNGKey) – Random key.
target (PosteriorTarget) – Posterior target produced by build_target().

Returns:

samples (Array) – Posterior samples of shape (n, dim).
info (object) – Sampler-specific diagnostics.

Raises:

NotImplementedError – This is an abstract method; subclasses must override it.

class gensbi.inference.SmcInfo[source]#

Diagnostics from an adaptive tempered SMC run.

Parameters:

log_evidence (float) – Log marginal likelihood estimate accumulated over tempering steps.
num_temperature_steps (int) – Number of temperature increments taken from beta=0 to beta=1.
final_tempering_param (float) – Final value of the tempering parameter beta (should be 1.0).

final_tempering_param: float#

log_evidence: float#

num_temperature_steps: int#

class gensbi.inference.TemperedSMC(*, num_particles=1000, target_ess=0.9, num_mcmc_steps=10, inner_kernel='mclmc', inner_step_size=0.1, inner_num_integration_steps=5, inner_inverse_mass_matrix=None)[source]#

Bases: Sampler

Adaptive tempered Sequential Monte Carlo for (possibly multimodal) posteriors.

Walks particles along p(theta) * q(x_o | theta)^beta for beta from 0 to 1, choosing the beta ladder adaptively to maintain target_ess. The inner rejuvenation kernel is adjusted MCLMC by default; fixed-trajectory HMC is available as an alternative. (NUTS is deliberately not offered: its data-dependent trajectory length does not vectorize cleanly across SMC particles, and rejuvenation does not need NUTS’s full-mixing guarantee.)

target_ess and num_mcmc_steps default to values calibrated for blackjax >= 1.6, whose adaptive-tempering ESS solver was corrected (upstream #914 fixed a sign bug in the bisection target). The fix changes how many temperature steps are needed to anneal from beta=0 to 1: with the pre-1.6-era target_ess=0.5 default, the corrected solver can collapse the schedule to a single step, leaving too few temperature increments for rejuvenation to move particles onto the posterior.

Parameters:

num_particles (int, optional) – Number of SMC particles. Default is 1000.
target_ess (float, optional) – Target effective sample size ratio in (0, 1), used to adapt the temperature increments. Default is 0.9.
num_mcmc_steps (int, optional) – Number of inner MCMC rejuvenation steps per temperature. Default is 10.
inner_kernel (str, optional) – Inner MCMC kernel: "mclmc" (adjusted MCLMC, default) or "hmc" (fixed-trajectory HMC).
inner_step_size (float, optional) – Step size for the inner MCMC kernel. Default is 0.1.
inner_num_integration_steps (int, optional) – Number of integration steps for the inner MCLMC or HMC kernel. Default is 5.
inner_inverse_mass_matrix (Array or None, optional) – Inverse mass matrix for the inner kernel. If None, defaults to a vector of ones of length target.dim. Default is None.

_inner(target)[source]#: Return (mcmc_step_fn, mcmc_init_fn, mcmc_parameters) for the inner kernel.

run(key, target)[source]#

Draw posterior samples using adaptive tempered SMC.

Parameters:

key (jax.random.PRNGKey) – Random key.
target (PosteriorTarget) – Posterior target produced by build_target().

Returns:

samples (Array) – Posterior samples of shape (num_particles, dim).
info (SmcInfo) – SMC diagnostics including log evidence and tempering statistics.

inner_inverse_mass_matrix = None#

inner_kernel = 'mclmc'#

inner_num_integration_steps = 5#

inner_step_size = 0.1#

num_mcmc_steps = 10#

num_particles = 1000#

target_ess = 0.9#