autofit.SMC

class SMC

Bases: AbstractMCMC

A BlackJAX adaptive tempered Sequential Monte Carlo (SMC) non-linear search with a gradient inner kernel.

SMC anneals a cloud of num_particles particles along a tempering path from a normalised starting distribution (lambda = 0) to the posterior (lambda = 1), moving the particles at each temperature with an inner MCMC kernel and reweighting/resampling between temperatures. Two properties make it the lead sampler of the JAX-native wave:

1. It yields the log evidence for free – the sum of every step’s tempering log_likelihood_increment is log Z, on the same scale a nested sampler reports.

2. Its inner kernel is gradient-based (MALA by default, HMC optionally), so it exploits the JAX-native differentiable likelihood rather than random-walking.

The autofit Analysis must therefore be constructed with use_jax=True (and JAX run in float64, JAX_ENABLE_X64=True); a clear error is raised at fit time otherwise.

For BlackJAX, see https://github.com/blackjax-devs/blackjax.

Whitening – and why it is the whole game. The inner kernel takes a single scalar step size, so the target must be close to isotropic and unit-scale for any step to work. The search therefore samples in a whitened space z, related to the physical parameters by the affine map x = shift + L @ z, and this is what inverse_mass_matrix sets:

• Prior-whitening does not whiten the posterior. Measured on a 15-parameter lens model, the posterior is 269x anisotropic in prior-whitened coordinates (one parameter’s prior scale is 8.0 against a posterior std of 2e-4). A scalar step tuned to the mean of that spread is ~88x too large for the tightest parameter and acceptance pins at exactly 0.000.

• Diagonal whitening is not enough. The posterior is correlated (condition number 568, correlation coefficient 0.95 between two parameters); diagonal whitening leaves a thin tilted ridge a spherical proposal walks off. Whitening uses the Cholesky factor of the full covariance.

• Cold, there is nothing better than the prior scale, which is what inverse_mass_matrix=None falls back to.

Warm starting – and how the evidence survives it. These samplers are meant to be handed a starting point near the maximum-likelihood solution by a JAX optimizer; cold-benchmarking them is not a fair test (on the lens model above, every cold arm sat ~190,000 log-units below the optimizer’s solution). To warm start:

search = af.SMC(
    initializer=result.start_point_from(),
    inverse_mass_matrix=result,
)

Simply dropping the particles at the previous best fit would destroy the log evidence, which is only valid when tempering starts from a normalised distribution. So the warm path does not shortcut lambda: it still starts the tempering at lambda = 0, but from a normalised Gaussian reference g = N(shift, L L^T) centred on the warm-start point rather than from the prior, bridging

log target_lambda = log g + lambda * (log prior + log L - log g)

onto BlackJAX’s logprior_fn + lambda * loglikelihood_fn. Because g integrates to 1, the accumulated increments still estimate log Z = log integral(prior * L) – the true evidence, directly comparable to Nautilus. The prior must carry its normalisation here (it cancels in a Metropolis ratio but not in the evidence), and the whitening Jacobian log|det L| is added back.

Judging a run. Never judge an SMC fit by “it reached lambda = 1”. When acceptance collapses to ~0, adaptive tempering can force a single huge increment straight to lambda = 1 and report a meaningless log_evidence while still looking converged. Judge it by the per-step acceptance trace and the max-log-likelihood progression, both recorded in search_internal and surfaced on SamplesSMC as acceptance_rate_list / lambda_list / ess_list.

Parameters:

• name (Optional[str]) – The name of the search, controlling the last folder results are output to.

• path_prefix (Optional[str]) – The path of folders prefixing the name folder where results are output.

• unique_tag (Optional[str]) – A unique tag for this model-fit, used as a folder between the path prefix and the search name and as the SQLite identifier.

• num_particles (int) – The number of SMC particles carried along the tempering path. This is also the number of posterior samples the fit returns.

