mici package#

MCMC samplers based on simulating Hamiltonian dynamics on a manifold.

mici.sample_constrained_hmc_chains(n_warm_up_iter, n_main_iter, init_states, neg_log_dens, constr, *, backend=None, seed=None, grad_neg_log_dens=None, jacob_constr=None, mhp_constr=None, dens_wrt_hausdorff=True, system_class=<class 'mici.systems.DenseConstrainedEuclideanMetricSystem'>, integrator_class=<class 'mici.integrators.ConstrainedLeapfrogIntegrator'>, sampler_class=<class 'mici.samplers.DynamicMultinomialHMC'>, system_kwargs=None, integrator_kwargs=None, sampler_kwargs=None, **kwargs)[source]#

Sample constrained Hamiltonian Monte Carlo chains for a given target distribution.

Samples one or more Markov chains with given initial states, with a stationary distribution on an implicitly-defined manifold embedded in an ambient Euclidean space, specified by functions evaluating the negative log density of the target distribution and specified by a constraint function for which the zero level-set of specifies the manifold the distribution is supported on. Each chain has zero or more adaptive warm-up iterations, during which the parameters of the chain transitions can be automatically tuned.

This function allows changing the types of the system, integrator and sampler classes used from their defaults and passing additional keyword arguments to their initialisers. For finer-grained control and more complex use cases it is recommended to initialise the objects directly and use the sample_chains method of the resulting sampler object.

Parameters:

n_warm_up_iter (int) – Number of adaptive warm up iterations per chain. Depending on the stagers.Stager instance specified by the stager argument the warm up iterations may be split between one or more adaptive stages. If zero, only a single non-adaptive stage is used.
n_main_iter (int) – Number of iterations (samples to draw) per chain during main (non-adaptive) sampling stage.
init_states (Iterable[ChainState | np.typing.NDArray]) – Initial chain states. Each state can be either an array specifying the state position component or a states.ChainState instance. If an array is passed or the mom attribute of the state is not set, a momentum component will be independently sampled from its conditional distribution. One chain will be run for each state in the iterable.
neg_log_dens (ScalarFunction) – Function which given a position array returns the negative logarithm of an unnormalized probability density on the position space with respect to a reference measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
constr (ArrayFunction) – Function which given a position array return as a 1D array the value of the (vector-valued) constraint function, the zero level-set of which implicitly defines the manifold the dynamic is simulated on.
backend (str | None) – Name of automatic differentiation backend to use. See autodiff subpackage documentation for details of available backends. If None (the default) no automatic differentiation fallback will be used and so all derivative functions must be specified explicitly.
seed (np.random.Generator | int | None) – Integer value to seed NumPy random generator with, or None to use default (non-fixed) seeding scheme, or an already seeded numpy.random.Generator instance.
grad_neg_log_dens (GradientFunction | None) – Function which given a position array returns the derivative of neg_log_dens with respect to the position array argument. Optionally the function may instead return a 2-tuple of values with the first being the array corresponding to the derivative and the second being the value of the neg_log_dens evaluated at the passed position array. If None is passed (the default) an automatic differentiation fallback will be used to attempt to construct the derivative of neg_log_dens automatically.
jacob_constr (JacobianFunction | None) – Function which given a position array computes the Jacobian (matrix / 2D array of partial derivatives) of the output of the constraint function c = constr(q) with respect to the position array argument q, returning the computed Jacobian as a 2D array jacob with jacob[i, j] = ∂c[i] / ∂q[j]. Optionally the function may instead return a 2-tuple of values with the first being the array corresponding to the Jacobian and the second being the value of constr evaluated at the passed position array. If None is passed (the default) an automatic differentiation fallback will be used to attempt to construct a function to compute the Jacobian (and value) of constr automatically.
mhp_constr (MatrixHessianProductFunction | None) – Function which given a position array returns another function which takes a 2D array as an argument and returns the matrix-Hessian-product (MHP) of the constraint function constr with respect to the position array argument. The MHP is here defined as a function of a (dim_constr, dim_pos) shaped 2D array m as mhp(m) = sum(m[:, :, None] * hess[:, :, :], axis=(0, 1)) where hess is the (dim_constr, dim_pos, dim_pos) shaped vector-Hessian of c = constr(q) with respect to q i.e. the array of second-order partial derivatives of such that hess[i, j, k] = ∂²c[i] / (∂q[j] ∂q[k]). Optionally the function may instead return a 3-tuple of values with the first a function to compute a MHP of constr, the second a 2D array corresponding to the Jacobian of constr, and the third the value of constr, all evaluated at the passed position array. If None is passed (the default) an automatic differentiation fallback will be used to attempt to construct a function which calculates the MHP (and Jacobian and value) of constr automatically. Only used if dens_wrt_hausdorff == False.
dens_wrt_hausdorff (bool) – If constr is specified, whether the neg_log_dens function specifies the (negative logarithm) of the density of the target distribution with respect to the Hausdorff measure on the manifold directly (True) or alternatively the negative logarithm of a density of a prior distriubtion on the unconstrained (ambient) position space with respect to the Lebesgue measure, with the target distribution then corresponding to the posterior distribution when conditioning on the event constr(pos) == 0 (False).
system_class (type[System]) – The Hamiltonian system class to use. One of the subclasses of systems.System in the systems module.
integrator_class (type[Integrator]) – The symplectic integrator class to use. One of the subclasses of integrators.Integrator in the integrators module.
sampler_class (type[HamiltonianMonteCarlo]) – The Hamiltonian Monte Carlo sampler class to use. One of the subclasses of samplers.HamiltonianMonteCarlo class in the samplers module.
system_kwargs (dict | None) – Any additional keyword arguments to system class initialiser.
integrator_kwargs (dict | None) – Any additional keyword arguments to integrator class intitialiser.
sampler_kwargs (dict | None) – Any additional keyword arguments to sampler class initialiser.

Keyword Arguments:

**kwargs – Additional keyword arguments to pass to the samplers.HamiltonianMonteCarlo.sample_chains() method called to sample the chains.

Returns:

Named tuple (final_states, traces, statistics) corresponding to states of chains after final iteration, dictionary of chain trace arrays and dictionary of chain statistics dictionaries.

Return type:

HMCSampleChainsOutputs

mici.sample_hmc_chains(n_warm_up_iter, n_main_iter, init_states, neg_log_dens, *, backend=None, seed=None, grad_neg_log_dens=None, system_class=<class 'mici.systems.EuclideanMetricSystem'>, integrator_class=<class 'mici.integrators.LeapfrogIntegrator'>, sampler_class=<class 'mici.samplers.DynamicMultinomialHMC'>, system_kwargs=None, integrator_kwargs=None, sampler_kwargs=None, **kwargs)[source]#

Sample Hamiltonian Monte Carlo chains for a given target distribution.

Samples one or more Markov chains with given initial states, with a stationary distribution specified by a function evaluating the negative log density of the target distribution on an unconstrained Euclidean space. Each chain has zero or more adaptive warm-up iterations, during which the parameters of the chain transitions can be automatically tuned.