pespace.inference.sampler.blackjax_smc

Adaptive-tempered SMC samplers with automatically tuned inner kernels.

Two samplers are provided, both running an adaptive-tempered SMC loop built from blackjax primitives (ESS-based tempering, systematic resampling) with an inner mutation kernel whose parameters are re-tuned from the particle cloud at every tempering level:

• BlackjaxSMCwithNUTS mutates the particles with NUTS on the unconstrained space prepared by NumpyroPriorSamplerBase (the prior bijections and the inverse mapping back to physical parameters are handled by the base class). Following SHARPy, a fixed shared step_size is used, and the per-particle dense inverse_mass_matrix is re-computed at every level as the inverse of the SoftAbs-regularized negative Hessian of the tempered log-posterior at each particle’s position (position-dependent Riemannian metric, frozen along each NUTS trajectory).

• BlackjaxSMCwithIRMH mutates the particles with an independent Metropolis-Hastings kernel directly in the physical parameter space; the Gaussian proposal mean / covariance are re-fitted to the particle cloud and the output is already physical.

For both samplers the number of mutation rounds per tempering level is a fixed number of rounds given by num_mcmc_steps.

With show_progress=True each level’s mutation is wrapped in blackjax.progress_bar() (requires the optional jaxtap package) and per-level diagnostics (temperature, tempering increment and kernel statistics) are printed. With save_diagnosis=True each level’s mutation and tuning info and the diagnostic figures are stored under outdir/_diagnosis_info/level_{i}_mutate/.

Functions

_hessian_matrix(fn, position)	Dense (d, d) Hessian of a scalar pytree function.
_softabs_metric(hessian, alpha)	SoftAbs-regularized metric of the Hessian (Betancourt 2013).
_square_mahal_distance(diff, inverse_mass_matrix)	Squared distance per row under the mass metric, dq^T M dq.
_stack_particles(particles)	Ravel each particle into a row of an (n_particles, dim) matrix.
_std_ratio(previous, particles)	Per-dimension ratio of the post/pre-mutation population std.

Classes

BlackjaxSMCwithIRMH(likelihood, prior_dict)	Adaptive-tempered SMC with an independent MH mutation kernel.
BlackjaxSMCwithNUTS(likelihood, prior_dict)	Adaptive-tempered SMC with a NUTS mutation kernel.
_TunedAdaptiveTemperedSMC(likelihood, prior_dict)	Tempering, resampling, and mutation shared by the SMC samplers.

class pespace.inference.sampler.blackjax_smc.BlackjaxSMCwithNUTS(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0, num_particles=4096, target_ess=0.9, num_mcmc_steps=1, batch_size=0, show_progress=False, save_diagnosis=False, step_size=0.2, softabs_alpha=0.3, hessian_batch_size=32, max_num_doublings=7, divergence_threshold=100)

Bases: _TunedAdaptiveTemperedSMC

Adaptive-tempered SMC with a NUTS mutation kernel.

NUTS requires an unbounded domain, so sampling happens on the unconstrained space prepared by NumpyroPriorSamplerBase (the prior bijections and the inverse transform back to physical parameters are handled by the base class).

At every tempering level the per-particle dense inverse_mass_matrix is re-computed as the inverse of the SoftAbs-regularized negative Hessian of the tempered log-posterior evaluated at each particle’s position (a position-dependent Riemannian metric, frozen along each NUTS trajectory, as in SHARPy). A single fixed step_size is used for all particles and levels; no step-size tuning is performed.

Parameters:

• likelihood (object) – Likelihood instance providing log_likelihood and the likelihood_param_type namedtuple type of the likelihood parameters.

• prior_dict (dict) – Prior declaration, see pespace.inference.sampler.base.NumpyroPriorSamplerBase.

• label (str, optional) – Label of the run, converted to str; also the JSON filename (outdir/{label}_results.json). Default is None (-> "None").