• kernel (str) – The gradient inner kernel used to move the particles at each temperature: "mala" (default) or "hmc". HMC holds acceptance far higher as lambda -> 1 but costs ~4x the likelihood/gradient evaluations; measured on the lens model, plain MALA matched its max log likelihood and log evidence at a quarter of the evaluations.

• num_mcmc_steps (int) – The number of inner-kernel updates applied to every particle at each temperature.

• num_integration_steps (int) – Leapfrog steps per HMC trajectory. Ignored when kernel="mala".

• target_ess (float) – The effective sample size, as a fraction of num_particles, that adaptive tempering targets when choosing each lambda increment. Higher means a finer schedule (more, smaller steps).

• step_size (Optional[float]) – The inner kernel’s step size, in whitened units. None (default) auto-scales it from the target width (see below). Units trap: MALA proposes x + eps*grad + sqrt(2*eps)*xi, so its step_size is a squared length and the proposal length is sqrt(2*eps); HMC’s step size is a length. The auto-scaling applies the optimal MALA rule ell = 2.38 * d^(-1/6) * sigma, eps = ell^2 / 2, and sigma * d^(-1/4) for HMC, so a hand-set step_size must respect the same convention. There is deliberately no per-temperature step adaptation: BlackJAX’s inner_kernel_tuning was measured to collapse acceptance (0.04-0.12, worse than a fixed step’s 0.84 -> 0.15 decay) on this problem.

• whiten_inflate (float) – The warm-start Gaussian reference is deliberately widened by this factor so it covers the posterior (the tempering path misses mass if the reference is narrower than the target). It also sets the auto step size: in whitened units the reference has width 1 but the posterior has width 1 / whiten_inflate, so the step targets the posterior width. Targeting the reference width instead overshoots by whiten_inflate^2 and pins acceptance at zero (measured: eps=1.148 -> acc 0.00, eps=0.1 -> acc 0.94).

• max_smc_steps (int) – A hard cap on the number of SMC temperatures, so a pathological run terminates instead of annealing forever. A run that hits it is reported as not converged.

• batch_size (int) – Forwarded to BlackJAX: when > 0, particles are processed in sequential batches of this size via jax.lax.map instead of one full jax.vmap, trading speed for peak GPU memory. 0 (default) keeps the full vmap.

• seed (int) – Integer seed passed to jax.random.PRNGKey.

• initializer (Optional[Initializer]) – Generates the initial particles. Defaults to InitializerPrior – cold SMC’s evidence requires the particles to be drawn from the prior, so the InitializerBall default the other MCMC searches use would silently invalidate log_evidence. When warm-started (inverse_mass_matrix given), the initializer supplies only the centre of the Gaussian reference and the particles are drawn from that reference instead – pass result.start_point_from() (equivalently InitializerParamStartPoints.from_result(result)).

• inverse_mass_matrix (Union[None, str, ndarray, object]) –

  The covariance the sampling space is whitened by (named for API parity with BlackJAXNUTS, where the same specification seeds window adaptation’s metric):

  • None (default): cold. Whiten by the per-parameter prior width.

  • a 1-D numpy array of shape (n_dim,): a diagonal covariance.

  • a 2-D numpy array of shape (n_dim, n_dim): a full covariance, Cholesky-factorised.

  • a Result or Samples object: its samples.covariance_matrix. Raises ValueError if that covariance looks MLE-only (too few samples, or non-finite / identity) – pass an explicit array in that case, e.g. from a Laplace approximation.

  The "diagonal" / "dense" strings BlackJAXNUTS accepts are rejected here: they name an adaptation strategy, and SMC does not adapt its metric, so they would silently give a cold run.

• auto_correlation_settings (AutoCorrelationsSettings) – Kept for API parity with the other MCMC searches. SMC particles are not a chain, so no auto-correlation diagnostics are computed and check_for_convergence defaults to False.

• iterations_per_full_update (Optional[int]) – Inherited from the autonerves config when None. SMC calls perform_update once per temperature (SMC steps are coarse and few), so this does not chunk the run.

