bartz.mcmcloop.Callback

class bartz.mcmcloop.Callback(*args, **kwargs)[source]

Callback type for run_mcmc.

__call__(*, key, state, burnin, i_total, callback_state, n_burn, n_save, n_skip, i_outer, inner_loop_length)[source]

Do an arbitrary action after an iteration of the MCMC.

Parameters:

  • key (Key[Array, '']) – A key for random number generation.

  • state (State) – The MCMC state just after updating it.

  • burnin (Bool[Array, '']) – Whether the last iteration was in the burn-in phase.

  • i_total (Int32[Array, '']) – The index of the last MCMC iteration (0-based).

  • callback_state (PyTree[Any, 'T']) – The callback state, initially set to the argument passed to run_mcmc, afterwards to the value returned by the last invocation of the callback.

  • n_burn (Int32[Array, ''])

  • n_save (Int32[Array, ''])

  • n_skip (Int32[Array, '']) – The corresponding run_mcmc arguments as-is.

  • i_outer (Int32[Array, '']) – The index of the last outer loop iteration (0-based).

  • inner_loop_length (Int32[Array, '']) – The number of MCMC iterations in the inner loop.

Returns:

  • state (State) – A possibly modified MCMC state. To avoid modifying the state, return the state argument passed to the callback as-is.

  • callback_state (CallbackState) – The new state to be passed on the next callback invocation.

Notes

For convenience, the callback may return None, and the states won’t be updated.