mici.systems module#

Hamiltonian systems encapsulating energy functions and their derivatives.

class mici.systems.CholeskyFactoredRiemannianMetricSystem(neg_log_dens, metric_chol_func, *, vjp_metric_chol_func=None, grad_neg_log_dens=None)[source]#

Bases: RiemannianMetricSystem

Riemannian-metric system with Cholesky-factored matrix representation.

Hamiltonian system with a position dependent metric matrix representation which is specified by its Cholesky factor by a matrix function metric_chol_func of the position q which outputs a lower-triangular matrix L = metric_chol_func(q) with the metric matrix representation then taken to be L @ L.T.

See documentation of RiemannianMetricSystem for more general details about Riemannian-metric Hamiltonian systems.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric_chol_func (ArrayFunction) – Function which given a position array returns a 2D array with zeros above the diagonal corresponding to the lower-triangular Cholesky-factor of the positive definite metric matrix representation.
vjp_metric_chol_func (Optional[VectorJacobianProductFunction]) –
Function which given a position array returns another function which takes a lower-triangular 2D array as an argument (any values in the array above the diagonal are ignored) and returns the vector-Jacobian-product (VJP) of metric_chol_func with respect to the position array argument. The VJP is here defined as a function of a 2D array v.
```
vjp(v) = sum(v[:, :, None] * jacob[:, :, :], axis=(0, 1))
```
where jacob is the (dim_pos, dim_pos, dim_pos) shaped Jacobian of L = metric_chol_func(q) with respect to q i.e. the array of partial derivatives of the function such that
```
jacob[i, j, k] = ∂L[i, j] / ∂q[k]
```
Optionally the function may instead return a 2-tuple of values with the first a function to compute a VJP of metric_chol_func and the second a 2D array containing the value of metric_chol_func, both 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 VJP (and value) of metric_chol_func automatically.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

metric(state)#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: VectorJacobianProduct

class mici.systems.ConstrainedEuclideanMetricSystem(neg_log_dens, constr, *, metric=None, dens_wrt_hausdorff=True, grad_neg_log_dens=None, jacob_constr=None)[source]#

Bases: ConstrainedTractableFlowSystem, EuclideanMetricSystem

Base class for Euclidean Hamiltonian systems subject to constraints.

The (constrained) position space is assumed to be a differentiable manifold embedded with a \(Q\)-dimensional ambient Euclidean space. The \(Q-C\) dimensional manifold \(\mathcal{M}\) is implicitly defined by an equation \(\mathcal{M} = \lbrace q \in \mathbb{R}^Q : c(q) = 0 \rbrace\) with \(c: \mathbb{R}^Q \to \mathbb{R}^C\) the constraint function.

The ambient Euclidean space is assumed to be equipped with a metric with constant positive-definite matrix representation \(M\) which further specifies the covariance of the zero-mean Gaussian distribution \(\mathcal{N}(0, M)\) on the unconstrained momentum (co-)vector \(p\) with corresponding \(h_2\) Hamiltonian component defined as

\[h_2(q, p) = \frac{1}{2} p^T M^{-1} p.\]

The time-derivative of the constraint equation implies a further set of constraints on the momentum \(p\) with \(\partial c(q) M^{-1} p = 0\) at all time points, corresponding to the momentum (velocity) being in the co-tangent space (tangent space) to the manifold.

The target distribution is either assumed to be directly specified with unnormalized density \(\exp(-\ell(q))\) with respect to the Hausdorff measure on the manifold (under the metric induced from the ambient metric) with in this case the \(h_1\) Hamiltonian component then simply

\[h_1(q) = \ell(q),\]

or alternatively it is assumed a prior distribution on the position \(q\) with density \(\exp(-\ell(q))\) with respect to the Lebesgue measure on the ambient space is specifed and the target distribution is the posterior distribution on \(q\) when conditioning on the event \(c(q) = 0\). The negative logarithm of the posterior distribution density with respect to the Hausdorff measure (and so \(h_1\) Hamiltonian component) is then

\[h_1(q) = \ell(q) + \frac{1}{2} \log\left|\partial c(q)M^{-1}\partial c(q)^T\right|\]

with an additional second Gram matrix determinant term to give the correct density with respect to the Hausdorff measure on the manifold.

Due to the requirement to enforce the constraints on the position and momentum, a constraint-preserving numerical integrator needs to be used when simulating the Hamiltonian dynamic associated with the system, e.g. mici.integrators.ConstrainedLeapfrogIntegrator.

References

Lelièvre, T., Rousset, M. and Stoltz, G., 2019. Hybrid Monte Carlo methods for sampling probability measures on submanifolds. Numerische Mathematik, 143(2), pp.379-421.
Graham, M.M. and Storkey, A.J., 2017. Asymptotically exact inference in differentiable generative models. Electronic Journal of Statistics, 11(2), pp.5105-5164.

Parameters:

