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.

  search.json records only the kind ("none", "array" or "result"), so a search loaded from it receives that string: "none" behaves as None, while "array" / "result" raise at fit time because the covariance was not saved.

• 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.
fitness_overrides	Fitness keyword arguments a run(ctx) search adds to make_fitness for this analysis (e.g. batched=True for a batched backend).
info_from	The samples_info entries of a run(ctx) search's internal state, the info its raw_samples_from puts on the RawSamples (autofit.non_linear.samples.adapter).
load_search_internal	The search's internal state as stored in its output folder, or None when its paths store nothing (NullPaths, DatabasePaths).
make_fitness	The Fitness this search evaluates its likelihood through, the one construction site every search shares.
make_pool	Make the pool instance used to parallelize a NonLinearSearch alongside a set of unique ids for every process in the 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	Archive search_internal through the search's checkpointer and, once a fit has completed, discard its resume state (there is nothing left to resume).
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.
raw_samples_from	The particle cloud pickled under search_internal/ as raw samples.
result_via_completed_fit	Returns the result of the non-linear search of a completed model-fit.
run	The backend loop of a search on the run(ctx) hook: run the backend through ctx and return its internal state (search_internal).
samples_from	Loads the samples of a non-linear search from its output files.
samples_info_from
samples_via_internal_from	The samples of the search, converted from its internal state.
start_points	The search's n initial points, drawn by its initializer, with the start-point plot made once, here, from the first of them.
start_resume_fit

Attributes

backend	The search-internal dict archived by _fit (search_internal.pickle).
backend_filename
batched
checkpointer	where the backend's final state is stored when a fit completes, for Result.search_internal and samples_via_internal_from.
citation_keys
gradient
honours_gradient_mode
install_extra
invalid_value
is_warm_start	Whether the run whitens (and tempers) from a warm-start covariance rather than from the prior.
jax_use
logger	Log 'msg % args' with severity 'DEBUG'.
name
objective_target
paths
posterior_kind
produces_evidence
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.
requires_fp64
resumable
resume_state	The file an interrupted run resumes from, for searches that genuinely resume (None otherwise).
retains_search_internal	always when the archive declares retain_after_completion (its results cannot be rebuilt without it), otherwise when the output.search_internal config is on.
should_plot_start_point
status
test_mode_budget
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.
upstream_url
warm_start

samples_cls

alias of SamplesSMC

jax_use = 'required'

gradient = 'uses'

batched = True

honours_gradient_mode = False

posterior_kind = 'weighted'

produces_evidence = True

resumable = False

warm_start = 'provider'

install_extra = 'optional'

upstream_url = 'https://github.com/blackjax-devs/blackjax'

citation_keys = ('blackjax',)

status = 'experimental'

test_mode_budget = {'max_smc_steps': 5, 'num_mcmc_steps': 2, 'num_particles': 16}

objective_target = (ObjectiveQuantity.LOG_LIKELIHOOD, CoordinateSpace.PHYSICAL)

invalid_value = -1e+99

checkpointer: Checkpointer = PickleCheckpointer('search_internal.pickle', retain_after_completion=True)

where the backend’s final state is stored when a fit completes, for Result.search_internal and samples_via_internal_from.

Type:

The archive strategy

apply_test_mode()

Override in subclasses to reduce sampler iterations for test mode.

The base class never calls this method: each search that overrides it calls it at the end of its own __init__ when is_test_mode() is true (test mode level 1). Overrides 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

The search-internal dict archived by _fit (search_internal.pickle).

samples_info_from(search_internal=None)

raw_samples_from(model, search_internal)

The particle cloud pickled under search_internal/ as raw samples.

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. They are normalised here, by the search, as they always were; the samples adapter never renormalises.