NUTSKernel#

class liesel.goose.nuts.NUTSKernel(position_keys, initial_step_size=None, initial_inverse_mass_matrix=None, max_treedepth=10, da_target_accept=0.8, da_gamma=0.05, da_kappa=0.75, da_t0=10, mm_diag=True)[source]#

Bases: ModelMixin, TransitionMixin[NUTSKernelState, NUTSTransitionInfo], TuningMixin[NUTSKernelState, DefaultTuningInfo]

A NUTS kernel with dual averaging and an inverse mass matrix tuner, implementing the Kernel protocol.

Methods

end_epoch(prng_key, kernel_state, ...)

Sets the step size as found by the dual averaging algorithm.

end_warmup(prng_key, kernel_state, ...)

Currently does nothing.

has_model()

Whether the model interface is set.

init_state(prng_key, model_state)

Initializes the kernel state with an identity inverse mass matrix and a reasonable step size (unless explicit arguments were provided by the user).

log_prob_fn(model_state)

Returns the log-probability function with the position as the only argument.

position(model_state)

Extracts the position from a model state.

set_model(model)

Sets the model interface.

start_epoch(prng_key, kernel_state, ...)

Resets the state of the dual averaging algorithm.

transition(prng_key, kernel_state, ...)

Handles one transition.

tune(prng_key, kernel_state, model_state, ...)

The method can perform automatic tuning of the kernel and is called after each adaptation epoch.

Attributes

error_book

Dict of error codes and their meaning.

identifier

Kernel identifier, set by EngineBuilder

model

Returns the model interface if it is set.

needs_history

Whether this kernel needs its history for tuning.

position_keys

Tuple of position keys handled by this kernel.