Improved sampling algorithm for mcsolve

doc/biblio.rst

@@ -73,4 +73,7 @@ Bibliography
     N Khaneja et. al. *Optimal control of coupled spin dynamics: Design of NMR pulse sequences by gradient ascent algorithms.* J. Magn. Reson. **172**, 296–305 (2005). :doi:`10.1016/j.jmr.2004.11.004`
 
 .. [Donvil22]
-    B. Donvil, P. Muratore-Ginanneschi, *Quantum trajectory framework for general time-local master equations*, Nat Commun **13**, 4140 (2022). :doi:`10.1038/s41467-022-31533-8`.
+    B. Donvil, P. Muratore-Ginanneschi, *Quantum trajectory framework for general time-local master equations*, Nat Commun **13**, 4140 (2022). :doi:`10.1038/s41467-022-31533-8`.
+
+.. [Abd19]
+    M. Abdelhafez, D. I. Schuster, J. Koch, *Gradient-based optimal control of open quantum systems using quantumtrajectories and automatic differentiation*, Phys. Rev. A **99**, 052327 (2019). :doi:`10.1103/PhysRevA.99.052327`.

doc/changes/2218.feature

@@ -0,0 +1 @@
+Improved sampling algorithm for mcsolve

doc/guide/dynamics/dynamics-monte.rst

@@ -148,6 +148,29 @@ Now, the Monte Carlo solver will calculate expectation values for both operators
 
 
 
+Using the Improved Sampling Algorithm
+-------------------------------------
+
+Oftentimes, quantum jumps are rare. This is especially true in the context of simulating gates
+for quantum information purposes, where typical gate times are orders of magnitude smaller than
+typical timescales for decoherence. In this case, using the standard monte-carlo sampling algorithm,
+we often repeatedly sample the no-jump trajectory. We can thus reduce the number of required runs
+by only sampling the no-jump trajectory once. We then extract the no-jump probability :math:`p`,
+and for all future runs we only sample random numbers :math:`r_1` where :math:`r_1>p`, thus ensuring
+that a jump will occur. When it comes time to compute expectation values, we weight the no-jump
+trajectory by :math:`p` and the jump trajectories by :math:`1-p`. This algorithm is described
+in [Abd19]_. This algorithm can be utilized by setting the flag ``improved_sampling=True`` in the call to
+``mcsolve``:
+
+.. plot::
+    :context: close-figs
+
+    data = mcsolve(H, psi0, times, [np.sqrt(0.1) * a], e_ops=[a.dag() * a, sm.dag() * sm], improved_sampling=True)
+
+where in this case the first run samples the no-jump trajectory, and the remaining 499 trajectories are all
+guaranteed to include (at least) one jump.
+
+
 .. _monte-reuse:
 
 Reusing Hamiltonian Data

qutip/solver/mcsolve.py

@@ -1,17 +1,18 @@
-__all__ = ['mcsolve', "MCSolver"]
+__all__ = ['mcsolve', "MCSolver", "MCSolverImprovedSampling"]
 
 import numpy as np
 from ..core import QobjEvo, spre, spost, Qobj, unstack_columns
-from .multitraj import MultiTrajSolver
+from .multitraj import MultiTrajSolver, MultiTrajSolverImprovedSampling
 from .solver_base import Solver, Integrator
-from .result import McResult, McTrajectoryResult
+from .result import McResult, McTrajectoryResult, McResultImprovedSampling
 from .mesolve import mesolve, MESolver
 import qutip.core.data as _data
 from time import time
 
 
 def mcsolve(H, state, tlist, c_ops=(), e_ops=None, ntraj=500, *,
-            args=None, options=None, seeds=None, target_tol=None, timeout=None):
+            args=None, options=None, seeds=None, target_tol=None, timeout=None,
+            improved_sampling=False):
     r"""
     Monte Carlo evolution of a state vector :math:`|\psi \rangle` for a
     given Hamiltonian and sets of collapse operators. Options for the
@@ -116,6 +117,10 @@ def mcsolve(H, state, tlist, c_ops=(), e_ops=None, ntraj=500, *,
         Maximum time for the evolution in second. When reached, no more
         trajectories will be computed.
 
+    improved_sampling: Bool
+        Whether to use the improved sampling algorithm from Abdelhafez et al.
+        PRA (2019)
+
     Returns
     -------
     results : :class:`qutip.solver.McResult`
@@ -148,8 +153,11 @@ def mcsolve(H, state, tlist, c_ops=(), e_ops=None, ntraj=500, *,
             "ntraj must be an integer. "
             "A list of numbers is not longer supported."
         )
+    if not improved_sampling:
+        mc = MCSolver(H, c_ops, options=options)
+    else:
+        mc = MCSolverImprovedSampling(H, c_ops, options=options)
 
-    mc = MCSolver(H, c_ops, options=options)
     result = mc.run(state, tlist=tlist, ntraj=ntraj, e_ops=e_ops,
                     seed=seeds, target_tol=target_tol, timeout=timeout)
     return result
@@ -171,7 +179,8 @@ def __init__(self, integrator, c_ops, n_ops, options=None):
         self._is_set = False
         self.issuper = c_ops[0].issuper
 
-    def set_state(self, t, state0, generator):
+    def set_state(self, t, state0, generator,
+                  no_jump=False, jump_prob_floor=0.0):
         """
         Set the state of the ODE solver.
 