neg_log_dens (ScalarFunction) – Function which given a position array returns the negative logarithm of an unnormalized probability density on the constrained position space with respect to the Hausdorff measure on the constraint manifold (if dens_wrt_hausdorff == True) or alternatively the negative logarithm of an unnormalized probability density on the unconstrained (ambient) position space with respect to the Lebesgue measure. In the former case the target distribution it is wished to draw approximate samples from is assumed to be directly specified by the density function on the manifold. In the latter case the density function is instead taken to specify a prior distribution on the ambient space with the target distribution then corresponding to the posterior distribution when conditioning on the (zero Lebesgue measure) event constr(q) == 0 where q is the position array. This target posterior distribution has support on the differentiable manifold implicitly defined by the constraint equation, with density with respect to the Hausdorff measure on the manifold corresponding to the ratio of the prior density (specified by neg_log_dens) and the square-root of the determinant of the Gram matrix defined by gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.
constr (ArrayFunction) – Function which given a position rray 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.
metric (Optional[MetricLike]) – Matrix object corresponding to matrix representation of metric on unconstrained position space and covariance of Gaussian marginal distribution on unconstrained momentum vector. If None is passed (the default), the identity matrix will be used. If a 1D array is passed then this is assumed to specify a metric with positive diagonal matrix representation and the array the matrix diagonal. If a 2D array is passed then this is assumed to specify a metric with a dense positive definite matrix representation specified by the array. Otherwise if the value is a mici.matrices.PositiveDefiniteMatrix subclass it is assumed to directly specify the metric matrix representation.
dens_wrt_hausdorff (bool) – 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). Note that in the former case the base Hausdorff measure on the manifold depends on the metric defined on the ambient space, with the Hausdorff measure being defined with respect to the metric induced on the manifold from this ambient metric.
grad_neg_log_dens (Optional[GradientFunction]) – 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 a function to compute the derivative (and value) of neg_log_dens automatically.
jacob_constr (Optional[JacobianFunction]) – 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.

constr(state)[source]#

Constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of constr(state.pos) as 1D array.
Return type:: ArrayLike

dh1_dpos(state)[source]#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh2_flow_dmom(state, dt)[source]#

Derivatives of h2_flow() flow map with respect to momentum argument.

Parameters:

state (ChainState) – State to evaluate derivatives of flow map at.
dt (ScalarLike) – Time interval flow simulated for.

Returns:

Tuple (dpos_dmom, dmom_dmom) with dpos_dmom a matrix representing derivative (Jacobian) of position output of h2_flow() with respect to the value of the momentum component of the initial input state and dmom_dmom a matrix representing derivative (Jacobian) of momentum output of h2_flow() with respect to the value of the momentum component of the initial input state.

Return type:

tuple[Matrix, Matrix]

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

abstract grad_log_det_sqrt_gram(state)[source]#

Derivative of half log-determinant of Gram matrix with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of log_det_sqrt_gram(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

gram(state)[source]#

Gram matrix at current position.

The Gram matrix as a position q is defined as

gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T

where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)[source]#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

h2_flow(state, dt)#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

inv_gram(state)[source]#

Inverse of Gram matrix at current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Inverse of Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

jacob_constr(state)[source]#

Jacobian of constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Jacobian of constr(state.pos) as 2D array.
Return type:: ArrayLike

abstract jacob_constr_inner_product(jacob_constr_1, inner_product_matrix, jacob_constr_2=None)#

Compute inner product of rows of constraint Jacobian matrices.

Computes jacob_constr_1 @ inner_product_matrix @ jacob_constr_2.T potentially exploiting any structure / sparsity in jacob_constr_1, jacob_constr_2 and inner_product_matrix.

Parameters:

jacob_constr_1 (MatrixLike) – First constraint Jacobian in product.
inner_product_matrix (PositiveDefiniteMatrix) – Positive-definite matrix defining inner-product between rows of two constraint Jacobians.
jacob_constr_2 (Optional[MatrixLike]) – Second constraint Jacobian in product. Defaults to jacob_constr_1 if set to None.

Returns:

Object corresponding to computed inner products of the constraint Jacobian rows.

Return type:

MatrixLike

log_det_sqrt_gram(state)[source]#

Value of (half of) log-determinant of Gram matrix.

Parameters:: state (ChainState) –
Return type:: ScalarLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

project_onto_cotangent_space(mom, state)[source]#

Project a momentum on to the co-tangent space at a position.

Parameters:

mom (ArrayLike) – Momentum (co-)vector as 1D array to project on to co-tangent space.
state (ChainState) – State definining position on the manifold to project in to the co-tangent space of.

Returns:

Projected momentum in the co-tangent space at state.pos.

Return type:

ArrayLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.ConstrainedTractableFlowSystem(neg_log_dens, *, grad_neg_log_dens=None)[source]#

Bases: TractableFlowSystem

Base class for Hamiltonian systems subject to constraints with tractable flows.

\[\mathcal{M} = \lbrace q \in \mathbb{R}^Q : c(q) = 0 \rbrace\]

with \(c: \mathbb{R}^Q \to \mathbb{R}^C\) the differentiable and surjective vector-valued constraint function.

The Hamiltonian function \(h\) is assumed to have the general form

\[h(q, p) = h_1(q) + h_2(q, p)\]

where \(q\) and \(p\) are the position and momentum variables respectively, and \(h_1\) and \(h_2\) Hamiltonian component functions. The exact unconstrained Hamiltonian flows for both the \(h_1\) and \(h_2\) components, respectively \(\Phi_1\) and \(\Phi_2\) are assumed to be tractable for subclasses of this class.

The constrained Hamiltonian dynamics are described by the system of differential algebraic equations

\[\dot{q} = \nabla_2 h(q, p), \quad \dot{p} = -\nabla_1 h(q, p) - \partial c(q)^T\lambda, \quad c(q) = 0.\]

where \(\lambda\) is a set of Lagrange multipliers of dimension equal to number of constraints, \(C\), and which are implicitly defined by the condition that the constraint equation \(c(q) = 0\) applies at all times.

The dynamics implicitly define a set of constraints on the momentum variables, differentiating the constraint equation with respect to time giving that

\[\partial c(q) \nabla_2 h(q, p) = \partial c(q) \nabla_2 h_2(q, p) = 0.\]

