sode.py

import numpy as np
import warnings
from . import _sode
from ..integrator.integrator import Integrator
from ..stochastic import StochasticSolver, SMESolver
from ._noise import Wiener, PreSetWiener

__all__ = ["SIntegrator", "PlatenSODE", "PredCorr_SODE"]


class SIntegrator(Integrator):
    """
    A wrapper around stochastic ODE solvers.

    Parameters
    ----------
    system: qutip.StochasticSystem
        Quantum system in which states evolve.

    options: dict
        Options for the integrator.

    Class Attributes
    ----------------
    name : str
        The name of the integrator.

    supports_blackbox : bool
        If True, then the integrator calls only ``system.matmul``,
        ``system.matmul_data``, ``system.expect``, ``system.expect_data`` and
        ``isconstant``, ``isoper`` or ``issuper``. This allows the solver using
        the integrator to modify the system in creative ways. In particular,
        the solver may modify the system depending on *both* the time ``t``
        *and* the current ``state`` the system is being applied to.

        If the integrator calls any other methods, set to False.

    supports_time_dependent : bool
        If True, then the integrator supports time dependent systems. If False,
        ``supports_blackbox`` should usually be ``False`` too.

    integrator_options : dict
        A dictionary of options used by the integrator and their default
        values. Once initiated, ``self.options`` will be a dict with the same
        keys, not the full options object passed to the solver. Options' keys
        included here will be supported by the :cls:SolverOdeOptions.
    """
    _is_set = False
    _stepper_options = []

    def set_state(self, t, state0, generator):
        """
        Set the state of the SODE solver.

        Parameters
        ----------
        t : float
            Initial time

        state0 : qutip.Data
            Initial state.

        generator : numpy.random.generator
            Random number generator.
        """
        self.t = t
        self.state = state0
        stepper_opt = {
            key: self.options[key]
            for key in self._stepper_options
            if key in self.options
        }

        if isinstance(generator, PreSetWiener):
            self.wiener = generator
            if (
                generator.is_measurement
                and "measurement_noise" not in self._stepper_options
            ):
                raise NotImplementedError(
                    f"{type(self).__name__} does not support running"
                    " the evolution from measurements."
                )
            stepper_opt["measurement_noise"] = generator.is_measurement
        elif isinstance(generator, Wiener):
            self.wiener = generator
        else:
            num_collapse = len(self.rhs.sc_ops)
            self.wiener = Wiener(
                t, self.options["dt"], generator,
                (self.N_dw, num_collapse)
            )
        self.rhs._register_feedback(self.wiener)
        rhs = self.rhs(self.options)
        self.step_func = self.stepper(rhs, **stepper_opt).run
        self._is_set = True

    def get_state(self, copy=True):
        return self.t, self.state, self.wiener

    def integrate(self, t, copy=True):
        """
        Evolve to t.

        Before calling `integrate` for the first time, the initial state should
        be set with `set_state`.

        Parameters
        ----------
        t : float
            Time to integrate to, should be larger than the previous time.

        copy : bool [True]
            Whether to return a copy of the state or the state itself.

        Returns
        -------
        (t, state, noise) : (float, qutip.Data, np.ndarray)
            The state of the solver at ``t``.
        """
        raise NotImplementedError

    def mcstep(self, t, copy=True):
        raise NotImplementedError

    def reset(self, hard=False):
        if self._is_set:
            state = self.get_state()
            self.set_state(*state)
        if hard:
            raise NotImplementedError(
                "Changing stochastic integrator "
                "options is not supported."
            )


class _Explicit_Simple_Integrator(SIntegrator):
    """
    Stochastic evolution solver
    """

    integrator_options = {
        "dt": 0.001,
        "tol": 1e-10,
    }
    stepper = None
    N_dw = 0

    def __init__(self, rhs, options):
        self._options = self.integrator_options.copy()
        self.options = options
        self.rhs = rhs

    def integrate(self, t, copy=True):
        delta_t = t - self.t
        dt = self.options["dt"]
        if delta_t < 0:
            raise ValueError("Integration time, can't be negative.")
        elif delta_t < 0.5 * dt:
            warnings.warn(
                f"Step under minimum step ({dt}), skipped.",
                RuntimeWarning
            )
            return self.t, self.state, np.zeros(self.N_dw)

        N, extra = np.divmod(delta_t, dt)
        N = int(N)
        if extra > 0.5 * dt:
            # Not a whole number of steps, round to higher
            N += 1
        dW = self.wiener.dW(self.t, N)

        self.state = self.step_func(self.t, self.state, dt, dW, N)
        self.t += dt * N

        return self.t, self.state, np.sum(dW[:, 0, :], axis=0)

    @property
    def options(self):
        """
        Supported options by Explicit Stochastic Integrators:

        dt : float, default: 0.001
            Internal time step.

        tol : float, default: 1e-10
            Tolerance for the time steps.
        """
        return self._options

    @options.setter
    def options(self, new_options):
        Integrator.options.fset(self, new_options)


