IWLSKernel

Contents

IWLSKernel#

class liesel.goose.IWLSKernel(position_keys, chol_info_fn=None, initial_step_size=0.01, da_tune_step_size=True, da_target_accept=0.8, da_gamma=0.05, da_kappa=0.75, da_t0=10, identifier='', fallback_chol_info='identity')[source]#

Bases: ModelMixin, TransitionMixin[IWLSKernelState, DefaultTransitionInfo], ReprMixin

An IWLS kernel with dual averaging and an (optional) user-defined function for computing the Cholesky decomposition of the Fisher information matrix, implementing the liesel.goose.types.Kernel protocol.

Parameters:

  • position_keys (Sequence[str]) – Sequence of position keys (variable names) handled by this kernel.

  • chol_info_fn (Callable[[Any], Any] | None, default: None) – A custom function that takes a model state and returns the Cholesky decomposition of the information matrix to produce the IWLS proposal. By default, this will be the Cholesky decomposition of the observed negative hessian at the current values, i.e. the current observed information.

  • initial_step_size (float, default: 0.01) – Value at which to start step size tuning.

  • da_tune_step_step_size – Whether to tune the step size using dual averaging.

  • da_target_accept (float, default: 0.8) – Target acceptance probability for dual averaging algorithm.

  • da_gamma (float, default: 0.05) – The adaptation regularization scale.

  • da_kappa (float, default: 0.75) – The adaptation relaxation exponent.

  • da_t0 (int, default: 10) – The adaptation iteration offset.

  • identifier (str, default: '') – An string acting as a unique identifier for this kernel.

  • fallback_chol_info (Literal['identity', 'chol_of_modified_info'] | None, default: 'identity') – What do do if the Cholesky decomposition of the observed information matrix fails. If "identity", uses an identity matrix as the Cholesky factor. If "chol_of_modified_info", performs an eigendecomposition of the negative Hessian and clips the eigenvalues to 1e-5. This can be interpreted as replacing the observed negative Hessian with a very similar positive definite matrix. This is slow, because it performs an eigendecomposition and two cholesky factorizations. If None, does nothing.

Notes

For more information on step size tuning via dual averaging, see da_step() and DAKernelState.

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.

init_state(prng_key, model_state)

Initializes the kernel state.

start_epoch(prng_key, kernel_state, ...)

Resets the state of the dual averaging algorithm.

tune(prng_key, kernel_state, model_state, epoch)

Currently does nothing.

untuned(position_keys[, chol_info_fn, ...])

Initializes an IWLS kernel that does not conduct step size tuning during warmup.

Attributes

error_book

Dict of error codes and their meaning.

fallback_chol_info

identifier

Kernel identifier, set by EngineBuilder

needs_history

Whether this kernel needs its history for tuning.

position_keys

Tuple of position keys handled by this kernel.
Contents