<commit_before>#
# Relativistic MCMC method
#
# This file is part of PINTS (https://github.com/pints-team/pints/) which is
# released under the BSD 3-clause license. See accompanying LICENSE.md for
# copyright notice and full license details.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import pints
import numpy as np
from scipy.interpolate import interp1d
from scipy.optimize import root


class RelativisticMCMC(pints.SingleChainMCMC):
    r"""
    Implements Relativistic Monte Carlo as described in [1]_.

    Uses a physical analogy of a particle moving across a landscape under
    Hamiltonian dynamics to aid efficient exploration of parameter space.
    Introduces an auxilary variable -- the momentum (``p_i``) of a particle
    moving in dimension ``i`` of negative log posterior space -- which
    supplements the position (``q_i``) of the particle in parameter space. The
    particle's motion is dictated by solutions to Hamilton's equations,

    .. math::
        dq_i/dt &=   \partial H/\partial p_i\\
        dp_i/dt &= - \partial H/\partial q_i.

    The Hamiltonian is given by,

    .. math::
        H(q,p) &=       U(q)       +        KE(p)\\
               &= -\text{log}(p(q|X)p(q)) +
                    mc^2 (\Sigma_{i=1}^{d} p_i^2 / (m^2 c^2) + 1)^{0.5}

    where ``d`` is the dimensionality of model, ``m`` is the scalar 'mass'
    given to each particle (chosen to be 1 as default) and ``c`` is the
    speed of light (chosen to be 10 by default).

    To numerically integrate Hamilton's equations, it is essential to use a
    sympletic discretisation routine, of which the most typical approach is
    the leapfrog method,

    .. math::
        p_i(t + \epsilon/2) &= p_i(t) - (\epsilon/2) d U(q_i(t))/dq_i\\
        q_i(t + \epsilon) &= q_i(t) +
                    \epsilon M^{-1}(p_i(t + \epsilon/2)) p_i(t + \epsilon/2)\\
        p_i(t + \epsilon) &= p_i(t + \epsilon/2) -
                                (\epsilon/2) d U(q_i(t + \epsilon))/dq_i

    where relativistic mass (a scalar) is,

    .. math::
        M(p) = m (\Sigma_{i=1}^{d} p_i^2 / (m^2 c^2) + 1)^{0.5}

    In particular, the algorithm we implement follows eqs. in section 2.1 of
    [1]_.

    Extends :class:`SingleChainMCMC`.

    References
    ----------
    .. [1] "Relativistic Monte Carlo". <NAME>, <NAME>,
            <NAME>, <NAME>, <NAME>,
            2017, Proceedings of Machine Learning Research.
    """
    def __init__(self, x0, sigma0=None):
        super(RelativisticMCMC, self).__init__(x0, sigma0)

        # Set initial state
        self._running = False
        self._ready_for_tell = False

        # Current point in the Markov chain
        self._current = None            # Aka current_q in the chapter
        self._current_energy = None     # Aka U(current_q) = -log_pdf
        self._current_gradient = None
        self._current_momentum = None   # Aka current_p

        # Current point in the leapfrog iterations
        self._momentum = None       # Aka p in the chapter
        self._position = None       # Aka q in the chapter
        self._gradient = None       # Aka grad_U(q) in the chapter

        # Iterations, acceptance monitoring, and leapfrog iterations
        self._mcmc_iteration = 0
        self._mcmc_acceptance = 0
        self._frog_iteration = 0

        # Default number of leapfrog iterations
        self._n_frog_iterations = 20

        # Default integration step size for leapfrog algorithm
        self._epsilon = 0.1
        self._step_size = None
        self._mass = 1
        self._c = 10
        self.set_leapfrog_step_size(np.diag(self._sigma0))

        # Divergence checking
        # Create a vector of divergent iterations
        self._divergent = np.asarray([], dtype='int')

        # Default threshold for Hamiltonian divergences
        # (currently set to match Stan)
        self._hamiltonian_threshold = 10**3

    def ask(self):
        """ See :meth:`SingleChainMCMC.ask()`. """
        # Check ask/tell pattern
        if self._ready_for_tell:
            raise RuntimeError('Ask() called when expecting call to tell().')

        # Initialise on first call
        if not self._running:
            self._running = True
            self._mc2 = self._mass * self._c**2
            self._m2c2 = self._mass**2 * self._c**2
            self._calculate_momentum_distribution()

        # Notes:
        #  Ask is responsible for updating the position, which is the point
        #   returned to the user
        #  Tell is then responsible for updating the momentum, which uses the
        #   gradient at this new point
        #  The MCMC step happens in tell, and does not require any new
        #   information (it uses the log_pdf and gradient of the final point
        #   in the leapfrog run).

        # Very first iteration
        if self._current is None:

            # Ask for the pdf and gradient of x0
            self._ready_for_tell = True
            return np.array(self._x0, copy=True)

        # First iteration of a run of leapfrog iterations
        if self._frog_iteration == 0:

            # Sample random momentum for current point
            self._current_momentum = self._sample_momentum()

            # First leapfrog position is the current sample in the chain
            self._position = np.array(self._current, copy=True)
            self._gradient = np.array(self._current_gradient, copy=True)
            self._momentum = np.array(self._current_momentum, copy=True)

            # Perform a half-step before starting iteration 0 below
            self._momentum -= self._scaled_epsilon * self._gradient * 0.5

        # Perform a leapfrog step for the position
        squared = np.sum(np.array(self._momentum)**2)
        relativistic_mass = self._mass * np.sqrt(squared / self._m2c2 + 1)
        self._position += (
            self._scaled_epsilon * self._momentum / relativistic_mass)

        # Ask for the pdf and gradient of the current leapfrog position
        # Using this, the leapfrog step for the momentum is performed in tell()
        self._ready_for_tell = True
        return np.array(self._position, copy=True)

    def _calculate_momentum_distribution(self):
        """Calculate an approximation to the CDF of momentum magnitude.

        The purpose of this method is to calculate self._inv_cdf, a function
        giving a numerical approximation to the inverse cumulative distribution
        function of the magnitude of the momentum vector. This function can be
        used to perform inverse transform sampling to generate samples of the
        momentum.

        This method should be called before any sampling is performed. As
        implemented, it is called upon the first call to self.ask.

        Numerical integration of the momentum density is performed using the
        trapezoidal rule.
        """
        # Calculate the value of p (momentum) corresponding to the maximum of
        # the pdf
        def logpdf_deriv(p):
            # logpdf_deriv(p) = 0 is the simplified form of d/dp(logpdf(p)) = 0
            d = np.sqrt(self._mass * (self._n_parameters - 1)) \
                * (p**2 + self._m2c2) ** 0.25 / np.sqrt(self._c * self._mass) \
                - p
            return d
