Weighted trajectories

doc/apidoc/classes.rst

@@ -236,15 +236,17 @@ Solver Options and Results
     :inherited-members:
     :exclude-members: add_processor, add, add_end_condition
 
+.. autoclass:: qutip.solver.result.TrajectoryResult
+    :show-inheritance:
+    :members:
+
 .. autoclass:: qutip.solver.result.McResult
+    :show-inheritance:
     :members:
-    :inherited-members:
-    :exclude-members: add_processor, add, add_end_condition
 
 .. autoclass:: qutip.solver.result.NmmcResult
+    :show-inheritance:
     :members:
-    :inherited-members:
-    :exclude-members: add_processor, add, add_end_condition
 
 .. _classes-piqs:
 

qutip/solver/mcsolve.py

@@ -4,7 +4,7 @@
 from ..core import QobjEvo, spre, spost, Qobj, unstack_columns
 from .multitraj import MultiTrajSolver, _MultiTrajRHS
 from .solver_base import Solver, Integrator, _solver_deprecation
-from .result import McResult, McTrajectoryResult, McResultImprovedSampling
+from .result import McResult
 from .mesolve import mesolve, MESolver
 from ._feedback import _QobjFeedback, _DataFeedback, _CollapseFeedback
 import qutip.core.data as _data
@@ -401,7 +401,7 @@ class MCSolver(MultiTrajSolver):
         Options for the evolution.
     """
     name = "mcsolve"
-    _trajectory_resultclass = McTrajectoryResult
+    _resultclass = McResult
     _mc_integrator_class = MCIntegrator
     solver_options = {
         "progress_bar": "text",
@@ -480,6 +480,9 @@ def _run_one_traj(self, seed, state, tlist, e_ops, **integrator_kwargs):
         """
         seed, result = super()._run_one_traj(seed, state, tlist, e_ops,
                                              **integrator_kwargs)
+        jump_prob_floor = integrator_kwargs.get('jump_prob_floor', 0)
+        if jump_prob_floor > 0:
+            result.add_relative_weight(1 - jump_prob_floor)
         result.collapse = self._integrator.collapses
         return seed, result
 
@@ -488,26 +491,23 @@ def run(self, state, tlist, ntraj=1, *,
         # Overridden 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
-        if not self.options.get("improved_sampling", False):
+        if not self.options["improved_sampling"]:
             return super().run(state, tlist, ntraj=ntraj, args=args,
                                e_ops=e_ops, timeout=timeout,
                                target_tol=target_tol, seeds=seeds)
-        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,
-            seeds=seeds,
+
+        seeds, result, map_func, map_kw, state0 = self._initialize_run(
+            state, ntraj, args=args, e_ops=e_ops,
+            timeout=timeout, target_tol=target_tol, seeds=seeds
         )
+
         # 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
+        no_jump_result.add_absolute_weight(no_jump_prob)
         result.add((seed0, no_jump_result))
         result.stats['no jump run time'] = time() - start_time
 
@@ -542,13 +542,6 @@ def _get_integrator(self):
         self._init_integrator_time = time() - _time_start
         return mc_integrator
 
-    @property
-    def _resultclass(self):
-        if self.options.get("improved_sampling", False):
-            return McResultImprovedSampling
-        else:
-            return McResult
-
     @property
     def options(self):
         """

qutip/solver/multitraj.py

@@ -1,4 +1,4 @@
-from .result import Result, MultiTrajResult
+from .result import TrajectoryResult, MultiTrajResult
 from .parallel import _get_map
 from time import time
 from .solver_base import Solver
@@ -50,7 +50,7 @@ class MultiTrajSolver(Solver):
     """
     name = "generic multi trajectory"
     _resultclass = MultiTrajResult
-    _trajectory_resultclass = Result
+    _trajectory_resultclass = TrajectoryResult
     _avail_integrators = {}
 
     # Class of option used by the solver
@@ -148,7 +148,7 @@ def _initialize_run(self, state, ntraj=1, args=None, e_ops=(),
         })
         state0 = self._prepare_state(state)
         stats['preparation time'] += time() - start_time
-        return stats, seeds, result, map_func, map_kw, state0
+        return seeds, result, map_func, map_kw, state0
 
     def run(self, state, tlist, ntraj=1, *,
             args=None, e_ops=(), timeout=None, target_tol=None, seeds=None):
