autofit.BlackJAXNUTS#

class BlackJAXNUTS[source]#

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.

  • 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.

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 BlackJAX chain pickled under search_internal/ into a standard SamplesMCMC.

start_resume_fit

Attributes

backend

Load the pickled search-internal dict written by _fit.

backend_filename

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()[source]#

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 backend_filename#
property backend: dict#

Load the pickled search-internal dict written by _fit.

output_search_internal(search_internal)[source]#

Pickle the search-internal dict.

BlackJAX has no native on-disk format (cf. emcee’s HDFBackend), so we round-trip the chain + diagnostics via pickle. We bypass self.paths.save_search_internal because the autofit dill path chokes on a few numpy/jax-backed members; 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)[source]#
samples_via_internal_from(model, search_internal=None)[source]#

Convert the BlackJAX chain pickled under search_internal/ into a standard SamplesMCMC. NUTS samples are unweighted draws from the posterior, so weights are 1.0.

auto_correlations_from(search_internal=None)[source]#

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.