autofit.BlackJAXNUTS

class BlackJAXNUTS

Bases: AbstractMCMC

A BlackJAX No-U-Turn Sampler (NUTS) non-linear search.

BlackJAX is a sampling library implemented on top of JAX. NUTS is its gradient-based MCMC kernel — an extension of HMC that adapts trajectory length on the fly. The autofit Analysis therefore must be constructed with use_jax=True so the log-likelihood is JAX-traceable end-to-end (and jax.grad of it can be taken). A clear error is raised at fit time otherwise.

For a full description of BlackJAX, see:

https://github.com/blackjax-devs/blackjax

The fit runs in two phases, both vmapped over num_chains independent chains:

1. blackjax.window_adaptation warmup, which tunes the leapfrog step size (dual averaging) and an inverse mass matrix (diagonal or dense, see inverse_mass_matrix below) per chain.

2. NUTS sampling with each chain’s tuned kernel, run inside a jax.lax.scan so the inner step is fully JIT-compiled. The scan is broken into iterations_per_full_update-sized chunks so partial state is persisted and perform_update runs periodically (samples.csv flush, plotting), mirroring the Emcee chunking pattern.

Warm starting. Both the starting point(s) and the mass matrix can be seeded from a previous Result, without ever modifying that result’s model’s priors (contrast with Result.model_centred*, which rewrite priors to new GaussianPrior`s). For example, to warm-start a 4-chain run from an earlier (e.g. single-chain, or MLE) fit’s ``result`:

search = af.BlackJAXNUTS(
    num_chains=4,
    initializer=result.start_point_from(n_points=4, jitter=0.05),
    inverse_mass_matrix=result,
)

result.start_point_from(...) (equivalently InitializerParamStartPoints.from_result(result, n_points=4, jitter=0.05)) maps the previous result’s best-fit point onto the (possibly different) target model by path, then jitters n_points physical-space starting vectors around it — one per chain. Passing inverse_mass_matrix=result (or result.samples) seeds warmup’s dense inverse mass matrix from that result’s samples.covariance_matrix; this raises a clear error if the result looks MLE-only (too few samples for a reliable covariance), in which case an explicit array (e.g. from a Laplace approximation) should be passed instead.

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_warmup (int) – Number of warmup steps used by blackjax.window_adaptation to tune the leapfrog step size and inverse mass matrix.

• num_samples (int) – Number of post-warmup samples drawn from the tuned NUTS kernel, per chain.

• num_chains (int) – Number of independent chains sampled in parallel via vmap.

• target_accept (float) – Target Metropolis acceptance rate for window adaptation (default 0.8 — the standard Stan setting).

• max_num_doublings (int) – NUTS doubling cap. 10 means at most 1024 leapfrog steps per sample, which is the standard ceiling.

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

• initializer (Optional[Initializer]) – Generates the starting point(s), one per chain. Defaults to InitializerBall(0.49, 0.51) from AbstractMCMC — small ball around the prior median in unit-cube coordinates. Pass result.start_point_from(n_points=num_chains, jitter=...) (or an InitializerParamStartPoints.from_result(...)) to warm-start from a previous result instead.

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

  Controls the metric adapted by blackjax.window_adaptation:

  • None (default): adapt a fresh diagonal inverse mass matrix, seeded at the identity (standard behaviour).

  • "diagonal" / "dense": adapt a diagonal / dense inverse mass matrix, still seeded at the identity.

  • a numpy array: seed the adaptation with this matrix – a 1-D array of shape (n_dim,) seeds a diagonal metric, a 2-D array of shape (n_dim, n_dim) seeds a dense metric.

  • a Result or Samples object: seed a dense metric from 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.

  search.json records only the kind ("none", "diagonal", "dense", "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 seed was not saved.

• mass_matrix_shrinkage (float) – Forwarded to blackjax.window_adaptation’s imm_shrinkage_to_previous: shrinkage of each warmup window’s adapted inverse mass matrix toward the previous window’s (in addition to Stan’s existing shrinkage toward the identity). 0.0 (default) reproduces Stan’s standard per-window-reset behaviour; a positive value lets a high-confidence inverse_mass_matrix seed persist further into warmup.

• share_adaptation (bool) – If True, after per-chain warmup the num_chains tuned inverse mass matrices are averaged and the tuned step sizes combined via their median, and every chain’s sampling phase uses this single shared, shared kernel. If False (default) each chain samples with its own independently-tuned kernel.

• auto_correlation_settings (AutoCorrelationsSettings) – Configures the per-parameter ESS-derived integrated auto-correlation diagnostics. check_for_convergence defaults to False here because NUTS is run to a fixed sample budget after warmup.

• iterations_per_full_update (Optional[int]) – Sample chunk size between perform_update calls. Inherited from the autonerves config when None.

• number_of_cores (int) – Currently unused — chains run vmapped on a single device. Kept for API parity with the other MCMC searches.

• 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.
auto_correlations_from	Synthesise the standard AutoCorrelations from BlackJAX's per-param ESS via the canonical identity τ_int = N / ESS.
check_model
copy_with_paths
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 BlackJAX chains 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
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

jax_use = 'required'

gradient = 'uses'

batched = True

honours_gradient_mode = False

posterior_kind = 'chain'

produces_evidence = False

resumable = False

warm_start = 'consumer'

install_extra = 'optional'

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

citation_keys = ('blackjax',)

status = 'stable'

test_mode_budget = {'num_chains': 2, 'num_samples': 20, 'num_warmup': 20}

objective_target = (ObjectiveQuantity.LOG_POSTERIOR, CoordinateSpace.PHYSICAL)

invalid_value = -inf

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 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 BlackJAX chains pickled under search_internal/ as raw samples. NUTS samples are unweighted draws from the posterior, so every weight is 1.

auto_correlations_from(search_internal=None)

Synthesise the standard AutoCorrelations from BlackJAX’s per-param ESS via the canonical identity τ_int = N / ESS. previous_times uses the chain truncated by check_size so the relative-change convergence metric stays meaningful.