@@ -211,7 +211,7 @@ def run(self, state, tlist, ntraj=1, *,
             The simulation will end when the first end condition is reached
             between ``ntraj``, ``timeout`` and ``target_tol``.
         """
-        stats, seeds, result, map_func, map_kw, state0 = self._initialize_run(
+        seeds, result, map_func, map_kw, state0 = self._initialize_run(
             state,
             ntraj,
             args=args,

qutip/solver/nm_mcsolve.py

@@ -1,15 +1,14 @@
 __all__ = ['nm_mcsolve', 'NonMarkovianMCSolver']
 
-import functools
 import numbers
 
 import numpy as np
 import scipy
 
+from .result import NmmcResult
 from .multitraj import MultiTrajSolver
 from .mcsolve import MCSolver, MCIntegrator
 from .mesolve import MESolver, mesolve
-from .result import NmmcResult, NmmcTrajectoryResult
 from .cy.nm_mcsolve import RateShiftCoefficient, SqrtRealCoefficient
 from ..core.coefficient import ConstantCoefficient
 from ..core import (
@@ -114,6 +113,9 @@ def nm_mcsolve(H, state, tlist, ops_and_rates=(), e_ops=None, ntraj=500, *,
             ``norm_tol`` are the tolerance in time and norm respectively.
             An error will be raised if the collapse could not be found within
             ``norm_steps`` tries.
+        - | improved_sampling : Bool
+          | Whether to use the improved sampling algorithm from Abdelhafez et
+            al. PRA (2019)
         - | mc_corr_eps : float
           | Small number used to detect non-physical collapse caused by
             numerical imprecision.
@@ -217,7 +219,10 @@ def initialize(self, t0, cache='clear'):
         #  to pre-compute the continuous contribution to the martingale
         self._t_prev = t0
         self._continuous_martingale_at_t_prev = 1
-        self._discrete_martingale = 1
+
+        # _discrete_martingale is a list of (time, factor) such that
+        # mu_d(t) is the product of all factors with time < t
+        self._discrete_martingale = []
 
         if np.array_equal(cache, 'clear'):
             self._precomputed_continuous_martingale = {}
@@ -239,7 +244,7 @@ def add_collapse(self, collapse_time, collapse_channel):
         rate = self._nm_solver.rate(collapse_time, collapse_channel)
         shift = self._nm_solver.rate_shift(collapse_time)
         factor = rate / (rate + shift)
-        self._discrete_martingale *= factor
+        self._discrete_martingale.append((collapse_time, factor))
 
     def value(self, t):
         if self._t_prev is None:
@@ -256,7 +261,13 @@ def value(self, t):
         self._t_prev = t
         self._continuous_martingale_at_t_prev = mu_c
 
-        return self._discrete_martingale * mu_c
+        # find value of discrete martingale at given time
+        mu_d = 1
+        for time, factor in self._discrete_martingale:
+            if t > time:
+                mu_d *= factor
+
+        return mu_d * mu_c
 
     def _compute_continuous_martingale(self, t1, t2):
         if t1 == t2:
@@ -339,15 +350,12 @@ class NonMarkovianMCSolver(MCSolver):
         "norm_steps": 5,
         "norm_t_tol": 1e-6,
         "norm_tol": 1e-4,
+        "improved_sampling": False,
         "completeness_rtol": 1e-5,
         "completeness_atol": 1e-8,
         "martingale_quad_limit": 100,
     }
 
-    # both classes will be partially initialized in constructor
-    _trajectory_resultclass = NmmcTrajectoryResult
-    _mc_integrator_class = NmMCIntegrator
-
     def __init__(
         self, H, ops_and_rates, args=None, options=None,
     ):
@@ -380,14 +388,11 @@ def __init__(
             for op, sqrt_shifted_rate
             in zip(self.ops, self._sqrt_shifted_rates)
         ]
-        self._trajectory_resultclass = functools.partial(
-            NmmcTrajectoryResult, __nm_solver=self,
-        )
-        self._mc_integrator_class = functools.partial(
-            NmMCIntegrator, __martingale=self._martingale,
-        )
         super().__init__(H, c_ops, options=options)
 
+    def _mc_integrator_class(self, *args):
+        return NmMCIntegrator(*args, __martingale=self._martingale)
+
     def _check_completeness(self, ops_and_rates):
         """
         Checks whether ``sum(Li.dag() * Li)`` is proportional to the identity
@@ -509,9 +514,8 @@ def sqrt_shifted_rate(self, t, i):
     # the run.
     #
     # Regarding (b), in the start/step-interface we just include the martingale
-    # in the step method. In order to include the martingale in the
-    # run-interface, we use a custom trajectory-resultclass that grabs the
-    # martingale value from the NonMarkovianMCSolver whenever a state is added.
+    # in the step method. In the run-interface, the martingale is added as a
+    # relative weight to the trajectory result at the end of `_run_one_traj`.
 
     def start(self, state, t0, seed=None):
         self._martingale.initialize(t0, cache='clear')
@@ -524,6 +528,17 @@ def step(self, t, *, args=None, copy=True):
             state = ket2dm(state)
         return state * self.current_martingale()
 
+    def _run_one_traj(self, seed, state, tlist, e_ops, **integrator_kwargs):
+        """
+        Run one trajectory and return the result.
+        """
+        seed, result = super()._run_one_traj(seed, state, tlist, e_ops,
+                                             **integrator_kwargs)
+        martingales = [self._martingale.value(t) for t in tlist]
+        result.add_relative_weight(martingales)
+        result.trace = martingales
+        return seed, result
+
     def run(self, state, tlist, ntraj=1, *, args=None, **kwargs):
         # update `args` dictionary before precomputing martingale
         self._argument(args)
@@ -596,6 +611,10 @@ def options(self):
         norm_steps: int, default: 5
             Maximum number of tries to find the collapse.
 
+        improved_sampling: Bool, default: False
+            Whether to use the improved sampling algorithm
+            of Abdelhafez et al. PRA (2019)
+
         completeness_rtol: float, default: 1e-5
             Used in determining whether the given Lindblad operators satisfy
             a certain completeness relation. If they do not, an additional