Merge pull request #2369 from pmenczel/weighted-trajectories

Weighted trajectories

doc/apidoc/classes.rst

@@ -231,20 +231,22 @@ Solver Options and Results
     :inherited-members:
     :exclude-members: add_processor, add
 
-.. autoclass:: qutip.solver.result.MultiTrajResult
+.. autoclass:: qutip.solver.multitrajresult.MultiTrajResult
     :members:
     :inherited-members:
     :exclude-members: add_processor, add, add_end_condition
 
-.. autoclass:: qutip.solver.result.McResult
+.. autoclass:: qutip.solver.result.TrajectoryResult
+    :show-inheritance:
     :members:
-    :inherited-members:
-    :exclude-members: add_processor, add, add_end_condition
 
-.. autoclass:: qutip.solver.result.NmmcResult
+.. autoclass:: qutip.solver.multitrajresult.McResult
+    :show-inheritance:
+    :members:
+
+.. autoclass:: qutip.solver.multitrajresult.NmmcResult
+    :show-inheritance:
     :members:
-    :inherited-members:
-    :exclude-members: add_processor, add, add_end_condition
 
 .. _classes-piqs:
 

doc/changes/2369.feature

@@ -0,0 +1 @@
+Weighted trajectories in trajectory solvers (enables improved sampling for nm_mcsolve)

doc/guide/dynamics/dynamics-intro.rst

@@ -45,14 +45,14 @@ quantum systems and indicates the type of object returned by the solver:
    * - Monte Carlo evolution
      - :func:`~qutip.solver.mcsolve.mcsolve`
      - :obj:`~qutip.solver.mcsolve.MCSolver`
-     - :obj:`~qutip.solver.result.McResult`
+     - :obj:`~qutip.solver.multitrajresult.McResult`
    * - Non-Markovian Monte Carlo
      - :func:`~qutip.solver.nm_mcsolve.nm_mcsolve`
      - :obj:`~qutip.solver.nm_mcsolve.NonMarkovianMCSolver`
-     - :obj:`~qutip.solver.result.NmmcResult`
+     - :obj:`~qutip.solver.multitrajresult.NmmcResult`
    * - Bloch-Redfield master equation
-     - :func:`~qutip.solver.mesolve.brmesolve`
-     - :obj:`~qutip.solver.mesolve.BRSolver`
+     - :func:`~qutip.solver.brmesolve.brmesolve`
+     - :obj:`~qutip.solver.brmesolve.BRSolver`
      - :obj:`~qutip.solver.result.Result`
    * - Floquet-Markov master equation
      - :func:`~qutip.solver.floquet.fmmesolve`
@@ -61,11 +61,11 @@ quantum systems and indicates the type of object returned by the solver:
    * - Stochastic Schrödinger equation
      - :func:`~qutip.solver.stochastic.ssesolve`
      - :obj:`~qutip.solver.stochastic.SSESolver`
-     - :obj:`~qutip.solver.result.MultiTrajResult`
+     - :obj:`~qutip.solver.multitrajresult.MultiTrajResult`
    * - Stochastic master equation
      - :func:`~qutip.solver.stochastic.smesolve`
      - :obj:`~qutip.solver.stochastic.SMESolver`
-     - :obj:`~qutip.solver.result.MultiTrajResult`
+     - :obj:`~qutip.solver.multitrajresult.MultiTrajResult`
    * - Transfer Tensor Method time-evolution
      - :func:`~qutip.solver.nonmarkov.transfertensor.ttmsolve`
      - None

doc/guide/dynamics/dynamics-nmmonte.rst

@@ -79,6 +79,9 @@ associated jump rates :math:`\Gamma_n(t)\geq0` appropriate for simulation.
 We conclude with a simple example demonstrating the usage of the ``nm_mcsolve``
 function. For more elaborate, physically motivated examples, we refer to the
 `accompanying tutorial notebook <https://github.com/qutip/qutip-tutorials/blob/main/tutorials-v5/time-evolution/013_nonmarkovian_monte_carlo.md>`_.
+Note that the example also demonstrates the usage of the ``improved_sampling``
+option (which is explained in the guide for the
+:ref:`Monte Carlo Solver<monte>`) in ``nm_mcsolve``.
 
 
 .. plot::
@@ -98,10 +101,11 @@ function. For more elaborate, physically motivated examples, we refer to the
     ops_and_rates = []
     ops_and_rates.append([a0.dag(), gamma1])
     ops_and_rates.append([a0,       gamma2])
+    nm_options = {'map': 'parallel', 'improved_sampling': True}
     MCSol = nm_mcsolve(H, psi0, times, ops_and_rates,
                        args={'kappa': 1.0 / 0.129, 'nth': 0.063},
                        e_ops=[a0.dag() * a0, a0 * a0.dag()],
-                       options={'map': 'parallel'}, ntraj=2500)
+                       options=nm_options, ntraj=2500)
 
     # mesolve integration for comparison
     d_ops = [[lindblad_dissipator(a0.dag(), a0.dag()), gamma1],

qutip/solver/__init__.py

@@ -1,4 +1,5 @@
 from .result import *
+from .multitrajresult import *
 from .options import *
 import qutip.solver.integrator as integrator
 from .integrator import IntegratorException

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 .multitrajresult 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,5 @@
-from .result import Result, MultiTrajResult
+from .result import TrajectoryResult
+from .multitrajresult import MultiTrajResult
 from .parallel import _get_map
 from time import time
 from .solver_base import Solver
@@ -50,7 +51,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 +149,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 +212,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,