mici.stagers module#

Classes for staging sampling of Markov chains.

class mici.stagers.ChainStage(n_iter, adapters, trace_funcs, record_stats)[source]#

Bases: NamedTuple

Parameters of chain stage.

Parameters:

  • n_iter (int) – Number of iterations in chain stage.

  • adapters (dict[str, Iterable[Adapter]]) – Dictionary of adapters to apply to each transition in chain stage.

  • trace_funcs (Optional[tuple[TraceFunction]]) – Functions defining chain variables to trace during chain stage.

  • record_stats (bool) – Whether to record statistics during chain stage.

class mici.stagers.Stager[source]#

Bases: ABC

Abstract chain iteration stager.

abstract stages(n_warm_up_iter, n_main_iter, adapters, trace_funcs, *, trace_warm_up=False)[source]#

Create dictionary specifying labels and parameters of chain sampling stages.

Constructs an ordered dictionary with entries corresponding to the sequence of sampling stages when running chains with one or more initial adaptation stages. The keys of each entry are a string label for the sampling stage and the value a ChainStage named tuple specifying the parameters of the stage.

Parameters:

  • n_warm_up_iter (int) – Number of adaptive warm up iterations per chain. Depending on the adapters specified by the adapters argument the warm up iterations may be split between one or more adaptive stages.

  • n_main_iter (int) – Number of iterations (samples to draw) per chain during main (non-adaptive) sampling stage.

  • trace_funcs (Iterable[TraceFunction]) – Iterable of functions which compute the variables to be recorded at each chain iteration, with each trace function being passed the current state and returning a dictionary of scalar or array values corresponding to the variable(s) to be stored. The keys in the returned dictionaries are used to index the trace arrays in the returned traces dictionary. If a key appears in multiple dictionaries only the the value corresponding to the last trace function to return that key will be stored. By default chains are only traced in the iterations of the final non-adaptive stage - this behaviour can be altered using the trace_warm_up argument.

  • adapters (dict[str, Iterable[Adapter]]) – Dictionary of iterables of mici.adapters.Adapter instances keyed by strings corresponding to the key of the transition in the sampler mici.samplers.MarkovChainMonteCarloMethod.transitions dictionary to apply the adapters to, to use to adaptatively set parameters of the transitions during the adaptive stages of the chains. Note that the adapter updates are applied in the order the adapters appear in the iterables and so if multiple adapters change the same parameter(s) the order will matter.

  • trace_warm_up (bool) – Whether to record chain traces and statistics during warm-up stage iterations (True) or only record traces and statistics in the iterations of the final non-adaptive stage (False, the default).

Returns:

Dictionary specifying chain stages, keyed by labels for stages.

Return type:

dict[str, ChainStage]

class mici.stagers.WarmUpStager[source]#

Bases: Stager

Chain iteration stager with a single adaptive warm up stage.

Sampling is split in to two stages:

  1. An adaptive warm up stage will all adapters active.

  2. A main sampling stage with no adapters active.

Only in the main sampling stage are traces of the chain state recorded by storing the outputs of functions of the sampled chain state after each iteration.

stages(n_warm_up_iter, n_main_iter, adapters, trace_funcs, *, trace_warm_up=False)[source]#

Create dictionary specifying labels and parameters of chain sampling stages.

Constructs an ordered dictionary with entries corresponding to the sequence of sampling stages when running chains with one or more initial adaptation stages. The keys of each entry are a string label for the sampling stage and the value a ChainStage named tuple specifying the parameters of the stage.

Parameters:

  • n_warm_up_iter (int) – Number of adaptive warm up iterations per chain. Depending on the adapters specified by the adapters argument the warm up iterations may be split between one or more adaptive stages.

  • n_main_iter (int) – Number of iterations (samples to draw) per chain during main (non-adaptive) sampling stage.

  • trace_funcs (Iterable[TraceFunction]) – Iterable of functions which compute the variables to be recorded at each chain iteration, with each trace function being passed the current state and returning a dictionary of scalar or array values corresponding to the variable(s) to be stored. The keys in the returned dictionaries are used to index the trace arrays in the returned traces dictionary. If a key appears in multiple dictionaries only the the value corresponding to the last trace function to return that key will be stored. By default chains are only traced in the iterations of the final non-adaptive stage - this behaviour can be altered using the trace_warm_up argument.

  • adapters (dict[str, Iterable[Adapter]]) – Dictionary of iterables of mici.adapters.Adapter instances keyed by strings corresponding to the key of the transition in the sampler mici.samplers.MarkovChainMonteCarloMethod.transitions dictionary to apply the adapters to, to use to adaptatively set parameters of the transitions during the adaptive stages of the chains. Note that the adapter updates are applied in the order the adapters appear in the iterables and so if multiple adapters change the same parameter(s) the order will matter.

  • trace_warm_up (bool) – Whether to record chain traces and statistics during warm-up stage iterations (True) or only record traces and statistics in the iterations of the final non-adaptive stage (False, the default).

Returns:

