gensbi.inference.samplers

blackjax samplers consumed by NLEPosterior. blackjax imported lazily in run().

Classes

MCLMC	Microcanonical Langevin Monte Carlo sampler.
MclmcInfo	Tuning parameters and diagnostics from an MCLMC run.
NestedSampler	Nested slice sampling (blackjax nss) posterior sampler.
NestedSamplerInfo	Diagnostics from a nested sampling run.
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.

Functions

_check_rescale_domain(mu)	Raise if the tuned mu = L / step_size is outside _rescale's domain.
_inference_loop(rng_key, step_fn, initial_state, ...)	Run a blackjax SamplingAlgorithm.step for num_samples via lax.scan.
_rescale(mu)	Map a mean trajectory length to a uniform-integer draw scale.

Module Contents

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

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)

_run_single(key, target)

_run_unadjusted(key, target)

run(key, target)

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.samplers.MclmcInfo

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.samplers.NestedSampler(*, num_live=500, num_delete=None, num_inner_steps=None, num_samples=1000, dlogz=-3.0, max_iterations=100000, num_rejuvenation_steps=0)

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)

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)

num_inner_steps if set, else blackjax’s max(5, 2 * dim).

run(key, target)

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.samplers.NestedSamplerInfo

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.samplers.Sampler

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)

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.samplers.SmcInfo

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.samplers.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)

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)

Return (mcmc_step_fn, mcmc_init_fn, mcmc_parameters) for the inner kernel.

run(key, target)

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

gensbi.inference.samplers._check_rescale_domain(mu)

Raise if the tuned mu = L / step_size is outside _rescale’s domain.

For mu < 1, floor(2 * mu - 1) == 0 and _rescale returns 0, so the integration-step draw ceil(U(0,1) * 0) is 0 — a chain that never moves. A host-side check on the tuned value turns that silent failure into an explicit error. (The in-tuning average is left to blackjax; this is a convenience sampler, not a fully hardened MCMC engine.)

gensbi.inference.samplers._inference_loop(rng_key, step_fn, initial_state, num_samples)

Run a blackjax SamplingAlgorithm.step for num_samples via lax.scan.

gensbi.inference.samplers._rescale(mu)

Map a mean trajectory length to a uniform-integer draw scale.

From blackjax’s adjusted_mclmc_dynamic: drawing the number of integration steps as ceil(U(0,1) * _rescale(L/step_size)) makes the average number of steps exactly mu = L / step_size.

mu must satisfy mu >= 1; see _check_rescale_domain() for the host-side guard applied to the tuned value.