<commit_msg>change sphinx inv_cdf<commit_after>#
# Relativistic MCMC method
#
# This file is part of PINTS (https://github.com/pints-team/pints/) which is
# released under the BSD 3-clause license. See accompanying LICENSE.md for
# copyright notice and full license details.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import pints
import numpy as np
from scipy.interpolate import interp1d
from scipy.optimize import root


class RelativisticMCMC(pints.SingleChainMCMC):
    r"""
    Implements Relativistic Monte Carlo as described in [1]_.

    Uses a physical analogy of a particle moving across a landscape under
    Hamiltonian dynamics to aid efficient exploration of parameter space.
    Introduces an auxilary variable -- the momentum (``p_i``) of a particle
    moving in dimension ``i`` of negative log posterior space -- which
    supplements the position (``q_i``) of the particle in parameter space. The
    particle's motion is dictated by solutions to Hamilton's equations,

    .. math::
        dq_i/dt &=   \partial H/\partial p_i\\
        dp_i/dt &= - \partial H/\partial q_i.

    The Hamiltonian is given by,

    .. math::
        H(q,p) &=       U(q)       +        KE(p)\\
               &= -\text{log}(p(q|X)p(q)) +
                    mc^2 (\Sigma_{i=1}^{d} p_i^2 / (m^2 c^2) + 1)^{0.5}

    where ``d`` is the dimensionality of model, ``m`` is the scalar 'mass'
    given to each particle (chosen to be 1 as default) and ``c`` is the
    speed of light (chosen to be 10 by default).

    To numerically integrate Hamilton's equations, it is essential to use a
    sympletic discretisation routine, of which the most typical approach is
    the leapfrog method,

    .. math::
        p_i(t + \epsilon/2) &= p_i(t) - (\epsilon/2) d U(q_i(t))/dq_i\\
        q_i(t + \epsilon) &= q_i(t) +
                    \epsilon M^{-1}(p_i(t + \epsilon/2)) p_i(t + \epsilon/2)\\
        p_i(t + \epsilon) &= p_i(t + \epsilon/2) -
                                (\epsilon/2) d U(q_i(t + \epsilon))/dq_i

    where relativistic mass (a scalar) is,

    .. math::
        M(p) = m (\Sigma_{i=1}^{d} p_i^2 / (m^2 c^2) + 1)^{0.5}

    In particular, the algorithm we implement follows eqs. in section 2.1 of
    [1]_.

    Extends :class:`SingleChainMCMC`.

    References
    ----------
    .. [1] "Relativistic Monte Carlo". <NAME>, <NAME>,
            <NAME>, <NAME>, <NAME>,
            2017, Proceedings of Machine Learning Research.
    """
    def __init__(self, x0, sigma0=None):
        super(RelativisticMCMC, self).__init__(x0, sigma0)

        # Set initial state
        self._running = False
        self._ready_for_tell = False

        # Current point in the Markov chain
        self._current = None            # Aka current_q in the chapter
        self._current_energy = None     # Aka U(current_q) = -log_pdf
        self._current_gradient = None
        self._current_momentum = None   # Aka current_p

        # Current point in the leapfrog iterations
        self._momentum = None       # Aka p in the chapter
        self._position = None       # Aka q in the chapter
        self._gradient = None       # Aka grad_U(q) in the chapter

        # Iterations, acceptance monitoring, and leapfrog iterations
        self._mcmc_iteration = 0
        self._mcmc_acceptance = 0
        self._frog_iteration = 0

        # Default number of leapfrog iterations
        self._n_frog_iterations = 20

        # Default integration step size for leapfrog algorithm
        self._epsilon = 0.1
        self._step_size = None
        self._mass = 1
        self._c = 10
        self.set_leapfrog_step_size(np.diag(self._sigma0))

        # Divergence checking
        # Create a vector of divergent iterations
        self._divergent = np.asarray([], dtype='int')

        # Default threshold for Hamiltonian divergences
        # (currently set to match Stan)
        self._hamiltonian_threshold = 10**3

    def ask(self):
        """ See :meth:`SingleChainMCMC.ask()`. """
        # Check ask/tell pattern
        if self._ready_for_tell:
            raise RuntimeError('Ask() called when expecting call to tell().')

        # Initialise on first call
        if not self._running:
            self._running = True
            self._mc2 = self._mass * self._c**2
            self._m2c2 = self._mass**2 * self._c**2
            self._calculate_momentum_distribution()

        # Notes:
        #  Ask is responsible for updating the position, which is the point
        #   returned to the user
        #  Tell is then responsible for updating the momentum, which uses the
        #   gradient at this new point
        #  The MCMC step happens in tell, and does not require any new
        #   information (it uses the log_pdf and gradient of the final point
        #   in the leapfrog run).

        # Very first iteration
        if self._current is None:

            # Ask for the pdf and gradient of x0
            self._ready_for_tell = True
            return np.array(self._x0, copy=True)

        # First iteration of a run of leapfrog iterations
        if self._frog_iteration == 0:

            # Sample random momentum for current point
            self._current_momentum = self._sample_momentum()

            # First leapfrog position is the current sample in the chain
            self._position = np.array(self._current, copy=True)
            self._gradient = np.array(self._current_gradient, copy=True)
            self._momentum = np.array(self._current_momentum, copy=True)

            # Perform a half-step before starting iteration 0 below
            self._momentum -= self._scaled_epsilon * self._gradient * 0.5

        # Perform a leapfrog step for the position
        squared = np.sum(np.array(self._momentum)**2)
        relativistic_mass = self._mass * np.sqrt(squared / self._m2c2 + 1)
        self._position += (
            self._scaled_epsilon * self._momentum / relativistic_mass)

        # Ask for the pdf and gradient of the current leapfrog position
        # Using this, the leapfrog step for the momentum is performed in tell()
        self._ready_for_tell = True
        return np.array(self._position, copy=True)

    def _calculate_momentum_distribution(self):
        """Calculate an approximation to the CDF of momentum magnitude.

        The purpose of this method is to calculate :member:_inv_cdf, a function
        giving a numerical approximation to the inverse cumulative distribution
        function of the magnitude of the momentum vector. This function can be
        used to perform inverse transform sampling to generate samples of the
        momentum.

        This method should be called before any sampling is performed. As
        implemented, it is called upon the first call to self.ask.

        Numerical integration of the momentum density is performed using the
        trapezoidal rule.
        """
        # Calculate the value of p (momentum) corresponding to the maximum of
        # the pdf
        def logpdf_deriv(p):
            # logpdf_deriv(p) = 0 is the simplified form of d/dp(logpdf(p)) = 0
            d = np.sqrt(self._mass * (self._n_parameters - 1)) \
                * (p**2 + self._m2c2) ** 0.25 / np.sqrt(self._c * self._mass) \
                - p
            return d

{'fraction_non_alphanumeric': 0.08783645655877342, 'fraction_numerical': 0.010766609880749574, 'mean_word_length': 2.207167832167832, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2107013', 'n_tokens_mistral': 4341, 'n_tokens_neox': 4053, 'n_words': 1466}