nifty.re.blackjax module#

class LogDensity(likelihood)[source]#

Bases: Model

__init__(likelihood)[source]#
likelihood: Likelihood#
blackjax_nuts(likelihood: Likelihood, position, key, n_warmup_steps=1000, n_samples=1000, estimate_effective_sample_size=True, state_and_parameters=None)[source]#

Run BlackJAX’s NUTS sampler with optional windowed adaptation for step size and mass matrix, unless these are provided via state_and_parameters.

Parameters:

  • likelihood (Likelihood) – Negative log-likelihood of the model.

  • position (PyTree) – Initial position in parameter space.

  • key (jax.random.PRNGKey) – JAX random key for warm-up and sampling.

  • n_warmup_steps (int, default=1000) – Number of warm-up (adaptation) steps used to tune the NUTS sampler. Ignored if state_and_parameters is provided.

  • n_samples (int, default=1000) – Number of MCMC samples to draw after warm-up.

  • estimate_effective_sample_size (bool, default=True) – Whether to estimate and report the effective sample size in the status message.

  • state_and_parameters (tuple, optional) – A tuple (state, parameters) returned from a previous call. If provided, warm-up is skipped and sampling resumes from this state using the given NUTS parameters. When resuming you need to provide a fresh random key, otherwise the same random numbers will be used.

Returns:

  • samples (Samples) – Container holding the drawn HMC samples.

  • state_and_parameters (tuple) – A tuple (final_state, parameters) that can be passed back into this function to resume sampling without re-running warm-up.

get_sample_size_estimate(samples)[source]#
get_status_message(samples, state, residual=None, estimate_effective_sample_size=True, *, name='', map='lmap') → str[source]#