diff --git a/src/mici/adapters.py b/src/mici/adapters.py index 02c8b71..23fb047 100644 --- a/src/mici/adapters.py +++ b/src/mici/adapters.py @@ -113,7 +113,7 @@ def is_fast(self) -> bool: An adapter which requires only local information to adapt the transition parameters should be classified as fast while one which requires more global information and so more chain iterations should be classified as slow i.e. - `is_fast == False`. + :code:`is_fast == False`. """ diff --git a/src/mici/integrators.py b/src/mici/integrators.py index bb18269..abffa4f 100644 --- a/src/mici/integrators.py +++ b/src/mici/integrators.py @@ -18,27 +18,28 @@ if TYPE_CHECKING: from typing import Any, Callable, Optional, Sequence from mici.states import ChainState - from mici.systems import System + from mici.systems import System, TractableFlowSystem from mici.types import NormFunction class Integrator(ABC): - """Base class for integrators for simulating Hamiltonian dynamics. + r"""Base class for integrators for simulating Hamiltonian dynamics. - For a Hamiltonian function `h` with position variables `q` and momentum variables - `p`, the canonical Hamiltonian dynamic is defined by the ordinary differential - equation system + For a Hamiltonian function :math:`h` with position variables :math:`q` and momentum + variables :math:`p`, the canonical Hamiltonian dynamic is defined by the ordinary + differential equation system - q̇ = ∇₂h(q, p) - ṗ = -∇₁h(q, p) + .. math:: + + \dot{q} = \nabla_2 h(q, p), \qquad \dot{p} = -\nabla_1 h(q, p) with the flow map corresponding to the solution of the corresponding initial value problem a time-reversible and symplectic (and by consequence volume-preserving) map. - Derived classes implement a `step` method which approximates the flow-map over some - small time interval, while conserving the properties of being time-reversible and - symplectic, with composition of this integrator step method allowing simulation of - time-discretised trajectories of the Hamiltonian dynamics. + Derived classes implement a :py:meth:`step` method which approximates the flow-map + over some small time interval, while conserving the properties of being + time-reversible and symplectic, with composition of this integrator step method + allowing simulation of time-discretised trajectories of the Hamiltonian dynamics. """ def __init__(self, system: System, step_size: Optional[float] = None): @@ -85,23 +86,26 @@ class TractableFlowIntegrator(Integrator): The Hamiltonian function is assumed to be expressible as the sum of two analytically tractable components for which the corresponding Hamiltonian flows can be exactly - simulated. Specifically it is assumed that the Hamiltonian function `h` takes the - form + simulated. Specifically it is assumed that the Hamiltonian function :math:`h` takes + the form - h(q, p) = h₁(q) + h₂(q, p) + .. math:: - where `q` and `p` are the position and momentum variables respectively, and `h₁` and - `h₂` are Hamiltonian component functions for which the exact flows can be computed. + h(q, p) = h_1(q) + h_2(q, p) + + where :math:`q` and :math:`p` are the position and momentum variables respectively, + and :math:`h_1` and :math:`h_2` are Hamiltonian component functions for which the + exact flows can be computed. """ - def __init__(self, system: System, step_size: Optional[float] = None): + def __init__(self, system: TractableFlowSystem, step_size: Optional[float] = None): """ Args: - system: Hamiltonian system to integrate the dynamics of. Must define both - `h1_flow` and `h2_flow` methods. - step_size: Integrator time step. If set to `None` it is assumed that a step - size adapter will be used to set the step size before calling the `step` - method. + system: Hamiltonian system to integrate the dynamics of with tractable + Hamiltonian component flows. + step_size: Integrator time step. If set to :code:`None` it is assumed that a + step size adapter will be used to set the step size before calling the + :py:meth:`step` method. """ if not hasattr(system, "h1_flow") or not hasattr(system, "h2_flow"): raise ValueError( @@ -117,25 +121,29 @@ def __init__(self, system: System, step_size: Optional[float] = None): class LeapfrogIntegrator(TractableFlowIntegrator): """Leapfrog integrator for Hamiltonian systems with tractable component flows. - For separable Hamiltonians of the form + For separable Hamiltonians of the - h(q, p) = h₁(q) + h₂(p) + .. math:: - where `h₁` is the potential energy and `h₂` is the kinetic energy, this integrator - corresponds to the classic (position) Störmer-Verlet method. + h(q, p) = h_1(q) + h_2(p) + + where :math:`h_1` is the potential energy and :math:`h_2` is the kinetic energy, + this integrator corresponds to the classic (position) Störmer-Verlet method. The integrator can also be applied to the more general Hamiltonian splitting - h(q, p) = h₁(q) + h₂(q, p) + .. math:: + + h(q, p) = h_1(q) + h_2(q, p) - providing the flows for `h₁` and `h₂` are both tractable. + providing the flows for :math:`h_1` and :math:`h_2` are both tractable. For more details see Sections 2.6 and 4.2.2 in Leimkuhler and Reich (2004). References: - Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). - Cambridge University Press. + 1. Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). + Cambridge University Press. """ def _step(self, state: ChainState, time_step: float): @@ -145,39 +153,61 @@ def _step(self, state: ChainState, time_step: float): class SymmetricCompositionIntegrator(TractableFlowIntegrator): - """Symmetric composition integrator for Hamiltonians with tractable component flows. + r"""Symmetric composition integrator for Hamiltonians with tractable component flows The Hamiltonian function is assumed to be expressible as the sum of two analytically tractable components for which the corresponding Hamiltonian flows can be exactly - simulated. Specifically it is assumed that the Hamiltonian function h takes the form + simulated. Specifically it is assumed that the Hamiltonian function :math:`h` takes + the form - h(q, p) = h₁(q) + h₂(q, p) + .. math:: - where `q` and `p` are the position and momentum variables respectively, and `h₁` and - `h₂` are Hamiltonian component functions for which the exact flows, respectively - `Φ₁` and `Φ₂`, can be computed. An alternating composition can then be formed as + h(q, p) = h_1(q) + h_2(q, p) + + where :math:`q` and :math:`p` are the position and momentum variables respectively, + and :math:`h_1` and :math:`h_2` are Hamiltonian component functions for which the + exact flows, respectively :math:`\Phi_1` and :math:`\Phi_2`, can be computed. An + alternating composition can then be formed as + + .. math:: - Ψ(t) = A(aₛt) ∘ B(bₛt) ∘ ⋯ ∘ A(a₁t) ∘ B(b₁t) ∘ A(a₀t) + \Psi(t) = + A(a_S t) \circ B(b_S t) \circ \dots \circ A(a_1 t) \circ B(b_1 t) \circ A(a_0 t) - where `A = Φ₁` and `B = Φ₂` or `A = Φ₂` and `B = Φ₁`, and `(a₀, ⋯, aₛ)` and - `(b₁, ⋯, bₛ)` are a set of coefficients to be determined and `s` is the number of - stages in the composition. + where :math:`A = \Phi_1` and :math:`B = \Phi_2` or :math:`A = \Phi_2` and :math:`B = + \Phi_1`, and :math:`(a_0,\dots, a_S)` and :math:`(b_1, \dots, b_S)` are a set of + coefficients to be determined and :math:`S` is the number of stages in the + composition. To ensure a consistency (i.e. the integrator is at least order one) we require that - sum(a₀, ⋯, aₛ) = sum(b₁, ⋯, bₛ) = 1 + .. math:: + + \sum_{s=0}^S a_s = \sum_{s=1}^S b_s = 1. For symmetric compositions we restrict that - aₛ₋ₘ = aₘ and bₛ₊₁₋ₘ = bₛ + .. math:: + + a_{S-m} = a_m, \quad b_{S+1-m} = b_s with symmetric consistent methods of at least order two. - The combination of the symmetry and consistency requirements mean that a `s`-stage - symmetric composition method can be described by `(s - 1)` 'free' coefficients + The combination of the symmetry and consistency requirements mean that a + :math:`S`-stage symmetric composition method can be described by :math:`S - 1` + 'free' coefficients + + .. math:: + + (a_0, b_1, a_1, \dots, a_K, b_K) + + with :math:`K = (S - 1) / 2` if :math:`S` is odd or + + .. math:: - (a₀, b₁, a₁, ⋯, aₖ, bₖ) with k = (s - 1) / 2 if s is odd or - (a₀, b₁, a₁, ⋯, aₖ) with k = (s - 2) / 2 if s is even + (a_0, b_1, a_1, \dots, a_K) + + with :math:`K = (S - 2) / 2` if :math:`S` is even. The Störmer-Verlet 'leapfrog' integrator is a special case corresponding to the unique (symmetric and consistent) 1-stage integrator. @@ -186,30 +216,33 @@ class SymmetricCompositionIntegrator(TractableFlowIntegrator): References: - Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). - Cambridge University Press. + 1. Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). + Cambridge University Press. """ def __init__( self, - system: System, + system: TractableFlowSystem, free_coefficients: Sequence[float], step_size: Optional[float] = None, initial_h1_flow_step: bool = True, ): - """ + r""" Args: - system: Hamiltonian system to integrate the dynamics of. - free_coefficients: Sequence of `s - 1` scalar values, where `s` is the - number of stages in the symmetric composition, specifying the free - coefficients `(a₀, b₁, a₁, ⋯, aₖ, bₖ)` with `k = (s - 1) / 2` if `s` is - odd or `(a₀, b₁, a₁, ⋯, aₖ)` with `k = (s - 2) / 2` if `s` is even. - step_size: Integrator time step. If set to `None` it is assumed that a step - size adapter will be used to set the step size before calling the `step` - method. - initial_h1_flow_step: Whether the initial 'A' flow in the composition should - correspond to the flow of the `h₁` Hamiltonian component (`True`) or to - the flow of the `h₂` component (`False`). + system: Hamiltonian system to integrate the dynamics of with tractable + Hamiltonian component flows. + free_coefficients: Sequence of :math:`S - 1` scalar values, where :math:`S` + is the number of stages in the symmetric composition, specifying the + free coefficients :math:`(a_0, b_1, a_1, \dots, a_K, b_K)` with :math:`K + = (S - 1) / 2` if :math:`S` is odd or :math:`(a_0, b_1, a_1, \dots, + a_K)` with :math:`k = (S - 2) / 2` if `S` is even. + step_size: Integrator time step. If set to :code:`None` it is assumed that a + step size adapter will be used to set the step size before calling the + :py:meth:`step` method. + initial_h1_flow_step: Whether the initial :math:`A` flow in the composition + should correspond to the flow of the `h_1` Hamiltonian component + (:code:`True`) or to the flow of the :math:`h_2` component + (:code:`False`). """ super().__init__(system, step_size) self.initial_h1_flow_step = initial_h1_flow_step @@ -238,18 +271,19 @@ class BCSSTwoStageIntegrator(SymmetricCompositionIntegrator): References: - Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). - Numerical integrators for the Hybrid Monte Carlo method. - SIAM Journal on Scientific Computing, 36(4), A1556-A1580. + 1. Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). + Numerical integrators for the Hybrid Monte Carlo method. + SIAM Journal on Scientific Computing, 36(4), A1556-A1580. """ - def __init__(self, system: System, step_size: Optional[float] = None): + def __init__(self, system: TractableFlowSystem, step_size: Optional[float] = None): """ Args: - system: Hamiltonian system to integrate the dynamics of. - step_size: Integrator time step. If set to `None` it is assumed that a step - size adapter will be used to set the step size before calling the `step` - method. + system: Hamiltonian system to integrate the dynamics of with tractable + Hamiltonian component flows. + step_size: Integrator time step. If set to :code:`None` it is assumed that a + step size adapter will be used to set the step size before calling the + :py:meth:`step` method. """ a_0 = (3 - 3 ** 0.5) / 6 super().__init__(system, (a_0,), step_size, True) @@ -262,18 +296,19 @@ class BCSSThreeStageIntegrator(SymmetricCompositionIntegrator): References: - Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). - Numerical integrators for the Hybrid Monte Carlo method. - SIAM Journal on Scientific Computing, 36(4), A1556-A1580. + 1. Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). + Numerical integrators for the Hybrid Monte Carlo method. + SIAM Journal on Scientific Computing, 36(4), A1556-A1580. """ - def __init__(self, system: System, step_size: Optional[float] = None): + def __init__(self, system: TractableFlowSystem, step_size: Optional[float] = None): """ Args: - system: Hamiltonian system to integrate the dynamics of. - step_size: Integrator time step. If set to `None` it is assumed that a step - size adapter will be used to set the step size before calling the `step` - method. + system: Hamiltonian system to integrate the dynamics of with tractable + Hamiltonian component flows. + step_size: Integrator time step. If set to :code:`None` it is assumed that a + step size adapter will be used to set the step size before calling the + :py:meth:`step` method. """ a_0 = 0.11888010966548 b_1 = 0.29619504261126 @@ -287,18 +322,19 @@ class BCSSFourStageIntegrator(SymmetricCompositionIntegrator): References: - Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). - Numerical integrators for the Hybrid Monte Carlo method. - SIAM Journal on Scientific Computing, 36(4), A1556-A1580. + 1. Blanes, S., Casas, F., & Sanz-Serna, J. M. (2014). + Numerical integrators for the Hybrid Monte Carlo method. + SIAM Journal on Scientific Computing, 36(4), A1556-A1580. """ - def __init__(self, system: System, step_size: Optional[float] = None): + def __init__(self, system: TractableFlowSystem, step_size: Optional[float] = None): """ Args: - system: Hamiltonian system to integrate the dynamics of. - step_size: Integrator time step. If set to `None` it is assumed that a step - size adapter will be used to set the step size before calling the `step` - method. + system: Hamiltonian system to integrate the dynamics of with tractable + Hamiltonian component flows. + step_size: Integrator time step. If set to :code:`None` it is assumed that + a step size adapter will be used to set the step size before calling the + :py:meth:`step` method. """ a_0 = 0.071353913450279725904 b_1 = 0.191667800000000000000 @@ -311,18 +347,20 @@ class ImplicitLeapfrogIntegrator(Integrator): Also known as the generalised leapfrog method. - The Hamiltonian function `h` is assumed to take the form + The Hamiltonian function :math:`h` is assumed to take the form - h(q, p) = h₁(q) + h₂(q, p) + .. math:: - where `q` and `p` are the position and momentum variables respectively, `h₁` is a - Hamiltonian component function for which the exact flow can be computed and `h₂` is - a Hamiltonian component function of the position and momentum variables, which may - be non-separable and for which exact simulation of the correspond Hamiltonian flow - may not be possible. + h(q, p) = h_1(q) + h_2(q, p) + + where :math:`q` and :math:`p` are the position and momentum variables respectively, + :math:`h_1` is a Hamiltonian component function for which the exact flow can be + computed and :math:`h_2` is a Hamiltonian component function of the position and + momentum variables, which may be non-separable and for which exact simulation of the + correspond Hamiltonian flow may not be possible. A pair of implicit component updates are used to approximate the flow due to the - `h₂` Hamiltonian component, with a fixed-point iteration used to solve the + :math:`h_2` Hamiltonian component, with a fixed-point iteration used to solve the non-linear system of equations. The resulting implicit integrator is a symmetric second-order method corresponding @@ -331,8 +369,8 @@ class ImplicitLeapfrogIntegrator(Integrator): References: - Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). - Cambridge University Press. + 1. Leimkuhler, B., & Reich, S. (2004). Simulating Hamiltonian Dynamics (No. 14). + Cambridge University Press. """ def __init__( diff --git a/src/mici/samplers.py b/src/mici/samplers.py index 555f650..a15145a 100644 --- a/src/mici/samplers.py +++ b/src/mici/samplers.py @@ -805,8 +805,18 @@ def __init__(self, rng: Generator, transitions: dict[str, Transition]): DeprecationWarning, ) rng = np.random.Generator(rng._bit_generator) - self.rng = rng - self.transitions = transitions + self._rng = rng + self._transitions = transitions + + @property + def transitions(self) -> dict[str, Transition]: + """Dictionary of transition kernels sampled from in each chain iteration.""" + return self._transitions + + @property + def rng(self) -> Generator: + """NumPy random number generator object.""" + return self._rng def sample_chains( self, @@ -1092,21 +1102,22 @@ def __init__( ): """ Args: - system Hamiltonian system to be simulated. + system Hamiltonian system to be simulated, corresponding to joint + distribution on augmented space. rng: Numpy random number generator. - integration_transition: Markov transition kernel which leaves canonical + integration_transition: Markov transition kernel which leaves joint distribution invariant and jointly updates the position and momentum components of the chain state by integrating the Hamiltonian dynamics of the system to propose new values for the state. momentum_transition: Markov transition kernel which leaves the conditional - distribution on the momentum under the canonical distribution invariant, + distribution on the momentum under the join distribution invariant, updating only the momentum component of the chain state. If set to - `None` the momentum transition operator - `mici.transitions.IndependentMomentumTransition` will be used, which - independently samples the momentum from its conditional distribution. + :py:const:`None` the momentum transition operator + :py:class:`mici.transitions.IndependentMomentumTransition` will be used, + which independently samples the momentum from its conditional + distribution. """ - self.system = system - self.rng = rng + self._system = system if momentum_transition is None: momentum_transition = IndependentMomentumTransition(system) super().__init__( @@ -1117,6 +1128,11 @@ def __init__( }, ) + @property + def system(self) -> System: + """Hamiltonian system corresponding to joint distribution on augmented space.""" + return self._system + def _preprocess_init_state( self, init_state: Union[NDArray, ChainState, dict] ) -> ChainState: diff --git a/src/mici/stagers.py b/src/mici/stagers.py index 5d8b576..8ee295f 100644 --- a/src/mici/stagers.py +++ b/src/mici/stagers.py @@ -1,4 +1,4 @@ -"""Classes for controlling sampling of Markov chains split into stages.""" +"""Classes for staging sampling of Markov chains.""" from __future__ import annotations @@ -13,12 +13,12 @@ class ChainStage(NamedTuple): """Parameters of chain stage. - + Parameters: n_iter: Number of iterations in chain stage. adapters: Dictionary of adapters to apply to each transition in chain stage. trace_funcs: Functions defining chain variables to trace during chain stage. - record_stats: Whether to record statistics and traces during chain stage. + record_stats: Whether to record statistics during chain stage. """ n_iter: int @@ -41,10 +41,10 @@ def stages( ) -> dict[str, ChainStage]: """Create dictionary specifying labels and parameters of chain sampling stages. - Constructs a (ordered) dictionary with entries corresponding to the sequence of + Constructs an ordered dictionary with entries corresponding to the sequence of sampling stages when running chains with one or more initial adaptation stages. The keys of each entry are a string label for the sampling stage and the value a - `ChainStage` named tuple specifying the parameters of the stage. + :py:class:`ChainStage` named tuple specifying the parameters of the stage. Args: n_warm_up_iter: Number of adaptive warm up iterations per chain. Depending @@ -62,16 +62,19 @@ def stages( stored. By default chains are only traced in the iterations of the final non-adaptive stage - this behaviour can be altered using the `trace_warm_up` argument. - adapters: Dictionary of iterables of `Adapter` instances keyed by strings - corresponding to the key of the transition in the sampler `transitions` + adapters: Dictionary of iterables of :py:class:`mici.adapters.Adapter` + instances keyed by strings corresponding to the key of the transition in + the sampler + :py:attr:`mici.samplers.MarkovChainMonteCarloMethod.transitions` dictionary to apply the adapters to, to use to adaptatively set parameters of the transitions during the adaptive stages of the chains. Note that the adapter updates are applied in the order the adapters appear in the iterables and so if multiple adapters change the same parameter(s) the order will matter. trace_warm_up: Whether to record chain traces and statistics during warm-up - stage iterations (`True`) or only record traces and statistics in the - iterations of the final non-adaptive stage (`False`, the default). + stage iterations (:code:`True`) or only record traces and statistics in + the iterations of the final non-adaptive stage (:code:`False`, the + default). Returns: Dictionary specifying chain stages, keyed by labels for stages. @@ -130,7 +133,8 @@ class WindowedWarmUpStager(Stager): information. The adapters to be used in both the fast and slow adaptation stages will be referred to as the *fast adapters* and the adapters to use in only the slow adaptation stages the *slow adapters*. Each adapter self identifies if it is a fast - adapter by whether the `is_fast` attribute is set to `True`. + adapter by whether the :py:attr:`mici.adapters.Adapter.is_fast` attribute is set to + :code:`True`. The adaptive warm up iterations are split into three stages: diff --git a/src/mici/systems.py b/src/mici/systems.py index 2a51635..0aeceb2 100644 --- a/src/mici/systems.py +++ b/src/mici/systems.py @@ -48,7 +48,7 @@ class System(ABC): the corresponding exact Hamiltonian flow may or may not be simulable. By default :math:`h_1` is assumed to correspond to the negative logarithm of an - unnormalized density on the position variables with respect to the Lebesgue measure, + 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. """ @@ -62,7 +62,7 @@ def __init__( Args: neg_log_dens: 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 + 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: Function which given a position array returns the @@ -216,7 +216,55 @@ def sample_momentum(self, state: ChainState, rng: Generator) -> ArrayLike: """ -class EuclideanMetricSystem(System): +class TractableFlowSystem(System): + r"""Base class for Hamiltonian systems with tractable component flows. + + The Hamiltonian function :math:`h` is assumed to have the general form + + .. math:: + + h(q, p) = h_1(q) + h_2(q, p) + + where :math:`q` and :math:`p` are the position and momentum variables respectively, + and :math:`h_1` and :math:`h_2` Hamiltonian component functions. The exact + Hamiltonian flows for both the :math:`h_1` and :math:`h_2` components are assumed to + be tractable for subclasses of this class. + + By default :math:`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. + """ + + @abstractmethod + def h2_flow(self, state: ChainState, dt: ScalarLike): + """Apply exact flow map corresponding to `h2` Hamiltonian component. + + `state` argument is modified in place. + + Args: + state: State to start flow at. + dt: Time interval to simulate flow for. + """ + + @abstractmethod + def dh2_flow_dmom(self, dt: ScalarLike) -> tuple[matrices.Matrix, matrices.Matrix]: + """Derivatives of `h2_flow` flow map with respect to momentum argument. + + Args: + dt: Time interval flow simulated for. + + Returns: + Tuple `(dpos_dmom, dmom_dmom)` with :code:`dpos_dmom` a matrix representing + derivative (Jacobian) of position output of :py:meth:`h2_flow` with respect + to the value of the momentum component of the initial input state and + :code:`dmom_dmom` a matrix representing derivative (Jacobian) of momentum + output of :py:meth:`h2_flow` with respect to the value of the momentum + component of the initial input state. + """ + + +class EuclideanMetricSystem(TractableFlowSystem): r"""Hamiltonian system with a Euclidean metric on the position space. Here Euclidean metric is defined to mean a metric with a fixed positive definite @@ -303,30 +351,9 @@ def dh_dpos(self, state: ChainState) -> ArrayLike: return self.dh1_dpos(state) def h2_flow(self, state: ChainState, dt: ScalarLike): - """Apply exact flow map corresponding to `h2` Hamiltonian component. - - `state` argument is modified in place. - - Args: - state: State to start flow at. - dt: Time interval to simulate flow for. - """ state.pos += dt * self.dh2_dmom(state) def dh2_flow_dmom(self, dt: ScalarLike) -> tuple[matrices.Matrix, matrices.Matrix]: - """Derivatives of `h2_flow` flow map with respect to input momentum. - - Args: - dt: Time interval flow simulated for. - - Returns: - dpos_dmom: Matrix representing derivative (Jacobian) of position output of - `h2_flow` with respect to the value of the momentum component of the - initial input state. - dmom_dmom: 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 (dt * self.metric.inv, matrices.IdentityMatrix(self.metric.shape[0])) def sample_momentum(self, state: ChainState, rng: Generator) -> ArrayLike: @@ -460,7 +487,7 @@ class ConstrainedEuclideanMetricSystem(EuclideanMetricSystem): 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 :math:`q` with :math:` \partial c(q) M^{-1} p = 0` at all time + on the momentum :math:`q` with :math:`\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. @@ -491,7 +518,7 @@ class ConstrainedEuclideanMetricSystem(EuclideanMetricSystem): 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`. + :py:class:`mici.integrators.ConstrainedLeapfrogIntegrator`. References: @@ -512,80 +539,75 @@ def __init__( grad_neg_log_dens: Optional[GradientFunction] = None, jacob_constr: Optional[JacobianFunction] = None, ): - """ + r""" Args: neg_log_dens: 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(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 + manifold (if :code:`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 + :code:`constr(q) == 0` where :code:`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 :code:`neg_log_dens`) and the square-root of the determinant of the Gram matrix defined by - - .. code-block:: - - 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. + :code:`gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T` + where :code:`jacob_constr` is the Jacobian of the constraint function + :code:`constr` and :code:`metric` is the matrix representation of the + metric on the ambient space. constr: 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: 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 + distribution on *unconstrained* momentum vector. If :code:`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: Whether the `neg_log_dens` function specifies the + value is a :py:class:`mici.matrices.PositiveDefiniteMatrix` subclass it + is assumed to directly specify the metric matrix representation. + dens_wrt_hausdorff: Whether the :code:`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 + respect to the Hausdorff measure on the manifold directly (:code:`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 `const(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. + to the posterior distribution when conditioning on the event + :code:`constr(pos) == 0` (:code:`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: Function which given a position array returns the - derivative of `neg_log_dens` with respect to the position array + derivative of :code:`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. + second being the value of the :code:`neg_log_dens` evaluated at the + passed position array. If :code:`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 :code:`neg_log_dens` + automatically. jacob_constr: 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 - - .. code-block:: - - 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` + constraint function :code:`c = constr(q)` with respect to the position + array argument :code:`q`, returning the computed Jacobian as a 2D array + :code:`jacob` with :code:`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 + :code:`constr` evaluated at the passed position array. If :code:`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. + value) of :code:`constr` + automatically. """ super().__init__( neg_log_dens=neg_log_dens, @@ -606,7 +628,7 @@ def constr(self, state: ChainState) -> ArrayLike: state: State to compute value at. Returns: - Value of `constr(state.pos)` as 1D array. + Value of :code:`constr(state.pos)` as 1D array. """ return self._constr(state.pos) @@ -618,7 +640,7 @@ def jacob_constr(self, state: ChainState) -> ArrayLike: state: State to compute value at. Returns: - Value of Jacobian of `constr(state.pos)` as 2D array. + Value of Jacobian of :code:`constr(state.pos)` as 2D array. """ return self._jacob_constr(state.pos) @@ -631,16 +653,16 @@ def jacob_constr_inner_product( ) -> MatrixLike: """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`. + Computes :code:`jacob_constr_1 @ inner_product_matrix @ jacob_constr_2.T` + potentially exploiting any structure / sparsity in :code:`jacob_constr_1`, + :code:`jacob_constr_2` and :code:`inner_product_matrix`. Args: jacob_constr_1: First constraint Jacobian in product. inner_product_matrix: Positive-definite matrix defining inner-product between rows of two constraint Jacobians. jacob_constr_2: Second constraint Jacobian in product. Defaults to - `jacob_constr_1` if set to `None`. + :code:`jacob_constr_1` if set to :code:`None`. Returns Object corresponding to computed inner products of the constraint Jacobian @@ -653,12 +675,13 @@ def gram(self, state: ChainState) -> matrices.PositiveDefiniteMatrix: The Gram matrix as a position `q` is defined as - .. code-block:: + .. code:: 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. + where :code:`jacob_constr` is the Jacobian of the constraint function + :code:`constr` and :code:`metric` is the matrix representation of the metric on + the ambient space. Args: state: State to compute value at. @@ -687,13 +710,14 @@ def log_det_sqrt_gram(self, state: ChainState) -> ScalarLike: @abstractmethod def grad_log_det_sqrt_gram(self, state: ChainState) -> ArrayLike: - """Derivative of (half of) log-determinant of Gram matrix wrt position. + """Derivative of half of log-determinant of Gram matrix with respect to position Args: state: State to compute value at. Returns: - Value of `log_det_sqrt_gram(state)` derivative with respect to `state.pos`. + Value of :code:`log_det_sqrt_gram(state)` derivative with respect to + :code:`state.pos`. """ def h1(self, state: ChainState) -> ScalarLike: @@ -719,7 +743,7 @@ def project_onto_cotangent_space( co-tangent space of. Returns: - Projected momentum in the co-tangent space at `state.pos`. + Projected momentum in the co-tangent space at :code:`state.pos`. """ # Use parenthesis to force right-to-left evaluation to avoid # matrix-matrix products @@ -737,7 +761,8 @@ def sample_momentum(self, state: ChainState, rng: Generator) -> ArrayLike: class DenseConstrainedEuclideanMetricSystem(ConstrainedEuclideanMetricSystem): r"""Euclidean Hamiltonian system subject to a dense set of constraints. - See `ConstrainedEuclideanMetricSystem` for more details about constrained systems. + See :py:class:`ConstrainedEuclideanMetricSystem` for more details about constrained + systems. """ def __init__( @@ -755,100 +780,88 @@ def __init__( neg_log_dens: 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(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 + manifold (if :code:`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 + :code:`constr(q) == 0` where :code:`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 :code:`neg_log_dens`) and the square-root of the determinant of the Gram matrix defined by - - .. code-block:: - - 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: Function which given a position array return as a 1D array the value + :code:`gram(q) = jacob_constr(q) @ inv(metric) @ jacob_constr(q).T` + where :code:`jacob_constr` is the Jacobian of the constraint function + :code:`constr` and :code:`metric` is the matrix representation of the + metric on the ambient space. + constr: 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: 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 + distribution on *unconstrained* momentum vector. If :code:`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: Whether the `neg_log_dens` function specifies the + value is a :py:class:`mici.matrices.PositiveDefiniteMatrix` subclass it + is assumed to directly specify the metric matrix representation. + dens_wrt_hausdorff: Whether the :code:`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 + respect to the Hausdorff measure on the manifold directly (:code:`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 `const(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. + to the posterior distribution when conditioning on the event + :code:`constr(pos) == 0` (:code:`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: Function which given a position array returns the - derivative of `neg_log_dens` with respect to the position array + derivative of :code:`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. + second being the value of the :code:`neg_log_dens` evaluated at the + passed position array. If :code:`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 :code:`neg_log_dens` + automatically. jacob_constr: 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 - - .. code-block:: - - 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` + constraint function :code:`c = constr(q)` with respect to the position + array argument :code:`q`, returning the computed Jacobian as a 2D array + :code:`jacob` with :code:`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 + :code:`constr` evaluated at the passed position array. If :code:`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. + value) of :code:`constr` + automatically. mhp_constr: 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` - - .. code-block:: - - 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 - - .. code-block:: - - 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. + *matrix-Hessian-product* (MHP) of the constraint function :code:`constr` + with respect to the position array argument. The MHP is here defined as + a function of a :code:`(dim_constr, dim_pos)` shaped 2D array :code:`m` + as :code:`mhp(m) = sum(m[:, :, None] * hess[:, :, :], axis=(0, 1))` + where :code:`hess` is the :code:`(dim_constr, dim_pos, dim_pos)` shaped + vector-Hessian of :code:`c = constr(q)` with respect to :code:`q` i.e. + the array of second-order partial derivatives of such that + :code:`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 :code:`constr`, the second a 2D array corresponding to + the Jacobian of :code:`constr`, and the third the value of + :code:`constr`, all evaluated at the passed position array. If + :code:`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 :code:`constr` + automatically. """ super().__init__( neg_log_dens=neg_log_dens,