Merge pull request #2318 from Ericgig/feature.stochastic_from_experim…

…ents

Add `run_from_experiment` to stochastic solvers.

doc/changes/2318.feature

@@ -0,0 +1,2 @@
+Create `run_from_experiment`, which allows to run stochastic evolution from
+know noise or measurements.

qutip/solver/multitraj.py

@@ -234,7 +234,10 @@ def run(self, state, tlist, ntraj=1, *,
     def _initialize_run_one_traj(self, seed, state, tlist, e_ops,
                                  **integrator_kwargs):
         result = self._trajectory_resultclass(e_ops, self.options)
-        generator = self._get_generator(seed)
+        if "generator" in integrator_kwargs:
+            generator = integrator_kwargs.pop("generator")
+        else:
+            generator = self._get_generator(seed)
         self._integrator.set_state(tlist[0], state, generator,
                                    **integrator_kwargs)
         result.add(tlist[0], self._restore_state(state, copy=False))
@@ -287,11 +290,6 @@ def _get_generator(self, seed):
         If the ``seed`` has a ``random`` method, it will be used as the
         generator.
         """
-        if hasattr(seed, 'random'):
-            # We check for the method, not the type to accept pseudo non-random
-            # generator for debug/testing purpose.
-            return seed
-
         if self.options['bitgenerator']:
             bit_gen = getattr(np.random, self.options['bitgenerator'])
             generator = np.random.Generator(bit_gen(seed))

qutip/solver/sode/_noise.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-__all__ = ["Wiener"]
+__all__ = ["Wiener", "PreSetWiener"]
 
 
 class Wiener:
@@ -10,31 +10,84 @@ class Wiener:
     def __init__(self, t0, dt, generator, shape):
         self.t0 = t0
         self.dt = dt
-        self.generator = generator
-        self.t_end = t0
         self.shape = shape
-        self.process = np.zeros((1,) + shape, dtype=float)
+        self.generator = generator
+        self.noise = np.zeros((0,) + shape, dtype=float)
+        self.last_W = np.zeros(shape[-1], dtype=float)
+        self.idx_last_0 = 0
 
-    def _extend(self, t):
-        N_new_vals = int((t - self.t_end + self.dt*0.01) // self.dt)
+    def _extend(self, idx):
+        N_new_vals = idx - self.noise.shape[0]
         dW = self.generator.normal(
             0, np.sqrt(self.dt), size=(N_new_vals,) + self.shape
         )
-        W = self.process[-1, :, :] + np.cumsum(dW, axis=0)
-        self.process = np.concatenate((self.process, W), axis=0)
-        self.t_end = self.t0 + (self.process.shape[0] - 1) * self.dt
+        self.noise = np.concatenate((self.noise, dW), axis=0)
 
     def dW(self, t, N):
-        if t + N * self.dt > self.t_end:
-            self._extend(t + N * self.dt)
-        idx0 = int((t - self.t0 + self.dt * 0.01) // self.dt)
-        return np.diff(self.process[idx0:idx0 + N + 1, :, :], axis=0)
+        # Find the index of t.
+        # Rounded to the closest step, but only multiple of dt are expected.
+        idx0 = round((t - self.t0) / self.dt)
+        if idx0 + N - 1 >= self.noise.shape[0]:
+            self._extend(idx0 + N)
+        return self.noise[idx0:idx0 + N, :, :]
 
     def __call__(self, t):
-        if t > self.t_end:
-            self._extend(t)
-        idx = int((t - self.t0 + self.dt * 0.01) // self.dt)
-        return self.process[idx, 0, :]
+        """
+        Return the Wiener process at the closest ``dt`` step to ``t``.
+        """
+        # The Wiener process is not used directly in the evolution, so it's
+        # less optimized than the ``dW`` method.
+
+        # Find the index of t.
+        # Rounded to the closest step, but only multiple of dt are expected.
+        idx = round((t - self.t0) / self.dt)
+        if idx >= self.noise.shape[0]:
+            self._extend(idx + 1)
+
+        if self.idx_last_0 > idx:
+            # Before last call, reseting
+            self.idx_last_0 = 0
+            self.last_W = np.zeros(self.shape[-1], dtype=float)
+
+        self.last_W = self.last_W + np.sum(
+            self.noise[self.idx_last_0:idx+1, 0, :], axis=0
+        )
+
+        self.idx_last_0 = idx
+        return self.last_W
+
+
+class PreSetWiener(Wiener):
+    def __init__(self, noise, tlist, n_sc_ops, heterodyne, is_measurement):
+        if heterodyne:
+            if noise.shape != (n_sc_ops/2, 2, len(tlist)-1):
+                raise ValueError(
+                    "Noise is not of the expected shape: "
+                    f"{(n_sc_ops/2, 2, len(tlist)-1)}"
+                )
+            noise = np.reshape(noise, (n_sc_ops, len(tlist)-1), "C")
+        else:
+            if noise.shape != (n_sc_ops, len(tlist)-1):
+                raise ValueError(
+                    "Noise is not of the expected shape: "
+                    f"{(n_sc_ops, len(tlist)-1)}"
+                )
+
+        self.t0 = tlist[0]
+        self.dt = tlist[1] - tlist[0]
+        self.shape = noise.shape[1:]
+        self.noise = noise.T[:, np.newaxis, :].copy()
+        self.last_W = np.zeros(self.shape[-1], dtype=float)
+        self.idx_last_0 = 0
+        self.is_measurement = is_measurement
+        if self.is_measurement:
+            # Measurements is scaled as <M> + dW / dt
+            self.noise *= self.dt
+            if heterodyne:
+                self.noise /= 2**0.5
+
+    def _extend(self, N):
+        raise ValueError("Requested time is outside the integration range.")
 
 
 class _Noise:

qutip/solver/sode/_sode.pyx

@@ -11,9 +11,11 @@ import numpy as np
 
 cdef class Euler:
     cdef _StochasticSystem system
+    cdef bint measurement_noise
 
-    def __init__(self, _StochasticSystem system):
+    def __init__(self, _StochasticSystem system, measurement_noise=False):
         self.system = system
+        self.measurement_noise = measurement_noise
 
     @cython.wraparound(False)
     def run(
@@ -37,9 +39,16 @@ cdef class Euler:
         """
         cdef int i
         cdef _StochasticSystem system = self.system
+        cdef list expect
 
         cdef Data a = system.drift(t, state)
         b = system.diffusion(t, state)
+
+        if self.measurement_noise:
+            expect = system.expect(t, state)
+            for i in range(system.num_collapse):
+                dW[0, i] -= expect[i].real * dt
+
         cdef Data new_state = _data.add(state, a, dt)
         for i in range(system.num_collapse):
             new_state = _data.add(new_state, b[i], dW[0, i])
@@ -72,6 +81,12 @@ cdef class Platen(Euler):
         cdef list d2 = system.diffusion(t, state)
         cdef Data Vt, out
         cdef list Vp, Vm
+        cdef list expect
+
+        if self.measurement_noise:
+            expect = system.expect(t, state)
+            for i in range(system.num_collapse):
+                dW[0, i] -= expect[i].real * dt
 
         out = _data.mul(d1, 0.5)
         Vt = d1.copy()
@@ -106,6 +121,9 @@ cdef class Platen(Euler):
 
 
 cdef class Explicit15(Euler):
+    def __init__(self, _StochasticSystem system):
+        self.system = system
+
     @cython.boundscheck(False)
     @cython.wraparound(False)
     @cython.cdivision(True)
@@ -235,9 +253,11 @@ cdef class Explicit15(Euler):
 
 cdef class Milstein:
     cdef _StochasticSystem system
+    cdef bint measurement_noise
 
-    def __init__(self, _StochasticSystem system):
-        self.system = system
+    def __init__(self, _StochasticSystem system, measurement_noise=False):
+            self.system = system
+            self.measurement_noise = measurement_noise
 
     @cython.wraparound(False)
     def run(self, double t, Data state, double dt, double[:, :, ::1] dW, int ntraj):
@@ -273,6 +293,11 @@ cdef class Milstein:
         iadd_dense(out, state, 1)
         iadd_dense(out, system.a(), dt)
 
+        if self.measurement_noise:
+            expect = system.expect(t, state)
+            for i in range(system.num_collapse):
+                dW[0, i] -= system.expect_i(i).real * dt
+
         for i in range(num_ops):
             iadd_dense(out, system.bi(i), dW[0, i])
 
@@ -289,11 +314,17 @@ cdef class PredCorr:
     cdef Dense euler
     cdef double alpha, eta
     cdef _StochasticSystem system
+    cdef bint measurement_noise
 
-    def __init__(self, _StochasticSystem system, double alpha=0., double eta=0.5):
+    def __init__(
+        self, _StochasticSystem system,
+        double alpha=0., double eta=0.5,
+        measurement_noise=False
+    ):
         self.system = system
         self.alpha = alpha
         self.eta = eta
+        self.measurement_noise = measurement_noise
 
     @cython.wraparound(False)
     def run(self, double t, Data state, double dt, double[:, :, ::1] dW, int ntraj):
@@ -324,6 +355,11 @@ cdef class PredCorr:
 
         system.set_state(t, state)
 
+        if self.measurement_noise:
+            expect = system.expect(t, state)
+            for i in range(system.num_collapse):
+                dW[0, i] -= system.expect_i(i).real * dt
+
         imul_dense(out, 0.)
         iadd_dense(out, state, 1)
         iadd_dense(out, system.a(), dt * (1-alpha))
@@ -350,6 +386,10 @@ cdef class PredCorr:
 
 
 cdef class Taylor15(Milstein):
+    def __init__(self, _StochasticSystem system):
+        self.system = system
+        self.measurement_noise = False
+
     @cython.boundscheck(False)
     @cython.wraparound(False)
     cdef Data step(self, double t, Dense state, double dt, double[:, :] dW, Dense out):
@@ -398,17 +438,17 @@ cdef class Milstein_imp:
     cdef double prev_dt
     cdef dict imp_opt
 
-    def __init__(self, _StochasticSystem system, imp_method=None, imp_options={}):
+    def __init__(self, _StochasticSystem system, solve_method=None, solve_options={}):
         self.system = system
         self.prev_dt = 0
-        if imp_method == "inv":
+        if solve_method == "inv":
             if not self.system.L.isconstant:
                 raise TypeError("The 'inv' integration method requires that the system Hamiltonian or Liouvillian be constant.")
             self.use_inv = True
             self.imp_opt = {}
         else:
             self.use_inv = False
-            self.imp_opt = {"method": imp_method, "options": imp_options}
+            self.imp_opt = {"method": solve_method, "options": solve_options}
 
 
     @cython.wraparound(False)

qutip/solver/sode/itotaylor.py

@@ -17,8 +17,13 @@ class EulerSODE(_Explicit_Simple_Integrator):
 
     - Order: 0.5
     """
+    integrator_options = {
+        "dt": 0.001,
+        "tol": 1e-10,
+    }
     stepper = _sode.Euler
     N_dw = 1
+    _stepper_options = ["measurement_noise"]
 
 
 class Milstein_SODE(_Explicit_Simple_Integrator):
@@ -30,8 +35,13 @@ class Milstein_SODE(_Explicit_Simple_Integrator):
 
     - Order strong 1.0
     """
+    integrator_options = {
+        "dt": 0.001,
+        "tol": 1e-10,
+    }
     stepper = _sode.Milstein
     N_dw = 1
+    _stepper_options = ["measurement_noise"]
 
 
 class Taylor1_5_SODE(_Explicit_Simple_Integrator):

qutip/solver/sode/sode.py

@@ -3,7 +3,7 @@
 from . import _sode
 from ..integrator.integrator import Integrator
 from ..stochastic import StochasticSolver, SMESolver
-from ._noise import Wiener
+from ._noise import Wiener, PreSetWiener
 
 __all__ = ["SIntegrator", "PlatenSODE", "PredCorr_SODE"]
 
@@ -65,7 +65,24 @@ def set_state(self, t, state0, generator):
         """
         self.t = t
         self.state = state0
-        if isinstance(generator, Wiener):
+        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)
@@ -74,8 +91,8 @@ def set_state(self, t, state0, generator):
                 (self.N_dw, num_collapse)
             )
         self.rhs._register_feedback(self.wiener)
-        opt = [self.options[key] for key in self._stepper_options]
-        self.step_func = self.stepper(self.rhs(self.options), *opt).run
+        rhs = self.rhs(self.options)
+        self.step_func = self.stepper(rhs, **stepper_opt).run
         self._is_set = True
 
     def get_state(self, copy=True):
@@ -228,9 +245,13 @@ class PlatenSODE(_Explicit_Simple_Integrator):
 
     - 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):
@@ -257,7 +278,7 @@ class PredCorr_SODE(_Explicit_Simple_Integrator):
     }
     stepper = _sode.PredCorr
     N_dw = 1
-    _stepper_options = ["alpha", "eta"]
+    _stepper_options = ["alpha", "eta", "measurement_noise"]
 
     @property
     def options(self):

qutip/solver/sode/ssystem.pxd

@@ -13,10 +13,13 @@ cdef class _StochasticSystem:
 
     cpdef list diffusion(self, t, Data state)
 
+    cpdef list expect(self, t, Data state)
+
     cpdef void set_state(self, double t, Dense state) except *
 
     cpdef Data a(self)
     cpdef Data bi(self, int i)
+    cpdef complex expect_i(self, int i)
     cpdef Data Libj(self, int i, int j)
     cpdef Data Lia(self, int i)
     cpdef Data L0bi(self, int i)