<commit_before>"""Symplectic integrators for simulation of Hamiltonian dynamics."""

from abc import ABC, abstractmethod
from mici.errors import NonReversibleStepError, AdaptationError
from mici.solvers import (maximum_norm, solve_fixed_point_direct,
                          solve_projection_onto_manifold_quasi_newton)

__pdoc__ = {}


class Integrator(ABC):
    """Base class for integrators."""

    def __init__(self, system, step_size=None):
        """
        Args:
            system (mici.systems.System): Hamiltonian system to integrate the
                dynamics of.
            step_size (float or None): Integrator time step. If set to `None`
                (the default) it is assumed that a step size adapter will be
                used to set the step size before calling the `step` method.
        """
        self.system = system
        self.step_size = step_size

    def step(self, state):
        """Perform a single integrator step from a supplied state.

        Args:
            state (mici.states.ChainState): System state to perform integrator
                step from.

        Returns:
            new_state (mici.states.ChainState): New object corresponding to
                stepped state.
        """
        if self.step_size is None:
            raise AdaptationError(
                'Integrator `step_size` is `None`. This value should only be '
                'used if a step size adapter is being used to set the step '
                'size.')
        state = state.copy()
        self._step(state, state.dir * self.step_size)
        return state

    @abstractmethod
    def _step(self, state, dt):
        """Implementation of single integrator step.

        Args:
            state (mici.states.ChainState): System state to perform integrator
                step from. Updated in place.
            dt (float): Integrator time step. May be positive or negative.
        """



class ExplicitLeapfrogIntegrator(Integrator):
    r"""
    Leapfrog integrator for Hamiltonian systems 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

    \[ 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\) are Hamiltonian component functions for which the
    exact flows can be computed.

    `LeapfrogIntegrator` is an alias for `ExplicitLeapfrogIntegrator`.
    """

    def __init__(self, system, step_size=None):
        if not hasattr(system, 'h1_flow') or not hasattr(system, 'h2_flow'):
            raise ValueError(
                'Explicit leapfrog integrator can only be used for systems '
                'with explicit `h1_flow` and `h2_flow` Hamiltonian component '
                'flow maps. For systems in which only `h1_flow` is available '
                'the `ImplicitLeapfrogIntegrator` class may be used instead.')
        super().__init__(system, step_size)

    def _step(self, state, dt):
        self.system.h1_flow(state, 0.5 * dt)
        self.system.h2_flow(state, dt)
        self.system.h1_flow(state, 0.5 * dt)


LeapfrogIntegrator = ExplicitLeapfrogIntegrator
__pdoc__['LeapfrogIntegrator'] = False


class ImplicitLeapfrogIntegrator(Integrator):
    r"""
    Implicit leapfrog integrator for Hamiltonian with non-separable component.

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

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

    where \(q\) and \(p\) are the position and momentum variables respectively,<commit_msg>Rename ExplicitLeapfrogIntegrator to LeapfrogIntegrator and remove alias<commit_after>"""Symplectic integrators for simulation of Hamiltonian dynamics."""

from abc import ABC, abstractmethod
from mici.errors import NonReversibleStepError, AdaptationError
from mici.solvers import (maximum_norm, solve_fixed_point_direct,
                          solve_projection_onto_manifold_quasi_newton)


class Integrator(ABC):
    """Base class for integrators."""

    def __init__(self, system, step_size=None):
        """
        Args:
            system (mici.systems.System): Hamiltonian system to integrate the
                dynamics of.
            step_size (float or None): Integrator time step. If set to `None`
                (the default) it is assumed that a step size adapter will be
                used to set the step size before calling the `step` method.
        """
        self.system = system
        self.step_size = step_size

    def step(self, state):
        """Perform a single integrator step from a supplied state.

        Args:
            state (mici.states.ChainState): System state to perform integrator
                step from.

        Returns:
            new_state (mici.states.ChainState): New object corresponding to
                stepped state.
        """
        if self.step_size is None:
            raise AdaptationError(
                'Integrator `step_size` is `None`. This value should only be '
                'used if a step size adapter is being used to set the step '
                'size.')
        state = state.copy()
        self._step(state, state.dir * self.step_size)
        return state

    @abstractmethod
    def _step(self, state, dt):
        """Implementation of single integrator step.

        Args:
            state (mici.states.ChainState): System state to perform integrator
                step from. Updated in place.
            dt (float): Integrator time step. May be positive or negative.
        """



class LeapfrogIntegrator(Integrator):
    r"""
    Leapfrog integrator for Hamiltonian systems 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

    \[ 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\) are Hamiltonian component functions for which the
    exact flows can be computed.

    `LeapfrogIntegrator` is an alias for `ExplicitLeapfrogIntegrator`.
    """

    def __init__(self, system, step_size=None):
        if not hasattr(system, 'h1_flow') or not hasattr(system, 'h2_flow'):
            raise ValueError(
                'Explicit leapfrog integrator can only be used for systems '
                'with explicit `h1_flow` and `h2_flow` Hamiltonian component '
                'flow maps. For systems in which only `h1_flow` is available '
                'the `ImplicitLeapfrogIntegrator` class may be used instead.')
        super().__init__(system, step_size)

    def _step(self, state, dt):
        self.system.h1_flow(state, 0.5 * dt)
        self.system.h2_flow(state, dt)
        self.system.h1_flow(state, 0.5 * dt)


class ImplicitLeapfrogIntegrator(Integrator):
    r"""
    Implicit leapfrog integrator for Hamiltonian with non-separable component.

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

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

    where \(q\) and \(p\) are the position and momentum variables respectively,
{'fraction_non_alphanumeric': 0.08018804566823372, 'fraction_numerical': 0.0048354600402955, 'mean_word_length': 2.429755872869645, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5777209', 'n_tokens_mistral': 2081, 'n_tokens_neox': 1991, 'n_words': 744}