The set of momentum variables satisfying the above for given position variables is termed the cotangent space of the manifold (at a position). Here we assume that the operation of projecting a momentum vector onto the cotangent space at a given position is tractable to compute.

Parameters:

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.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

abstract constr(state)[source]#

Constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of constr(state.pos) as 1D array.
Return type:: ArrayLike

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ScalarLike

abstract dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

abstract dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

abstract dh2_flow_dmom(state, dt)[source]#

Derivatives of h2_flow() flow map with respect to momentum argument.

Parameters:

state (ChainState) – State to evaluate derivatives of flow map at.
dt (ScalarLike) – Time interval flow simulated for.

Returns:

Return type:

tuple[Matrix, Matrix]

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

abstract h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

abstract h2_flow(state, dt)#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

abstract jacob_constr(state)[source]#

Jacobian of constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Jacobian of constr(state.pos) as 2D array.
Return type:: ArrayLike

abstract jacob_constr_inner_product(jacob_constr_1, inner_product_matrix, jacob_constr_2=None)[source]#

Compute inner product of rows of constraint Jacobian matrices.

Computes jacob_constr_1 @ inner_product_matrix @ jacob_constr_2.T potentially exploiting any structure / sparsity in jacob_constr_1, jacob_constr_2 and inner_product_matrix.

Parameters:

jacob_constr_1 (MatrixLike) – First constraint Jacobian in product.
inner_product_matrix (PositiveDefiniteMatrix) – Positive-definite matrix defining inner-product between rows of two constraint Jacobians.
jacob_constr_2 (Optional[MatrixLike]) – Second constraint Jacobian in product. Defaults to jacob_constr_1 if set to None.

Returns:

Object corresponding to computed inner products of the constraint Jacobian rows.

Return type:

MatrixLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

abstract project_onto_cotangent_space(mom, state)[source]#

Project a momentum on to the co-tangent space at a position.

Parameters:

mom (ArrayLike) – Momentum (co-)vector as 1D array to project on to co-tangent space.
state (ChainState) – State definining position on the manifold to project in to the co-tangent space of.

Returns:

Projected momentum in the co-tangent space at state.pos.

Return type:

ArrayLike

sample_momentum(state, rng)[source]#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.DenseConstrainedEuclideanMetricSystem(neg_log_dens, constr, *, metric=None, dens_wrt_hausdorff=True, grad_neg_log_dens=None, jacob_constr=None, mhp_constr=None)[source]#

Bases: ConstrainedEuclideanMetricSystem

Euclidean Hamiltonian system subject to a dense set of constraints.

See ConstrainedEuclideanMetricSystem for more details about constrained systems.

Parameters:

neg_log_dens (ScalarFunction) – Function which given a position array returns the negative logarithm of an unnormalized probability density on the constrained position space with respect to the Hausdorff measure on the constraint manifold (if dens_wrt_hausdorff == True) or alternatively the negative logarithm of an unnormalized probability density on the unconstrained (ambient) position space with respect to the Lebesgue measure. In the former case the target distribution it is wished to draw approximate samples from is assumed to be directly specified by the density function on the manifold. In the latter case the density function is instead taken to specify a prior distribution on the ambient space with the target distribution then corresponding to the posterior distribution when conditioning on the (zero Lebesgue measure) event constr(q) == 0 where q is the position array. This target posterior distribution has support on the differentiable manifold implicitly defined by the constraint equation, with density with respect to the Hausdorff measure on the manifold corresponding to the ratio of the prior density (specified by neg_log_dens) and the square-root of the determinant of the Gram matrix defined by gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.
constr (ArrayFunction) – Function which given a position rray 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.
metric (Optional[MetricLike]) – Matrix object corresponding to matrix representation of metric on unconstrained position space and covariance of Gaussian marginal distribution on unconstrained momentum vector. If None is passed (the default), the identity matrix will be used. If a 1D array is passed then this is assumed to specify a metric with positive diagonal matrix representation and the array the matrix diagonal. If a 2D array is passed then this is assumed to specify a metric with a dense positive definite matrix representation specified by the array. Otherwise if the value is a mici.matrices.PositiveDefiniteMatrix subclass it is assumed to directly specify the metric matrix representation.
dens_wrt_hausdorff (bool) – 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). Note that in the former case the base Hausdorff measure on the manifold depends on the metric defined on the ambient space, with the Hausdorff measure being defined with respect to the metric induced on the manifold from this ambient metric.
grad_neg_log_dens (Optional[GradientFunction]) – 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 a function to compute the derivative (and value) of neg_log_dens automatically.
jacob_constr (Optional[JacobianFunction]) – 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 (Optional[MatrixHessianProductFunction]) – 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.

constr(state)#

Constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of constr(state.pos) as 1D array.
Return type:: ArrayLike

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh2_flow_dmom(state, dt)#

Derivatives of h2_flow() flow map with respect to momentum argument.

Parameters:

state (ChainState) – State to evaluate derivatives of flow map at.
dt (ScalarLike) – Time interval flow simulated for.

Returns:

Return type:

tuple[Matrix, Matrix]

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_log_det_sqrt_gram(state)[source]#