• outdir (str, optional) – Output directory for the saved result. It must not exist yet: the base class creates it and refuses to reuse an existing path, so that the results of an earlier run are never overwritten. Default is None, which resolves to an outdir_{label} directory under the current working directory.

• rkey_seed (int, optional) – Seed for the JAX random number generator key. Default is 0.

• num_particles (int, optional) – Number of SMC particles. Default is 4096.

• target_ess (float, optional) – Target ESS ratio of the adaptive tempering schedule. Default is 0.9.

• num_mcmc_steps (int, optional) – Fixed number of mutation rounds per tempering level. Default is 1.

• batch_size (int, optional) – Number of particles processed per sequential batch when positive. 0 (default) keeps the plain jax.vmap behaviour. A positive value bounds the peak memory of the likelihood intermediates and the mutation updates, at the cost of cross-batch parallelism. The per-particle Hessian tuning is chunked separately via hessian_batch_size.

• show_progress (bool, optional) – When True, wrap each level’s mutation in blackjax.progress_bar() (requires the optional jaxtap package) and print per-level diagnostics (temperature, tempering increment and the step-size / leapfrog / acceptance statistics below). Default is False.

• save_diagnosis (bool, optional) – When True, save each level’s full mutation info and the diagnostic figures under outdir/_diagnosis_info/level_{i}_mutate/ (mutate_info.h5, hist_num_integration_steps.png, particles_Mahal.png and particles_corner.png; no per-particle tuning info is written). Default is False.

• step_size (float, optional) – Fixed NUTS step size shared by all particles and levels (no step-size tuning is performed). Default is 0.2.

• softabs_alpha (float, optional) – SoftAbs regularization parameter of the position-dependent metric (Betancourt 2013): the metric eigenvalues are lambda / tanh(alpha * lambda), so near-flat directions get a finite mass 1 / alpha and the metric stays positive definite everywhere. Default is 0.3.

• hessian_batch_size (int, optional) – Number of particles per sequential chunk when the per-particle Hessians of the mass-matrix tuning are evaluated. The Hessian evaluation is far more memory-hungry than a plain likelihood evaluation (roughly dim forward tangents are materialised per particle), so it is chunked separately from batch_size: 0 keeps the plain jax.vmap behaviour and a positive value caps the chunk (and hence the peak memory) of the tuning evaluation. Lower it when the GPU runs out of memory. Default is 32.

• max_num_doublings (int, optional) – Maximum number of trajectory doublings of the NUTS kernel (at most 2**max_num_doublings leapfrog steps per trajectory). Default is 7.

• divergence_threshold (int, optional) – Absolute energy error above which a transition is flagged as divergent. Default is 100.

_initial_particles(key, num_particles)

_tune_mcmc_parameters(key, particles, tempering_param)

Build the per-particle Hessian-based mass matrices.

The per-particle dense inverse_mass_matrix is SoftAbs(-hessian)^{-1} of the tempered log-posterior at each (resampled) particle’s position (blackjax convention: the quadratic form entering the kinetic energy). SoftAbs keeps the metric positive definite everywhere and gives near-flat directions a finite mass 1 / alpha; the matrix is frozen along each NUTS trajectory, so the standard velocity-Verlet integrator and the detailed-balance argument remain valid. The evaluation is chunked with hessian_batch_size to bound the peak memory (one Hessian costs roughly dim gradient evaluations per particle). The fixed step size is rebuilt here and no probe move is performed, so the tuning returns the empty tuple as diagnostics.

_mutate_eval_counts(mutate_info)

NUTS: one value-and-gradient evaluation per particle for the init, plus one per leapfrog step of every trajectory.

_tune_eval_counts(tune_info)

NUTS: one value evaluation per particle for the level’s mass matrices, plus roughly dim gradient evaluations per particle for the Hessians (jax.hessian is forward-over-reverse).

_print_diagnose_info(mcmc_params, mutate_info, previous, particles, cum_ll, cum_grad, elapsed)

