The mcmc submodule#

This submodule contains the MCMC class that runs Markov Chain Monte Carlo sampling for MHNs.

class mhn.mcmc.mcmc.MCMC(*, optimizer: ~mhn.optimizers.Optimizer | None = None, mhn_model: ~mhn.model.oMHN | ~mhn.model.cMHN | None = None, data: ~numpy.ndarray | None = None, penalty: ~mhn.optimizers.Penalty | tuple[~typing.Callable] | None = None, log_prior: tuple[~typing.Callable] | None = None, n_chains: int = 10, step_size: ~typing.Literal['auto'] | float | ~numpy.ndarray = 'auto', kernel_class: ~mhn.mcmc.kernels.Kernel = <class 'mhn.mcmc.kernels.MALAKernel'>, thin: int = 100, seed: int | None = None)#

Markov chain Monte Carlo sampler for MHN.

The simplest way to create an MCMC sampler is to provide a trained Optimizer. This already includes the dataset, the regularization strength lambda, the penalty and its gradient.

This is enough to create an RWM or MALA sampler. For an smMALA sampler, the Hessian of the penalty is also needed, which is not stored in the optimizer. In this case, please provide the penalty and its derivatives directly via penalty or log_prior.

Alternatively, you can create an MCMC sampler from a trained MHN model, the data used to train it, and the penalty used for training. Depending on the Kernel (RWM, MALA, or smMALA), you have to provide the penalty, its gradient, and its Hessian,

  • by providing penalty either as a Penalty (only for RWM/MALA) or as a tuple of callables of length 1, 2, or 3, or

  • by providing log_prior as a tuple of callables of length 1, 2, or 3.

The difference between penalty and log_prior is that the former is unscaled by \(\lambda\) and positive (such as it is stored in an Optimizer object), while the latter is scaled by lambda and negative.

When using a custom penalty or prior, initial chain values must be set manually by setting Sampler.initial_step to an array of shape (n_chains, 1, m) with m the number of parameters.

Parameters:

  • optimizer (Optimizer, optional) – Trained Optimizer.

  • mhn_model (oMHN | cMHN, optional) – MHN model.

  • data (np.ndarray, optional) – Data used to train the MHN model.

  • penalty (Penalty | tuple[Callable[[np.ndarray], float], Callable[[np.ndarray], np.ndarray]], optional) – Penalty used during training. If not Penalty, penalty[0] gives the penalty (positive and unscaled by \(\lambda\)), penalty[1] its gradient and penalty[2] its Hessian.

  • log_prior – Log-prior used during training. log_prior[0] gives the log_prior (negative and scaled by \(\lambda\)), penalty[1] its gradient and penalty[2] its Hessian.

  • n_chains (int, optional) – Number of parallel chains to run. Defaults to 10.

  • step_size – (Literal[“auto”] | int | np.ndarray | ) Stepsize of the Kernel. See the documentation of the kernels (RWMKernel, MALAKernel, smMALAKernel) for more details. If "auto", the stepsize is tuned automatically before the first run using tune_stepsize(). If array, every chain uses a different step size. Defaults to "auto".

  • kernel_class (Kernel, optional) – Kernel class to use for MCMC sampling. Must be one of RWMKernel, MALAKernel, `smMALAKernel. Defaults to MALAKernel.

  • thin (int, optional) – Thinning factor for MCMC sampling. Defaults to 100.

  • seed (int, optional) – Random seed for reproducibility. Defaults to None.

acceptance(burn_in: int | float = 0.2, chain_id: int | None = None) → ndarray | float#

Calculate the acceptance rate for the chains.

Parameters:

  • burn_in (int | float, optional) – Number of steps to discard as burn-in. If a float, it is interpreted as a fraction of the total steps. Defaults to 0.2.

  • chain_id (int | None, optional) – Chain ID to return acceptance rate for. If None, returns acceptance rates for all chains. Defaults to None.

Returns:

The acceptance rate(s) for the specified chain(s).

Return type:

np.ndarray | float

ess(burn_in: int | float = 0.2, **kwargs)#

Calculate the effective sample size (ESS).

Parameters:

  • burn_in (int | float, optional) – Number of steps to discard as burn-in. If a float, it is interpreted as a fraction of the total steps. Defaults to 0.2.

  • **kwargs – Additional keyword arguments to pass to arviz.ess(). See the ArviZ documentation for more details.

Returns:

The effective sample size for each parameter.

Return type:

np.ndarray

rhat(burn_in: int | float = 0.2, **kwargs) → ndarray#

Calculate the Gelman-Rubin potential scale reduction factor \(\hat R\).

Parameters:

  • burn_in (int | float, optional) – Number of steps to discard as burn-in. If a float, it is interpreted as a fraction of the total steps. Defaults to 0.2.

  • **kwargs – Additional keyword arguments to pass to arviz.rhat(). See the ArviZ documentation for more details.

Returns:

The Gelman-Rubin R-hat values for each parameter.

Return type:

np.ndarray

run(stopping_crit: Literal['r_hat', 'ESS'] | Callable | None = 'r_hat', max_steps: int | None = None, check_interval: int | None = 1000, burn_in: int | float = 0.2, verbose: bool = True) → ndarray#

Run MCMC sampling until the stopping criterion is met or the maximum number of steps is reached.

Parameters:

  • stopping_crit (Literal["r_hat", "ESS"] | Callable | None, optional) – The stopping criterion to use. If "r_hat", runs until the the Gelman-Rubin potential scale reduction factor \(\hat R\) is below 1.01. If "ESS", runs until the effective sample size (ESS) is above 100. If a callable, it should take in the log_thetas and return a boolean indicating whether to stop. If None, runs until max_steps is reached. Burn-in is discarded before checking the stopping criterion and only every check_interval steps. Defaults to "r_hat".

  • max_steps (int | None, optional) – Maximum number of steps to run. If None, runs a maximum of 1,000,000 steps for RWM and MALA kernels and 100,000 steps for smMALA kernels. Defaults to None.

  • check_interval (int | None, optional) – Number of steps between checking the stopping criterion. Defaults to 1000.

  • burn_in (int | float, optional) – Number of steps to discard as burn-in. If a float, it is interpreted as a fraction of the total steps. Defaults to 0.2.

  • verbose (bool, optional) – Whether to print progress. Defaults to True.

Returns:

The log-thetas for each chain.

Return type:

np.ndarray

tune_stepsize(n_steps: int = 100, burn_in: float | int = 0.6, target_acceptance: float | Literal['auto'] = 'auto', max_trials: int = 10, verbose: bool = True, tol: float = 0.02) → float#

Automatically infer an appropriate step size epsilon for MCMC sampling.

Parameters:

  • n_steps (int, optional) – Number of steps to run for inference. Defaults to 100.

  • burn_in (float | int, optional) – Burn-in period. If float, fraction of n_steps. If int, number of steps. Defaults to 0.6.

  • target_acceptance (float | Literal["auto"], optional) – Target acceptance rate. If “auto”, set to 0.234 for RWM kernels, 0.574 for MALA kernels and 0.7 for smMALA kernels. Defaults to “auto”.

  • max_trials (int, optional) – Maximum number of trials to run. Defaults to 10.

  • verbose (bool, optional) – Whether to print progress. Defaults to True.

  • tol (float, optional) – Tolerance for acceptance rate. If the acceptance rate is within tol of the target acceptance rate, the step size is accepted. Defaults to 0.02.

Returns:

Inferred step sizet.

Return type:

float