@@ -185,10 +194,25 @@ def set_state(self, t, state0, generator):
 
         generator : numpy.random.generator
             Random number generator.
+
+        no_jump: Bool
+            whether or not to sample the no-jump trajectory.
+            If so, the "random number" should be set to zero
+
+        jump_prob_floor: float
+            if no_jump == False, this is set to the no-jump
+            probability. This setting ensures that we sample
+            a trajectory with jumps
         """
         self.collapses = []
         self._generator = generator
-        self.target_norm = self._generator.random()
+        if no_jump:
+            self.target_norm = 0.0
+        else:
+            self.target_norm = (
+                    self._generator.random() * (1 - jump_prob_floor)
+                    + jump_prob_floor
+            )
         self._integrator.set_state(t, state0)
         self._is_set = True
 
@@ -298,6 +322,9 @@ def _do_collapse(self, collapse_time, state):
         else:
             state_new = _data.mul(state_new, 1 / new_norm)
             self.collapses.append((collapse_time, which))
+            # this does not need to be modified for improved sampling:
+            # as noted in Abdelhafez PRA (2019),
+            # after a jump we reset to the full range [0, 1)
             self.target_norm = self._generator.random()
         self._integrator.set_state(collapse_time, state_new)
 
@@ -417,15 +444,19 @@ def _argument(self, args):
         for n_op in self._n_ops:
             n_op.arguments(args)
 
-    def _run_one_traj(self, seed, state, tlist, e_ops):
+    def _run_one_traj(self, seed, state, tlist, e_ops,
+                      no_jump=False, jump_prob_floor=0.0):
         """
         Run one trajectory and return the result.
         """
         # The integrators is reused, but non-reentrant. They are are fine for
         # multiprocessing, but will fail with multithreading.
         # If a thread base parallel map is created, eahc trajectory should use
         # a copy of the integrator.
-        seed, result = super()._run_one_traj(seed, state, tlist, e_ops)
+        seed, result = super()._run_one_traj(
+            seed, state, tlist, e_ops,
+            no_jump=no_jump, jump_prob_floor=jump_prob_floor
+        )
         result.collapse = self._integrator.collapses
         return seed, result
 
@@ -518,3 +549,18 @@ def avail_integrators(cls):
             **Solver.avail_integrators(),
             **cls._avail_integrators,
         }
+
+
+class MCSolverImprovedSampling(MultiTrajSolverImprovedSampling, MCSolver):
+    r"""
+    Monte Carlo Solver of a state vector :math:`|\psi \rangle` for a
+    given Hamiltonian and sets of collapse operators using the improved
+    sampling algorithm. See MCSolver for further details and documentation
+    """
+    name = "mcsolve improved sampling"
+    resultclass = McResultImprovedSampling
+    trajectory_resultclass = McTrajectoryResult
+    mc_integrator_class = MCIntegrator
+
+    def __init__(self, H, c_ops, *, options=None):
+        MCSolver.__init__(self, H, c_ops, options=options)

qutip/solver/multitraj.py

@@ -1,10 +1,10 @@
-from .result import Result, MultiTrajResult
+from .result import Result, MultiTrajResult, MultiTrajResultImprovedSampling
 from .parallel import _get_map
 from time import time
 from .solver_base import Solver
 import numpy as np
 
-__all__ = ["MultiTrajSolver"]
+__all__ = ["MultiTrajSolver", "MultiTrajSolverImprovedSampling"]
 
 
 class MultiTrajSolver(Solver):
@@ -102,6 +102,28 @@ def step(self, t, *, args=None, copy=True):
         _, state = self._integrator.integrate(t, copy=False)
         return self._restore_state(state, copy=copy)
 
+    def _initialize_run(self, state, ntraj=1, args=None, e_ops=(),
+                        timeout=None, target_tol=None, seed=None):
+        start_time = time()
+        self._argument(args)
+        stats = self._initialize_stats()
+        seeds = self._read_seed(seed, ntraj)
+
+        result = self.resultclass(
+            e_ops, self.options, solver=self.name, stats=stats
+        )
+        result.add_end_condition(ntraj, target_tol)
+
+        map_func = _get_map[self.options['map']]
+        map_kw = {
+            'timeout': timeout,
+            'job_timeout': self.options['job_timeout'],
+            'num_cpus': self.options['num_cpus'],
+        }
+        state0 = self._prepare_state(state)
+        stats['preparation time'] += time() - start_time
+        return stats, seeds, result, map_func, map_kw, state0
+
     def run(self, state, tlist, ntraj=1, *,
             args=None, e_ops=(), timeout=None, target_tol=None, seed=None):
         """