NUTS diagnostics for one level: the fixed step size, leapfrog length, acceptance rate, divergent-transition fraction, Mahalanobis distances, the fraction of stationary particles (unchanged positions), per-dimension std ratios and cumulative evaluation counts.

_diagnosis_plots(previous, particles, mcmc_params, mutate_info, tune_info, level_dir)

NUTS diagnostic figures: the integration-step histogram, the Mahalanobis distance histogram and the pre/post-mutation corner plot.

_abc_impl = <_abc._abc_data object>

class pespace.inference.sampler.blackjax_smc.BlackjaxSMCwithIRMH(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0, num_particles=4096, target_ess=0.9, num_mcmc_steps=100, batch_size=0, show_progress=False, save_diagnosis=False, covariance_regularization=1e-06)

Bases: _TunedAdaptiveTemperedSMC

Adaptive-tempered SMC with an independent MH mutation kernel.

Sampling is performed directly in the physical parameter space, so the output is already physical and needs no post-processing. The Gaussian proposal mean / covariance are re-fitted to the particle cloud at every tempering level.

Parameters:

• likelihood (object) – Likelihood instance providing log_likelihood and the likelihood_param_type namedtuple type of the likelihood parameters.

• prior_dict (dict) – Prior declaration, see pespace.inference.sampler.base.NumpyroPriorSamplerBase.

• label (str, optional) – Label of the run, converted to str; also the JSON filename (outdir/{label}_results.json). Default is None (-> "None").

• outdir (str, optional) – Output directory for the saved result. It must not exist yet: the base class creates it and refuses to reuse an existing path, so that the results of an earlier run are never overwritten. Default is None, which resolves to an outdir_{label} directory under the current working directory.

• rkey_seed (int, optional) – Seed for the JAX random number generator key. Default is 0.

• num_particles (int, optional) – Number of SMC particles. Default is 4096.

• target_ess (float, optional) – Target ESS ratio of the adaptive tempering schedule. Default is 0.9.

• num_mcmc_steps (int, optional) – Fixed number of mutation rounds per tempering level. Default is 100.

• batch_size (int, optional) – Number of particles processed per sequential batch when positive. 0 (default) keeps the plain jax.vmap behaviour. A positive value bounds the peak memory of the likelihood intermediates and the mutation updates, at the cost of cross-batch parallelism.

• show_progress (bool, optional) – When True, wrap each level’s mutation in blackjax.progress_bar() (requires the optional jaxtap package) and print per-level diagnostics (temperature, tempering increment and the acceptance / proposal-scale statistics below). Default is False.

• save_diagnosis (bool, optional) – When True, save each level’s full mutation info under outdir/_diagnosis_info/level_{i}_mutate/ (mutate_info.h5; IRMH has no per-particle tuning info, so no tune_info.h5 is written). No diagnostic figures are drawn yet for IRMH. Default is False.

• covariance_regularization (float, optional) – Diagonal regularization added to the fitted proposal covariance. Default is 1e-6.

_irmh_step_fn(rng_key, state, logdensity_fn, mean, covariance)

IRMH kernel with a shared Gaussian proposal N(mean, covariance).

_initial_particles(key, num_particles)

_tune_mcmc_parameters(key, particles, tempering_param)

_tune_eval_counts(tune_info)

IRMH: the proposal is fitted without a measurement move, so the tuning probe contributes no evaluations.

_mutate_eval_counts(mutate_info)

IRMH: one log-density evaluation per particle for the init plus one per mutation round; no gradients.

_print_diagnose_info(mcmc_params, mutate_info, previous, particles, cum_ll, cum_grad, elapsed)

IRMH diagnostics for one level: acceptance rate and the proposal scale range.

_diagnosis_plots(previous, particles, mcmc_params, mutate_info, tune_info, level_dir)

Placeholder: no IRMH diagnostic figures implemented yet.

_fitted_proposal(particles)

Gaussian proposal parameters fitted to the current particle cloud.

_abc_impl = <_abc._abc_data object>