Dictionary specifying chain stages, keyed by labels for stages.

Return type:

dict[str, ChainStage]

class mici.stagers.WindowedWarmUpStager(n_init_slow_window_iter=25, n_init_fast_stage_iter=75, n_final_fast_stage_iter=50, slow_window_multiplier=2.0)[source]#

Bases: Stager

Chain iteration stager with a hierarchy of adaptive warm up stages.

Following the approach of [Stan](https://mc-stan.org) the adaptive stages are split in to two types - ‘fast’ adaptation stages which adjust only transition parameters which can be adapted quickly using local information and ‘slow’ adaptation stages which addtionally adjust transition parameters which require more global information. The adapters to be used in both the fast and slow adaptation stages will be referred to as the fast adapters and the adapters to use in only the slow adaptation stages the slow adapters. Each adapter self identifies if it is a fast adapter by whether the mici.adapters.Adapter.is_fast attribute is set to True.

The adaptive warm up iterations are split into three stages:

  1. An initial fast adaptive stage with only fast adapters active.

  2. A slow adaptive stage with both slow and fast adapters active.

  3. A final adaptive stage with only fast adapters active.

The slow sampling stage (2) is further split in to a sequence of growing, memoryless windows with the adapter stages reset at the beginning of each window, and the number of iterations in each window increasing (by default doubling). The split of the iterations in each of these stages can be controlled using the keyword arguments n_init_fast_stage_iter, n_init_slow_window_iter, n_final_fast_stage_iter and slow_window_multiplier (see descriptions below).

After the initial adaptive warm up stages a subsequent main sampling stage with no further adaptation is performed. Only in this main sampling stage are traces of the chain state recorded by storing the outputs of functions of the sampled chain state after each iteration.

Parameters:

  • n_init_slow_window_iter (int) – Number of iterations in the initial (smallest) window in the slow adaptation stage. Defaults to 25. If the sum of n_init_slow_window_iter, n_init_fast_stage_iter and n_final_fast_stage_iter is more than n_warm_up_iter then n_init_slow_window_iter is set to approximately 75% of n_warm_up_iter (with a single window being used in the slow adaptation stage in this case).

  • n_init_fast_stage_iter (int) – Number of iterations in the initial fast adaptation stage. Defaults to 75. If the sum of n_init_slow_window_iter, n_init_fast_stage_iter` and n_final_fast_stage_iter is more than n_warm_up_iter then n_init_fast_stage_iter is set to approximately 15% of n_warm_up_iter.

  • n_final_fast_stage_iter (int) – Number of iterations in the final fast adaptation stage. Defaults to 50. If the sum of n_init_slow_window_iter, n_init_fast_stage_iter and n_final_fast_stage_iter is more than n_warm_up_iter then n_init_fast_stage_iter is set to approximately 10% of n_warm_up_iter.

  • slow_window_multiplier (float) – Multiplier by which to increase the number of iterations of each subsequent slow adaptation window by. Defaults to 2 such that each window doubles in size.

stages(n_warm_up_iter, n_main_iter, adapters, trace_funcs, *, trace_warm_up=False)[source]#

Create dictionary specifying labels and parameters of chain sampling stages.

Constructs an ordered dictionary with entries corresponding to the sequence of sampling stages when running chains with one or more initial adaptation stages. The keys of each entry are a string label for the sampling stage and the value a ChainStage named tuple specifying the parameters of the stage.

Parameters:

  • n_warm_up_iter (int) – Number of adaptive warm up iterations per chain. Depending on the adapters specified by the adapters argument the warm up iterations may be split between one or more adaptive stages.

  • n_main_iter (int) – Number of iterations (samples to draw) per chain during main (non-adaptive) sampling stage.

  • trace_funcs (Iterable[TraceFunction]) – Iterable of functions which compute the variables to be recorded at each chain iteration, with each trace function being passed the current state and returning a dictionary of scalar or array values corresponding to the variable(s) to be stored. The keys in the returned dictionaries are used to index the trace arrays in the returned traces dictionary. If a key appears in multiple dictionaries only the the value corresponding to the last trace function to return that key will be stored. By default chains are only traced in the iterations of the final non-adaptive stage - this behaviour can be altered using the trace_warm_up argument.

  • adapters (dict[str, Iterable[Adapter]]) – Dictionary of iterables of mici.adapters.Adapter instances keyed by strings corresponding to the key of the transition in the sampler mici.samplers.MarkovChainMonteCarloMethod.transitions dictionary to apply the adapters to, to use to adaptatively set parameters of the transitions during the adaptive stages of the chains. Note that the adapter updates are applied in the order the adapters appear in the iterables and so if multiple adapters change the same parameter(s) the order will matter.

  • trace_warm_up (bool) – Whether to record chain traces and statistics during warm-up stage iterations (True) or only record traces and statistics in the iterations of the final non-adaptive stage (False, the default).

Returns:

Dictionary specifying chain stages, keyed by labels for stages.

Return type:

dict[str, ChainStage]