Derivative of half log-determinant of Gram matrix with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of log_det_sqrt_gram(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

gram(state)#

Gram matrix at current position.

The Gram matrix as a position q is defined as

gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T

where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

h2_flow(state, dt)#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

inv_gram(state)#

Inverse of Gram matrix at current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Inverse of Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

jacob_constr(state)#

Jacobian of constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Jacobian of constr(state.pos) as 2D array.
Return type:: ArrayLike

jacob_constr_inner_product(jacob_constr_1, inner_product_matrix, jacob_constr_2=None)[source]#

Compute inner product of rows of constraint Jacobian matrices.

Computes jacob_constr_1 @ inner_product_matrix @ jacob_constr_2.T potentially exploiting any structure / sparsity in jacob_constr_1, jacob_constr_2 and inner_product_matrix.

Parameters:

jacob_constr_1 (MatrixLike) – First constraint Jacobian in product.
inner_product_matrix (PositiveDefiniteMatrix) – Positive-definite matrix defining inner-product between rows of two constraint Jacobians.
jacob_constr_2 (Optional[MatrixLike]) – Second constraint Jacobian in product. Defaults to jacob_constr_1 if set to None.

Returns:

Object corresponding to computed inner products of the constraint Jacobian rows.

Return type:

MatrixLike

log_det_sqrt_gram(state)#

Value of (half of) log-determinant of Gram matrix.

Parameters:: state (ChainState) –
Return type:: ScalarLike

mhp_constr(state)[source]#

Parameters:: state (ChainState) –
Return type:: MatrixHessianProduct

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

project_onto_cotangent_space(mom, state)#

Project a momentum on to the co-tangent space at a position.

Parameters:

mom (ArrayLike) – Momentum (co-)vector as 1D array to project on to co-tangent space.
state (ChainState) – State definining position on the manifold to project in to the co-tangent space of.

Returns:

Projected momentum in the co-tangent space at state.pos.

Return type:

ArrayLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.DenseRiemannianMetricSystem(neg_log_dens, metric_func, *, vjp_metric_func=None, grad_neg_log_dens=None)[source]#

Bases: RiemannianMetricSystem

Riemannian-metric system with dense matrix representation.

Hamiltonian system with a position dependent metric matrix representation which is specified to be a dense matrix function metric_func of the position q which is guaranteed to be positive definite almost-everywhere, with M = metric_func(q) then the metric matrix representation.

See documentation of RiemannianMetricSystem for more general details about Riemannian-metric Hamiltonian systems.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric_func (ArrayFunction) – Function which given a position array returns a 2D array corresponding to the positive definite metric matrix representation. The returned matrices (2D arrays) are assumed to be positive-definite for all input positions and a LinAlgError exception may be raised if this fails to be the case.
vjp_metric_func (Optional[VectorJacobianProductFunction]) –
Function which given a position array returns another function which takes a 2D array as an argument and returns the vector-Jacobian-product (VJP) of metric_func with respect to the position array argument. The VJP is here defined as a function of a 2D array v.
```
vjp(v) = sum(v[:, :, None] * jacob[:, :, :], axis=(0, 1))
```
where jacob is the (dim_pos, dim_pos, dim_pos) shaped Jacobian of M = metric_func(q) with respect to q i.e. the array of partial derivatives of the function such that
```
jacob[i, j, k] = ∂M[i, j] / ∂q[k]
```
Optionally the function may instead return a 2-tuple of values with the first a function to compute a VJP of metric_func and the second a 2D array containing the value of metric_func, both 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 VJP (and value) of metric_func automatically.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

metric(state)#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: VectorJacobianProduct

class mici.systems.DiagonalRiemannianMetricSystem(neg_log_dens, metric_diagonal_func, *, vjp_metric_diagonal_func=None, grad_neg_log_dens=None)[source]#

Bases: RiemannianMetricSystem

Riemannian-metric system with diagonal matrix representation.

Hamiltonian system with a position dependent diagonal metric matrix representation which is specified by a vector-valued function metric_diagonal_func of the position q which outputs a 1D array with strictly positive elements d = metric_diagonal_func(q) with the metric matrix representation then taken to be diag(d).

See documentation of RiemannianMetricSystem for more general details about Riemannian-metric Hamiltonian systems.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric_diagonal_func (ArrayFunction) – Function which given a position array returns a 1D array with strictly positive values corresponding to the diagonal values (left-to-right) of the diagonal metric matrix representation.
vjp_metric_diagonal_func (Optional[VectorJacobianProductFunction]) –
Function which given a position array returns another function which takes a 1D array as an argument and returns the vector-Jacobian-product (VJP) of metric_diagonal_func with respect to the position array argument. The VJP is here defined as a function of a 1D array v.
```
vjp(v) = sum(v[:, None] * jacob[:, :], axis=0)
```
where jacob is the (dim_pos, dim_pos) shaped Jacobian of d = metric_diagonal_func(q) with respect to q i.e. the array of partial derivatives of the function such that
```
jacob[i, j] = ∂d[i] / ∂q[j]
```
Optionally the function may instead return a 2-tuple of values with the first a function to compute a VJP of metric_diagonal_func and the second a 1D array containing the value of metric_diagonal_func, both 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 VJP (and value) of metric_diagonal_func automatically.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

metric(state)#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: VectorJacobianProduct

class mici.systems.EuclideanMetricSystem(neg_log_dens, *, metric=None, grad_neg_log_dens=None)[source]#

Bases: TractableFlowSystem

Hamiltonian system with a Euclidean metric on the position space.

Here Euclidean metric is defined to mean a metric with a fixed positive definite matrix representation \(M\). The momentum variables are taken to be independent of the position variables and with a zero-mean Gaussian marginal distribution with covariance specified by \(M\), so that the \(h_2\) Hamiltonian component is

\[h_2(q, p) = \frac{1}{2} p^T M^{-1} p\]

where \(q\) and \(p\) are the position and momentum variables respectively.

The \(h_1\) Hamiltonian component function is

\[h_1(q) = \ell(q)\]

where \(\ell(q)\) is the negative log (unnormalized) density of the target distribution with respect to the Lebesgue measure.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric (Optional[MetricLike]) – Matrix object corresponding to matrix representation of metric on position space and covariance of Gaussian marginal distribution on momentum vector. If None is passed (the default), the identity matrix will be used. If a 1D array is passed then this is assumed to specify a metric with positive diagonal matrix representation and the array the matrix diagonal. If a 2D array is passed then this is assumed to specify a metric with a dense positive definite matrix representation specified by the array. Otherwise if the value is a subclass of mici.matrices.PositiveDefiniteMatrix it is assumed to directly specify the metric matrix representation.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ScalarLike

dh2_dmom(state)[source]#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)[source]#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)[source]#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)[source]#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

h2_flow(state, dt)[source]#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)[source]#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.GaussianDenseConstrainedEuclideanMetricSystem(neg_log_dens, constr, *, metric=None, grad_neg_log_dens=None, jacob_constr=None, mhp_constr=None)[source]#

Bases: GaussianEuclideanMetricSystem, DenseConstrainedEuclideanMetricSystem

Gaussian Euclidean Hamiltonian system subject to a dense set of constraints.

See ConstrainedEuclideanMetricSystem for more details about constrained systems and GaussianEuclideanMetricSystem for Gaussian Euclidean metric systems.

Parameters:

neg_log_dens (ScalarFunction) –
Function which given a position array returns the negative logarithm of an unnormalized probability density on the unconstrained (ambient) position space with respect to the standard Gaussian measure. The density function is taken to specify a prior distribution on the ambient space with the target distribution then corresponding to the posterior distribution when conditioning on the (zero Lebesgue measure) event constr(pos) == 0. This target posterior distribution has support on the differentiable manifold implicitly defined by the constraint equation, with density with respect to the Hausdorff measure on the manifold corresponding to the ratio of the prior density (specified by neg_log_dens) and the square-root of the determinant of the Gram matrix defined by.
```
gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T
```
where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.
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.
metric (Optional[MetricLike]) – Matrix object corresponding to matrix representation of metric on unconstrained position space and covariance of Gaussian marginal distribution on unconstrained momentum vector. If None is passed (the default), the identity matrix will be used. If a 1D array is passed then this is assumed to specify a metric with positive diagonal matrix representation and the array the matrix diagonal. If a 2D array is passed then this is assumed to specify a metric with a dense positive definite matrix representation specified by the array. Otherwise if a subclass of mici.matrices.PositiveDefiniteMatrix it is assumed to directly specify the metric matrix representation.
grad_neg_log_dens (Optional[GradientFunction]) – 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 a function to compute the derivative (and value) of neg_log_dens automatically.
jacob_constr (Optional[JacobianFunction]) –
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 neg_log_dens automatically.
mhp_constr (Optional[MatrixHessianProductFunction]) –
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
```
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.

constr(state)#

Constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of constr(state.pos) as 1D array.
Return type:: ArrayLike

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh2_flow_dmom(state, dt)[source]#

Derivatives of h2_flow() flow map with respect to momentum argument.

Parameters:

state (ChainState) – State to evaluate derivatives of flow map at.
dt (ScalarLike) – Time interval flow simulated for.

Returns:

Return type:

tuple[Matrix, Matrix]

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_log_det_sqrt_gram(state)#

Derivative of half log-determinant of Gram matrix with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of log_det_sqrt_gram(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

gram(state)#

Gram matrix at current position.

The Gram matrix as a position q is defined as

gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T

where jacob_constr is the Jacobian of the constraint function constr and metric is the matrix representation of the metric on the ambient space.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

h2_flow(state, dt)#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

inv_gram(state)#

Inverse of Gram matrix at current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Inverse of Gram matrix as matrix object.
Return type:: PositiveDefiniteMatrix

jacob_constr(state)#

Jacobian of constraint function at the current position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Jacobian of constr(state.pos) as 2D array.
Return type:: ArrayLike

jacob_constr_inner_product(jacob_constr_1, inner_product_matrix, jacob_constr_2=None)[source]#

Compute inner product of rows of constraint Jacobian matrices.

Computes jacob_constr_1 @ inner_product_matrix @ jacob_constr_2.T potentially exploiting any structure / sparsity in jacob_constr_1, jacob_constr_2 and inner_product_matrix.

Parameters:

jacob_constr_1 (MatrixLike) – First constraint Jacobian in product.
inner_product_matrix (PositiveDefiniteMatrix) – Positive-definite matrix defining inner-product between rows of two constraint Jacobians.
jacob_constr_2 (Optional[MatrixLike]) – Second constraint Jacobian in product. Defaults to jacob_constr_1 if set to None.

Returns:

Object corresponding to computed inner products of the constraint Jacobian rows.

Return type:

MatrixLike

log_det_sqrt_gram(state)#

Value of (half of) log-determinant of Gram matrix.

Parameters:: state (ChainState) –
Return type:: ScalarLike

mhp_constr(state)#

Parameters:: state (ChainState) –
Return type:: MatrixHessianProduct

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

project_onto_cotangent_space(mom, state)#

Project a momentum on to the co-tangent space at a position.

Parameters:

mom (ArrayLike) – Momentum (co-)vector as 1D array to project on to co-tangent space.
state (ChainState) – State definining position on the manifold to project in to the co-tangent space of.

Returns:

Projected momentum in the co-tangent space at state.pos.

Return type:

ArrayLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.GaussianEuclideanMetricSystem(neg_log_dens, *, metric=None, grad_neg_log_dens=None)[source]#

Bases: EuclideanMetricSystem

Euclidean Hamiltonian system with a tractable Gaussian component.

Additionally the target distribution on the position variables is assumed to be defined by an unnormalized density with respect to the standard Gaussian measure on the position space (with identity covariance and zero mean), with the Hamiltonian component \(h_1\) corresponding to the negative logarithm of this density rather than the density with respect to the Lebesgue measure on the position space, i.e.

\[h_1(q) = \ell(q) - \frac{1}{2} q^T q\]

where \(q\) is the position and \(\ell(q)\) is the negative log (unnormalized) density of the target distribution with respect to the Lebesgue measure at \(q\). The Hamiltonian component function \(h_2\) is then assumed to have the form

\[h_2(q, p) = \frac{1}{2} q^T q + \frac{1}{2} p^T M^{-1} p\]

where \(p\) is the momentum. In this case the Hamiltonian flow due to the quadratic \(h_2\) component can be solved for analytically, allowing an integrator to be defined using this alternative splitting of the Hamiltonian [1].

References

Shahbaba, B., Lan, S., Johnson, W.O. and Neal, R.M., 2014. Split Hamiltonian Monte Carlo. Statistics and Computing, 24(3), pp.339-349.

Parameters:

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 the standard Gaussian measure on the position space, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric (Optional[MetricLike]) – Matrix object corresponding to matrix representation of metric on position space and covariance of Gaussian marginal distribution on momentum vector. If None is passed (the default), the identity matrix will be used. If a 1D array is passed then this is assumed to specify a metric with positive diagonal matrix representation and the array the matrix diagonal. If a 2D array is passed then this is assumed to specify a metric with a dense positive definite matrix representation specified by the array. Otherwise if the value is a subclass of mici.matrices.PositiveDefiniteMatrix it is assumed to directly specify the metric matrix representation.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ScalarLike

dh2_dmom(state)[source]#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)[source]#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)[source]#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

h2_flow(state, dt)[source]#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.RiemannianMetricSystem(neg_log_dens, metric_matrix_class, metric_func, *, vjp_metric_func=None, grad_neg_log_dens=None, metric_kwargs=None)[source]#

Bases: System

Riemannian Hamiltonian system with a position-dependent metric.

This class allows for metric matrix representations of any generic type. In most cases a specialized subclass such as DenseRiemannianMetricSystem, CholeskyFactoredRiemannianMetricSystem, DiagonalRiemannianMetricSystem, ScalarRiemannianMetricSystem or SoftAbsRiemannianMetricSystem will provide a simpler method of constructng a system with a metric matrix representation of a specific type.

The position space is assumed to be a Riemannian manifold with a metric with position-dependent positive definite matrix-representation \(M(q)\) where \(q\) is a position vector. The momentum \(p\) is then taken to have a zero-mean Gaussian conditional distribution given the position \(q\), with covariance \(M(q)\), i.e. \(p \sim \mathcal{N}(0, M(q))\) [1].

The \(h_1\) Hamiltonian component is then

\[h_1(q) = \ell(q) + \frac{1}{2}\log\left|M(q)\right|\]

where \(\ell(q)\) is the negative log (unnormalized) density of the target distribution with respect to the Lebesgue measure at \(q\). The \(h_2\) Hamiltonian component is

\[h_2(q, p) = \frac{1}{2} p^T (M(q))^{-1} p.\]

Due to the coupling between the position and momentum variables in \(h_2\), the Hamiltonian system is non-separable, requiring use of a numerical integrator with implicit steps when simulating the Hamiltonian dynamic associated with the system, e.g. mici.integrators.ImplicitLeapfrogIntegrator.

References

Girolami, M. and Calderhead, B., 2011. Riemann manifold Langevin and Hamiltonian Monte Varlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2), pp.123-214.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric_matrix_class (type[PositiveDefiniteMatrix]) –
Class (or factory function returning an instance of the class) which defines type of matrix representation of metric. The class initializer should take a single positional argument which will be passed the array outputted by metric_func, and which is assumed to be a parameter which fully defines the resulting matrix (e.g. the diagonal of a mici.matrices.DiagonalMatrix). The class initializer may also optionally take one or more keyword arguments, with the metric_kwargs argument used to specify the value of these, if any. Together this means the metric matrix representation at a position pos is constructed as.
```
metric = metric_matrix_class(metric_func(pos), **metric_kwargs)
```
The mici.matrices.PositiveDefiniteMatrix subclass should as a minimum define inv, log_abs_det, grad_log_abs_det, grad_quadratic_form_inv, __matmul__ and __rmatmul__ methods / properties (see documentation of mici.matrices.PositiveDefiniteMatrix and mici.matrices.DifferentiableMatrix for definitions of the expected behaviour of these methods).
metric_func (ArrayFunction) – Function which given a position array returns an array containing the parameter value of the metric matrix representation passed as the single positional argument to the metric_matrix_class initializer.
vjp_metric_func (Optional[VectorJacobianProductFunction]) –
Function which given a position array returns another function which takes an array as an argument and returns the vector-Jacobian-product (VJP) of metric_func with respect to the position array argument. The VJP is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to
```
vjp(v) = sum(v[..., None] * jacob, tuple(range(v.ndim))
```
where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that
```
jacob[..., i] = ∂m[...] / ∂q[i]
```
Optionally the function may instead return a 2-tuple of values with the first a function to compute a VJP of metric_func and the second an array containing the value of metric_func, both 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 VJP (and value) of metric_func automatically.
grad_neg_log_dens (Optional[GradientFunction]) – 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.
metric_kwargs (Optional[dict[str, Any]]) – An optional dictionary of any additional keyword arguments to the initializer of metric_matrix_class.

dh1_dpos(state)[source]#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)[source]#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)[source]#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)[source]#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)[source]#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)[source]#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

metric(state)[source]#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)[source]#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)[source]#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)[source]#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: VectorJacobianProduct

class mici.systems.ScalarRiemannianMetricSystem(neg_log_dens, metric_scalar_func, *, vjp_metric_scalar_func=None, grad_neg_log_dens=None)[source]#

Bases: RiemannianMetricSystem

Riemannian-metric system with scaled identity matrix representation.

Hamiltonian system with a position dependent scaled identity metric matrix representation which is specified by a scalar function metric_scalar_function of the position q which outputs a strictly positive scalar s = metric_scalar_func(q) with the the metric matrix representation then taken to be s * identity(q.shape[0]).

See documentation of RiemannianMetricSystem for more general details about Riemannian-metric Hamiltonian systems.

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
metric_scalar_func (ScalarFunction) – Function which given a position array returns a strictly positive scalar corresponding to the parameter value of the scaled identity metric matrix representation.
vjp_metric_scalar_func (Optional[VectorJacobianProductFunction]) –
Function which given a position array returns another function which takes a scalar as an argument and returns the vector-Jacobian-product (VJP) of metric_scalar_func with respect to the position array argument. The VJP is here defined as a function of a scalar v.
```
vjp(v) = v @ grad
```
where grad is the (dim_pos,) shaped Jacobian (gradient) of s = metric_scalar_func(q) with respect to q i.e. the array of partial derivatives of the function such that
```
grad[i] = ∂s / ∂q[i]
```
Optionally the function may instead return a 2-tuple of values with the first a function to compute a VJP of metric_scalar_func and the second a float containing the value of metric_scalar_func, both 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 VJP (and value) of metric_scalar_func automatically.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

metric(state)[source]#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: VectorJacobianProduct

class mici.systems.SoftAbsRiemannianMetricSystem(neg_log_dens, *, grad_neg_log_dens=None, hess_neg_log_dens=None, mtp_neg_log_dens=None, softabs_coeff=1.0)[source]#

Bases: RiemannianMetricSystem

SoftAbs Riemmanian metric Hamiltonian system.

Hamiltonian system with a position dependent metric matrix representation which is specified to be an eigenvalue-regularized transformation of the Hessian of the negative log density function (the symmetric matrix of second derivatives the negative log density function with respect to the position array components. Specifically if hess_neg_log_dens is a symmetric 2D square array valued function of the position q, with H = hess_neg_log_dens(q) then if an eigenvalue decomposition of H is computed, i.e. eigval, eigvec = eigh(H), with eigval a 1D array of real eigenvalues, and eigvec the corresponding 2D array (orthogonal matrix) with eigenvectors as columns, then the resulting positive-definite metric matrix representation M is computed as

M = eigvec @ diag(softabs(eigval, softabs_coeff)) @ eigvec.T

with softabs(x, softabs_coeff) = x / tanh(x * softabs_coeff) an elementwise function which acts as a smooth approximation to the absolute function (ensuring all the eigenvalues of M are strictly positive) with the additional scalar parameter softabs_coeff controlling the smoothness of the approximation, with softabs tending to the piecewise linear abs function as softabs_coeff tends to infinity, and becoming increasingly smooth as softabs_coeff tends to zero.

See documentation of RiemannianMetricSystem for more general details about Riemannian-metric Hamiltonian systems.

References

Betancourt, M., 2013. A general metric for Riemannian manifold Hamiltonian Monte Carlo. In Geometric science of information (pp. 327-334).

Parameters:

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 the Lebesgue measure, with the corresponding distribution on the position space being the target distribution it is wished to draw approximate samples from.
grad_neg_log_dens (Optional[GradientFunction]) – 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.
hess_neg_log_dens (Optional[HessianFunction]) – Function which given a position array returns the Hessian of neg_log_dens with respect to the position array argument as a 2D array. Optionally the function may instead return a 3-tuple of values with the first a 2D array containting the Hessian of neg_log_dens, the second a 1D array containing the gradient of neg_log_dens and the third the value of neg_log_dens, 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 Hessian (and gradient and value) of neg_log_dens automatically.
mtp_neg_log_dens (Optional[MatrixTressianProductFunction]) –
Function which given a position array returns another function which takes a 2D array (matrix) as an argument and returns the matrix-Tressian-product (MTP) of neg_log_dens with respect to the position array argument. The MTP is here defined as a function of a matrix m corresponding to.
```
mtp(m) = sum(m[:, :, None] * tress[:, :, :], axis=(0, 1))
```
where tress is the ‘Tressian’ of f = neg_log_dens(q) wrt q i.e. the 3D array of third-order partial derivatives of the scalar-valued function such that
```
tress[i, j, k] = ∂³f / (∂q[i] ∂q[j] ∂q[k])
```
Optionally the function may instead return a 4-tuple of values with the first a function to compute a MTP of neg_log_dens, the second a 2D array containing the Hessian of neg_log_dens, the third a 1D array containing the gradient of neg_log_dens and the fourth the value of neg_log_dens, 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 MTP (and Hesisan and gradient and value) of neg_log_dens automatically.
softabs_coeff (ScalarLike) – Positive regularisation coefficient for smooth approximation to absolute value used to regularize Hessian eigenvalues in metric matrix representation. As the value tends to infinity the approximation becomes increasingly close to the absolute function.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ArrayLike

dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

hess_neg_log_dens(state)[source]#

Hessian of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: 2D array of neg_log_dens(state) second derivatives with respect to state.pos, with hessian[i, j] the second derivative of neg_log_dens(state) with respect to state.pos[i] and state.pos[j].
Return type:: ArrayLike

metric(state)#

Function computing the metric matrix representation.

The returned type of this function is that specified by the metric_matrix_class argument to the initializer.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Metric matrix representation.
Return type:: PositiveDefiniteMatrix

metric_func(state)[source]#

Function computing the parameter of the metric matrix representation.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of metric_func(state.pos).
Return type:: ArrayLike

mtp_neg_log_dens(state)[source]#

Generate MTP of negative log density with respect to position.

The matrix-Tressian-product (MTP) is here defined as a function of a matrix m corresponding to

mtp(m) = sum(m[:, :, None] * tress[:, :, :], axis=(0, 1))

where tress is the ‘Tressian’ of f = neg_log_dens(q) with respect to q = state.pos i.e. the 3D array of third-order partial derivatives of the scalar-valued function such that

tress[i, j, k] = ∂³f / (∂q[i] ∂q[j] ∂q[k])

Parameters:: state (ChainState) – State to compute value at.
Returns:: Function which accepts a 2D array of shape (state.pos.shape[0], state.pos.shape[0]) as an argument and returns an array of shape state.pos.shape containing the computed MTP value.
Return type:: MatrixTressianProduct

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

vjp_metric_func(state)[source]#

Function constructing a vector-Jacobian-product for metric_func.

The vector-Jacobian-product is here defined as a function of an array v (of the same shape as the output of metric_func) corresponding to

vjp(v) = sum(v[..., None] * jacob, axis=tuple(range(v.ndim))

where jacob is the Jacobian of m = metric_func(q) wrt q i.e. the array of partial derivatives of the function such that

jacob[..., i] = ∂m[...] / ∂q[i]

Parameters:: state (ChainState) – State to compute VJP at.
Returns:: Vector-Jacobian-product function.
Return type:: MatrixTressianProduct

class mici.systems.System(neg_log_dens, *, grad_neg_log_dens=None)[source]#

Bases: ABC

Base class for Hamiltonian systems.

The Hamiltonian function \(h\) is assumed to have the general form

\[h(q, p) = h_1(q) + h_2(q, p)\]

where \(q\) and \(p\) are the position and momentum variables respectively, and \(h_1\) and \(h_2\) Hamiltonian component functions. The exact Hamiltonian flow for the \(h_1\) component can be always be computed as it depends only on the position variable however depending on the form of \(h_2\) the corresponding exact Hamiltonian flow may or may not be simulable.

By default \(h_1\) is assumed to correspond to the negative logarithm of an unnormalized density on the position variables 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.

Parameters:

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.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)[source]#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ScalarLike

abstract dh2_dmom(state)[source]#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

abstract dh2_dpos(state)[source]#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)[source]#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)[source]#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)[source]#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)[source]#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)[source]#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)[source]#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

abstract h2(state)[source]#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

neg_log_dens(state)[source]#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

abstract sample_momentum(state, rng)[source]#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike

class mici.systems.TractableFlowSystem(neg_log_dens, *, grad_neg_log_dens=None)[source]#

Bases: System

Base class for Hamiltonian systems with tractable component flows.

The Hamiltonian function \(h\) is assumed to have the general form

\[h(q, p) = h_1(q) + h_2(q, p)\]

where \(q\) and \(p\) are the position and momentum variables respectively, and \(h_1\) and \(h_2\) Hamiltonian component functions. The exact Hamiltonian flows for both the \(h_1\) and \(h_2\) components are assumed to be tractable for subclasses of this class.

Parameters:

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.
grad_neg_log_dens (Optional[GradientFunction]) – 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.

dh1_dpos(state)#

Derivative of h1 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed h1 derivative.
Return type:: ScalarLike

abstract dh2_dmom(state)#

Derivative of h2 Hamiltonian component with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.mom.
Return type:: ArrayLike

abstract dh2_dpos(state)#

Derivative of h2 Hamiltonian component with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2(state) derivative with respect to state.pos.
Return type:: ArrayLike

dh_dmom(state)#

Derivative of Hamiltonian with respect to momentum.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.mom.
Return type:: ArrayLike

dh_dpos(state)#

Derivative of Hamiltonian with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h(state) derivative with respect to state.pos.
Return type:: ArrayLike

grad_neg_log_dens(state)#

Derivative of negative log density with respect to position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of neg_log_dens(state) derivative with respect to state.pos.
Return type:: ArrayLike

h(state)#

Hamiltonian function for system.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of Hamiltonian.
Return type:: ScalarLike

h1(state)#

Hamiltonian component depending only on position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h1 Hamiltonian component.
Return type:: ScalarLike

h1_flow(state, dt)#

Apply exact flow map corresponding to h1 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

abstract h2(state)#

Hamiltonian component depending on momentum and optionally position.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of h2 Hamiltonian component.
Return type:: ScalarLike

abstract h2_flow(state, dt)[source]#

Apply exact flow map corresponding to h2 Hamiltonian component.

state argument is modified in place.

Parameters:

state (ChainState) – State to start flow at.
dt (ScalarLike) – Time interval to simulate flow for.

neg_log_dens(state)#

Negative logarithm of unnormalized density of target distribution.

Parameters:: state (ChainState) – State to compute value at.
Returns:: Value of computed negative log density.
Return type:: ScalarLike

abstract sample_momentum(state, rng)#

Sample a momentum from its conditional distribution given a position.

Parameters:

state (ChainState) – State defining position to condition on.
rng (Generator) –

Returns:

Sampled momentum.

Return type:

ArrayLike