• number_of_cores (int) – Used only when generating the initial particles. The tempering loop itself runs vmapped on a single device.

• silence (bool) – If True, the default print output of the non-linear search is silenced.

• session (Optional[Session]) – An SQLalchemy session instance.

Methods

apply_test_mode	Override in subclasses to reduce sampler iterations for test mode.
check_model
copy_with_paths
effective_step_size	The inner kernel's step size in whitened units: step_size if given, otherwise auto-scaled from the target width.
exact_fit
fit	Fit a model, M with some function f that takes instances of the class represented by model M and gives a score for their fitness.
make_pool	Make the pool instance used to parallelize a NonLinearSearch alongside a set of unique ids for every process in the pool.
make_sneakier_pool
make_sneaky_pool	Create a pool for multiprocessing that uses slight-of-hand to avoid copying the fitness function between processes multiple times.
optimise	Perform optimisation for expectation propagation.
output_search_internal	Pickle the search-internal dict.
perform_update	Perform an update of the non-linear search's model-fitting results.
perform_visualization	Perform visualization of the non-linear search's model-fitting results.
plot_results
plot_start_point	Visualize the starting point of the non-linear search, using an instance of the model at the starting point of the maximum likelihood estimator.
post_fit_output	Cleans up the output folderds after a completed non-linear search.
pre_fit_output	Outputs attributes of fit before the non-linear search begins.
result_via_completed_fit	Returns the result of the non-linear search of a completed model-fit.
samples_from	Loads the samples of a non-linear search from its output files.
samples_info_from
samples_via_internal_from	Convert the particle cloud pickled under search_internal/ into a SamplesSMC.
start_resume_fit

Attributes

backend	Load the pickled search-internal dict written by _fit.
backend_filename
is_warm_start	Whether the run whitens (and tempers) from a warm-start covariance rather than from the prior.
logger	Log 'msg % args' with severity 'DEBUG'.
name
paths
quick_update_message	One line, logged at the start of every search, telling the user the real cadence of the on-the-fly maximum-likelihood updates.
should_plot_start_point
timer	Returns the timer of the search, which is used to output informaiton such as how long the search took and how much parallelization sped up the search time.

apply_test_mode()

Override in subclasses to reduce sampler iterations for test mode.

Called during __init__ when test mode is active (level 1). Subclasses should directly mutate instance attributes to minimize the number of iterations the sampler performs.

property is_warm_start: bool

Whether the run whitens (and tempers) from a warm-start covariance rather than from the prior. Set by passing a Result / Samples / covariance array as inverse_mass_matrix.

effective_step_size(n_dim)

The inner kernel’s step size in whitened units: step_size if given, otherwise auto-scaled from the target width.

Warm-started, the target width is the posterior width, which in whitened units is 1 / whiten_inflate (the reference is inflated to cover the posterior, so scaling to the reference’s own width of 1 overshoots by whiten_inflate^2). Cold, the posterior width is unknown and _COLD_TARGET_WIDTH is used.

Parameters:

n_dim (int) – The number of free model parameters, which sets the dimension scaling of the optimal step.

property backend_filename

property backend: dict

Load the pickled search-internal dict written by _fit.

output_search_internal(search_internal)

Pickle the search-internal dict.

BlackJAX has no native on-disk format (cf. emcee’s HDFBackend), so we round-trip the particle cloud + tempering diagnostics via pickle, bypassing self.paths.save_search_internal for the same reason BlackJAXNUTS does: the autofit dill path chokes on a few numpy/jax-backed members, and a direct pickle of already-numpy data is robust.

NullPaths (no name/path_prefix) sets search_internal_path to None to suppress disk output – skip silently in that case.

samples_info_from(search_internal=None)

samples_via_internal_from(model, search_internal=None)

Convert the particle cloud pickled under search_internal/ into a SamplesSMC.

SMC particles are weighted samples: each carries the normalised importance weight blackjax assigns it at the current temperature, so the weights (not a uniform 1.0) are what the PDF, medians and errors are computed from.