@@ -163,25 +185,15 @@ def run(self, state, tlist, ntraj=1, *,
             The simulation will end when the first end condition is reached
             between ``ntraj``, ``timeout`` and ``target_tol``.
         """
-        start_time = time()
-        self._argument(args)
-        stats = self._initialize_stats()
-        seeds = self._read_seed(seed, ntraj)
-
-        result = self.resultclass(
-            e_ops, self.options, solver=self.name, stats=stats
+        stats, seeds, result, map_func, map_kw, state0 = self._initialize_run(
+            state,
+            ntraj,
+            args=args,
+            e_ops=e_ops,
+            timeout=timeout,
+            target_tol=target_tol,
+            seed=seed,
         )
-        result.add_end_condition(ntraj, target_tol)
-
-        map_func = _get_map[self.options['map']]
-        map_kw = {
-            'timeout': timeout,
-            'job_timeout': self.options['job_timeout'],
-            'num_cpus': self.options['num_cpus'],
-        }
-        state0 = self._prepare_state(state)
-        stats['preparation time'] += time() - start_time
-
         start_time = time()
         map_func(
             self._run_one_traj, seeds,
@@ -193,13 +205,16 @@ def run(self, state, tlist, ntraj=1, *,
         result.stats['run time'] = time() - start_time
         return result
 
-    def _run_one_traj(self, seed, state, tlist, e_ops):
+    def _run_one_traj(self, seed, state, tlist, e_ops,
+                      no_jump=False, jump_prob_floor=0.0):
         """
         Run one trajectory and return the result.
         """
         result = self.trajectory_resultclass(e_ops, self.options)
         generator = self._get_generator(seed)
-        self._integrator.set_state(tlist[0], state, generator)
+        self._integrator.set_state(tlist[0], state, generator,
+                                   no_jump=no_jump,
+                                   jump_prob_floor=jump_prob_floor)
         result.add(tlist[0], self._restore_state(state, copy=False))
         for t in tlist[1:]:
             t, state = self._integrator.integrate(t, copy=False)
@@ -254,3 +269,58 @@ def _get_generator(self, seed):
     @classmethod
     def avail_integrators(cls):
         return cls._avail_integrators
+
+
+class MultiTrajSolverImprovedSampling(MultiTrajSolver):
+    """
+    Class for multi-trajectory evolutions using the improved sampling
+    algorithm. See docstring for MultiTrajSolver for further documentation
+    """
+    name = "multi trajectory efficient"
+    resultclass = MultiTrajResultImprovedSampling
+    trajectory_resultclass = Result
+
+    def __init__(self, rhs, *, options=None):
+        super().__init__(rhs, options=options)
+
+    def run(self, state, tlist, ntraj=1, *,
+            args=None, e_ops=(), timeout=None, target_tol=None, seed=None):
+        """
+        Do the evolution of the Quantum system.
+        See the overridden method for further details. The modification
+        here is to sample the no-jump trajectory first. Then, the no-jump
+        probability is used as a lower-bound for random numbers in future
+        monte carlo runs
+        """
+        stats, seeds, result, map_func, map_kw, state0 = self._initialize_run(
+            state,
+            ntraj,
+            args=args,
+            e_ops=e_ops,
+            timeout=timeout,
+            target_tol=target_tol,
+            seed=seed,
+        )
+        # first run the no-jump trajectory
+        start_time = time()
+        seed0, no_jump_result = self._run_one_traj(seeds[0], state0, tlist,
+                                                   e_ops, no_jump=True)
+        _, state, _ = self._integrator.get_state(copy=False)
+        no_jump_prob = self._integrator._prob_func(state)
+        result.no_jump_prob = no_jump_prob
+        result.add((seed0, no_jump_result))
+        result.stats['no jump run time'] = time() - start_time
+
+        # run the remaining trajectories with the random number floor
+        # set to the no jump probability such that we only sample
+        # trajectories with jumps
+        start_time = time()
+        map_func(
+            self._run_one_traj, seeds[1:],
+            (state0, tlist, e_ops, False, no_jump_prob),
+            reduce_func=result.add, map_kw=map_kw,
+            progress_bar=self.options["progress_bar"],
+            progress_bar_kwargs=self.options["progress_kwargs"]
+        )
+        result.stats['run time'] = time() - start_time
+        return result