class _Implicit_Simple_Integrator(_Explicit_Simple_Integrator):
    """
    Stochastic evolution solver
    """

    integrator_options = {
        "dt": 0.001,
        "tol": 1e-10,
        "solve_method": None,
        "solve_options": {},
    }
    _stepper_options = ["solve_method", "solve_options"]
    stepper = None
    N_dw = 0

    @property
    def options(self):
        """
        Supported options by Implicit Stochastic Integrators:

        dt : float, default: 0.001
            Internal time step.

        tol : float, default: 1e-10
            Tolerance for the time steps.

        solve_method : str, default: None
            Method used for solver the ``Ax=b`` of the implicit step.
            Accept methods supported by :func:`qutip.core.data.solve`.
            When the system is constant, the inverse of the matrix ``A`` can be
            used by entering ``inv``.

        solve_options : dict, default: {}
            Options to pass to the call to :func:`qutip.core.data.solve`.
        """
        return self._options

    @options.setter
    def options(self, new_options):
        Integrator.options.fset(self, new_options)


class PlatenSODE(_Explicit_Simple_Integrator):
    """
    Explicit scheme, creates the Milstein using finite differences
    instead of analytic derivatives. Also contains some higher order
    terms, thus converges better than Milstein while staying strong
    order 1.0.  Does not require derivatives. See eq. (7.47) of chapter 7 of
    H.-P. Breuer and F. Petruccione, *The Theory of Open Quantum Systems*.

    - Order: strong 1, weak 2
    """
    integrator_options = {
        "dt": 0.001,
        "tol": 1e-10,
    }
    stepper = _sode.Platen
    N_dw = 1
    _stepper_options = ["measurement_noise"]


class PredCorr_SODE(_Explicit_Simple_Integrator):
    """
    Generalization of the trapezoidal method to stochastic differential
    equations. More stable than explicit methods.  See eq. (5.4) of
    chapter 15.5 of Peter E. Kloeden and Exkhard Platen,
    *Numerical Solution of Stochastic Differential Equations*.

    - Order strong 0.5, weak 1.0
    - Codes to only correct the stochastic part (:math:`\\alpha=0`,
      :math:`\\eta=1/2`): ``'pred-corr'``, ``'predictor-corrector'`` or
      ``'pc-euler'``
    - Codes to correct both the stochastic and deterministic parts
      (:math:`\\alpha=1/2`, :math:`\\eta=1/2`): ``'pc-euler-imp'``,
      ``'pc-euler-2'`` or ``'pred-corr-2'``
    """

    integrator_options = {
        "dt": 0.001,
        "tol": 1e-10,
        "alpha": 0.0,
        "eta": 0.5,
    }
    stepper = _sode.PredCorr
    N_dw = 1
    _stepper_options = ["alpha", "eta", "measurement_noise"]

    @property
    def options(self):
        """
        Supported options by Explicit Stochastic Integrators:

        dt : float, default: 0.001
            Internal time step.

        tol : float, default: 1e-10
            Tolerance for the time steps.

        alpha : float, default: 0.
            Implicit factor to the drift.
            eff_drift ~= drift(t) * (1-alpha) + drift(t+dt) * alpha

        eta : float, default: 0.5
            Implicit factor to the diffusion.
            eff_diffusion ~= diffusion(t) * (1-eta) + diffusion(t+dt) * eta
        """
        return self._options

    @options.setter
    def options(self, new_options):
        Integrator.options.fset(self, new_options)


StochasticSolver.add_integrator(PlatenSODE, "platen")
SMESolver.add_integrator(PredCorr_SODE, "pred_corr")