bartz.mcmcloop.make_tqdm_callback

bartz.mcmcloop.make_tqdm_callback(state, *, update_every=1, report_every=100, average=True, **tqdm_kwargs)[source]

Prepare a tqdm progress-bar callback for run_mcmc.

The callback shows a progress bar that advances with the MCMC iterations, optionally annotated with the proposal acceptance statistics.

Parameters:

  • state (State) – The MCMC state to use the callback with, used to determine device sharding.

  • update_every (int, default: 1) – The bar position is refreshed every update_every MCMC iterations (tqdm further throttles the actual redraw rate on its own).

  • report_every (int | None, default: 100) – The acceptance statistics shown next to the bar are refreshed every report_every MCMC iterations, None to omit them.

  • average (bool, default: True) – If True, the statistics shown are averaged over the iterations since the previous refresh; if False, they reflect the current iteration only. Ignored when report_every is None.

  • **tqdm_kwargs (Any) – Additional keyword arguments forwarded to the tqdm.tqdm constructor, e.g., desc, file, or disable.

Returns:

dict[str, Any] – A dictionary with the arguments to pass to run_mcmc as keyword arguments to set up the callback.

Notes

Works with chains sharded across multiple devices. If the run is interrupted (e.g. with ^C), the bar is left as-is; the next make_tqdm_callback call closes it, so a subsequent run starts from a clean line.

Examples

>>> run_mcmc(key, state, ..., **make_tqdm_callback(state, ...))