Repository: qutip/qutip

Issue:
MemoryError when updating to latest qutip
### Bug Description

Might just be me mis-using the new 5.X branch of Qutip. I get a memory error when upgrading to the new branch.

It appears to fail when performing self._to_dm(state) for each state. My states are vectors of length 800 - so it runs out of memory. Is it possible to request the solver only stores the wavefunction at each timestep? I want the wavefunction for each trajectory for each timestep. Alternatively one could introduce a new "proj" data storage to make this more efficient.

### Code to Reproduce the Bug

```shell
result = qutip.ssesolve(
        hamiltonian_qobj,
        initial_state_qobj,
        times.times,
        sc_ops=sc_ops,
        e_ops=[],
        options={
            "method": "euler",
            "progress_bar": "enhanced",
            "store_states": True,
            "keep_runs_results": True,
            "improved_sampling": False,
            "dt": times.delta_t / times.fundamental_n,
        },
        ntraj=n_trajectories,  # cspell:disable-line
    )
```


### Code Output

```shell
File "...", line 270, in solve_stochastic_schrodinger_equation
    result = qutip.ssesolve(
             ^^^^^^^^^^^^^^^
  File "...\.venv\Lib\site-packages\qutip\solver\stochastic.py", line 461, in ssesolve
    return sol.run(
           ^^^^^^^^
  File "...\.venv\Lib\site-packages\qutip\solver\multitraj.py", line 198, in run
    map_func(
  File "...\.venv\Lib\site-packages\qutip\solver\parallel.py", line 122, in serial_map
    remaining_ntraj = reduce_func(result)
                      ^^^^^^^^^^^^^^^^^^^
  File "...\.venv\Lib\site-packages\qutip\solver\result.py", line 645, in add   
    op(trajectory)
  File "...\.venv\Lib\site-packages\qutip\solver\result.py", line 519, in _reduce_states
    accu + self._to_dm(state)
    ~~~~~^~~~~~~~~~~~~~~~~~~~
  File "...\.venv\Lib\site-packages\qutip\core\qobj.py", line 134, in out
    return method(self, other)
           ^^^^^^^^^^^^^^^^^^^
  File "...\.venv\Lib\site-packages\qutip\core\qobj.py", line 408, in __add__   
    return Qobj(_data.add(self._data, other._data),
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "qutip\\core\\data\\dispatch.pyx", line 392, in qutip.core.data.dispatch.Dispatcher.__call__
  File "qutip\\core\\data\\add.pyx", line 118, in qutip.core.data.add.add_csr
  File "qutip\\core\\data\\add.pyx", line 154, in qutip.core.data.add.add_csr
  File "qutip\\core\\data\\csr.pyx", line 159, in qutip.core.data.csr.CSR.copy
  File "qutip\\core\\data\\csr.pyx", line 542, in qutip.core.data.csr.empty_like
  File "qutip\\core\\data\\csr.pyx", line 533, in qutip.core.data.csr.empty
MemoryError
```


### Expected Behaviour

Code should not have to store the full density matrix, and should not lead to a memory error

### Your Environment

```shell
QuTiP Version:      5.0.0a2
Numpy Version:      1.26.3
Scipy Version:      1.12.0
Cython Version:     None
Matplotlib Version: 3.8.2
Python Version:     3.12.1
Number of CPUs:     4
BLAS Info:          Generic
INTEL MKL Ext:      False
Platform Info:      Windows (AMD64)
Installation path:  ...
```


### Additional Context

_No response_


Discussion:
About only storing kets, but not averaging into a dm, we can't do that yet, but it's easy to implement so if useful, we can add the option.

The error is not in the `_to_dm`, but in the addition, which is done in the sparse CSR format. No dense conversion are made. But the `proj` operator does not have a tidyup step, so any non-zero are kept, however small they are. This is also an easy fix that could help.

The `MemoryError` is also not that helpful, we should improve the error message to tell what size of arrays are allocated.

Thank you for reporting these and trying the alpha version.

I see you don't have cython, do you have a way to compile the source yourself? 
Once fixed, I would be curious to have you try the development version without waiting for a release.




Thanks for the quick response - currently I've just pip-installed the package, but I would be happy to use the main branch if you would be able to fix this. I would find this very useful - for the quantities I am interested in calculating I need much less than 800 repeats to get convergence. I understand it is possible to calculate the expectation directly in ssesolve but I find it is useful to save the states to a file so I have more freedom to calculate a range of different operators at a later date.
Saving to a file without averaging is possible.

`e_ops` accept any function `f(t: float, state: Qobj) -> number`. That function can save the state to a file and return a dummy value. However, it does not know which trajectories it belong to so it would only work with `serial_map`.

_Expectation functions_ are always called, it would remove the need to keep trajectories and store states.

`Qobj.data_as("CSR")` will extract the state as a scipy CSR matrix (if already in `csr`, does not do the conversion).






Yeah I'm not really interested in using e_ops this way - it feels like a hack. I really just want the states of each trajectory so I can do what I want, and it seems silly I have to also pay the price of storing the density matrix. Wouldn't there be significant overhead calling to python each step and opening/writing to a file 
Yes, it's a hack, but it's probably the best option you have until we fix it. So if you are in an hurry for your research, just know the method exist. 

Resolving patch:
diff --git a/qutip/solver/heom/bofin_solvers.py b/qutip/solver/heom/bofin_solvers.py
--- a/qutip/solver/heom/bofin_solvers.py
+++ b/qutip/solver/heom/bofin_solvers.py
@@ -385,7 +385,7 @@ def _post_init(self):
 
         self.store_ados = self.options["store_ados"]
         if self.store_ados:
-            self.final_ado_state = None
+            self._final_ado_state = None
             self.ado_states = []
 
     def _e_op_func(self, e_op):
@@ -407,9 +407,17 @@ def _store_state(self, t, ado_state):
             self.ado_states.append(ado_state)
 
     def _store_final_state(self, t, ado_state):
-        self.final_state = ado_state.rho
+        self._final_state = ado_state.rho
         if self.store_ados:
-            self.final_ado_state = ado_state
+            self._final_ado_state = ado_state
+
+    @property
+    def final_ado_state(self):
+        if self._final_ado_state is not None:
+            return self._final_state
+        if self.ado_states:
+            return self.ado_states[-1]
+        return None
 
 
 def heomsolve(
diff --git a/qutip/solver/mcsolve.py b/qutip/solver/mcsolve.py
--- a/qutip/solver/mcsolve.py
+++ b/qutip/solver/mcsolve.py
@@ -167,6 +167,7 @@ class _MCSystem(_MTSystem):
     """
     Container for the operators of the solver.
     """
+
     def __init__(self, rhs, c_ops, n_ops):
         self.rhs = rhs
         self.c_ops = c_ops
@@ -247,8 +248,8 @@ def set_state(self, t, state0, generator,
             self.target_norm = 0.0
         else:
             self.target_norm = (
-                    self._generator.random() * (1 - jump_prob_floor)
-                    + jump_prob_floor
+                self._generator.random() * (1 - jump_prob_floor)
+                + jump_prob_floor
             )
         self._integrator.set_state(t, state0)
         self._is_set = True
diff --git a/qutip/solver/result.py b/qutip/solver/result.py
--- a/qutip/solver/result.py
+++ b/qutip/solver/result.py
@@ -1,9 +1,17 @@
 """ Class for solve function results"""
+
+from typing import TypedDict
 import numpy as np
 from ..core import Qobj, QobjEvo, expect, isket, ket2dm, qzero_like
 
-__all__ = ["Result", "MultiTrajResult", "McResult", "NmmcResult",
-           "McTrajectoryResult", "McResultImprovedSampling"]
+__all__ = [
+    "Result",
+    "MultiTrajResult",
+    "McResult",
+    "NmmcResult",
+    "McTrajectoryResult",
+    "McResultImprovedSampling",
+]
 
 
 class _QobjExpectEop:
@@ -16,6 +24,7 @@ class _QobjExpectEop:
     op : :obj:`.Qobj`
         The expectation value operator.
     """
+
     def __init__(self, op):
         self.op = op
 
@@ -46,6 +55,7 @@ class ExpectOp:
     op : object
         The original object used to define the e_op operation.
     """
+
     def __init__(self, op, f, append):
         self.op = op
         self._f = f
@@ -70,6 +80,7 @@ class _BaseResult:
     """
     Common method for all ``Result``.
     """
+
     def __init__(self, options, *, solver=None, stats=None):
         self.solver = solver
         if stats is None:
@@ -81,12 +92,12 @@ def __init__(self, options, *, solver=None, stats=None):
 
         # make sure not to store a reference to the solver
         options_copy = options.copy()
-        if hasattr(options_copy, '_feedback'):
+        if hasattr(options_copy, "_feedback"):
             options_copy._feedback = None
         self.options = options_copy
 
     def _e_ops_to_dict(self, e_ops):
-        """ Convert the supplied e_ops to a dictionary of Eop instances. """
+        """Convert the supplied e_ops to a dictionary of Eop instances."""
         if e_ops is None:
             e_ops = {}
         elif isinstance(e_ops, (list, tuple)):
@@ -118,6 +129,11 @@ def add_processor(self, f, requires_copy=False):
         self._state_processors_require_copy |= requires_copy
 
 
+class ResultOptions(TypedDict):
+    store_states: bool
+    store_final_state: bool
+
+
 class Result(_BaseResult):
     """
     Base class for storing solver results.
@@ -199,7 +215,18 @@ class Result(_BaseResult):
     options : dict
         The options for this result class.
     """
-    def __init__(self, e_ops, options, *, solver=None, stats=None, **kw):
+
+    options: ResultOptions
+
+    def __init__(
+        self,
+        e_ops,
+        options: ResultOptions,
+        *,
+        solver=None,
+        stats=None,
+        **kw,
+    ):
         super().__init__(options, solver=solver, stats=stats)
         raw_ops = self._e_ops_to_dict(e_ops)
         self.e_data = {k: [] for k in raw_ops}
@@ -211,7 +238,7 @@ def __init__(self, e_ops, options, *, solver=None, stats=None, **kw):
 
         self.times = []
         self.states = []
-        self.final_state = None
+        self._final_state = None
 
         self._post_init(**kw)
 
@@ -243,29 +270,29 @@ def _post_init(self):
         Sub-class ``.post_init()`` implementation may take additional keyword
         arguments if required.
         """
-        store_states = self.options['store_states']
-        store_final_state = self.options['store_final_state']
-
-        if store_states is None:
-            store_states = len(self.e_ops) == 0
+        store_states = self.options["store_states"]
+        store_states = store_states or (
+            len(self.e_ops) == 0 and store_states is None
+        )
         if store_states:
             self.add_processor(self._store_state, requires_copy=True)
 
-        if store_states or store_final_state:
+        store_final_state = self.options["store_final_state"]
+        if store_final_state and not store_states:
             self.add_processor(self._store_final_state, requires_copy=True)
 
     def _store_state(self, t, state):
-        """ Processor that stores a state in ``.states``. """
+        """Processor that stores a state in ``.states``."""
         self.states.append(state)
 
     def _store_final_state(self, t, state):
-        """ Processor that writes the state to ``.final_state``. """
-        self.final_state = state
+        """Processor that writes the state to ``._final_state``."""
+        self._final_state = state
 
     def _pre_copy(self, state):
-        """ Return a copy of the state. Sub-classes may override this to
-            copy a state in different manner or to skip making a copy
-            altogether if a copy is not necessary.
+        """Return a copy of the state. Sub-classes may override this to
+        copy a state in different manner or to skip making a copy
+        altogether if a copy is not necessary.
         """
         return state.copy()
 
@@ -311,10 +338,7 @@ def __repr__(self):
         ]
         if self.stats:
             lines.append("  Solver stats:")
-            lines.extend(
-                f"    {k}: {v!r}"
-                for k, v in self.stats.items()
-            )
+            lines.extend(f"    {k}: {v!r}" for k, v in self.stats.items())
         if self.times:
             lines.append(
                 f"  Time interval: [{self.times[0]}, {self.times[-1]}]"
@@ -334,6 +358,20 @@ def __repr__(self):
     def expect(self):
         return [np.array(e_op) for e_op in self.e_data.values()]
 
+    @property
+    def final_state(self):
+        if self._final_state is not None:
+            return self._final_state
+        if self.states:
+            return self.states[-1]
+        return None
+
+
+class MultiTrajResultOptions(TypedDict):
+    store_states: bool
+    store_final_state: bool
+    keep_runs_results: bool
+
 
 class MultiTrajResult(_BaseResult):
     """
@@ -455,7 +493,18 @@ class MultiTrajResult(_BaseResult):
     options : :obj:`~SolverResultsOptions`
         The options for this result class.
     """
-    def __init__(self, e_ops, options, *, solver=None, stats=None, **kw):
+
+    options: MultiTrajResultOptions
+
+    def __init__(
+        self,
+        e_ops,
+        options: MultiTrajResultOptions,
+        *,
+        solver=None,
+        stats=None,
+        **kw,
+    ):
         super().__init__(options, solver=solver, stats=stats)
         self._raw_ops = self._e_ops_to_dict(e_ops)
 
@@ -471,15 +520,29 @@ def __init__(self, e_ops, options, *, solver=None, stats=None, **kw):
         self._target_tols = None
 
         self.average_e_data = {}
-        self.e_data = {}
         self.std_e_data = {}
         self.runs_e_data = {}
 
         self._post_init(**kw)
 
+    @property
+    def _store_average_density_matricies(self) -> bool:
+        return (
+            self.options["store_states"]
+            or (self.options["store_states"] is None and self._raw_ops == {})
+        ) and not self.options["keep_runs_results"]
+
+    @property
+    def _store_final_density_matrix(self) -> bool:
+        return (
+            self.options["store_final_state"]
+            and not self._store_average_density_matricies
+            and not self.options["keep_runs_results"]
+        )
+
     @staticmethod
     def _to_dm(state):
-        if state.type == 'ket':
+        if state.type == "ket":
             state = state.proj()
         return state
 
@@ -489,10 +552,12 @@ def _add_first_traj(self, trajectory):
         """
         self.times = trajectory.times
 
-        if trajectory.states:
-            self._sum_states = [qzero_like(self._to_dm(state))
-                                for state in trajectory.states]
-        if trajectory.final_state:
+        if trajectory.states and self._store_average_density_matricies:
+            self._sum_states = [
+                qzero_like(self._to_dm(state)) for state in trajectory.states
+            ]
+
+        if trajectory.final_state and self._store_final_density_matrix:
             state = trajectory.final_state
             self._sum_final_states = qzero_like(self._to_dm(state))
 
@@ -507,13 +572,10 @@ def _add_first_traj(self, trajectory):
 
         self.average_e_data = {
             k: list(avg_expect)
-            for k, avg_expect
-            in zip(self._raw_ops, self._sum_expect)
+            for k, avg_expect in zip(self._raw_ops, self._sum_expect)
         }
-        self.e_data = self.average_e_data
-        if self.options['keep_runs_results']:
+        if self.options["keep_runs_results"]:
             self.runs_e_data = {k: [] for k in self._raw_ops}
-            self.e_data = self.runs_e_data
 
     def _store_trajectory(self, trajectory):
         self.trajectories.append(trajectory)
@@ -521,8 +583,7 @@ def _store_trajectory(self, trajectory):
     def _reduce_states(self, trajectory):
         self._sum_states = [
             accu + self._to_dm(state)
-            for accu, state
-            in zip(self._sum_states, trajectory.states)
+            for accu, state in zip(self._sum_states, trajectory.states)
         ]
 
     def _reduce_final_state(self, trajectory):
@@ -569,7 +630,7 @@ def _fixed_end(self):
         """
         ntraj_left = self._target_ntraj - self.num_trajectories
         if ntraj_left == 0:
-            self.stats['end_condition'] = 'ntraj reached'
+            self.stats["end_condition"] = "ntraj reached"
         return ntraj_left
 
     def _average_computer(self):
@@ -586,40 +647,36 @@ def _target_tolerance_end(self):
         if self.num_trajectories <= 1:
             return np.inf
         avg, avg2 = self._average_computer()
-        target = np.array([
-            atol + rtol * mean
-            for mean, (atol, rtol)
-            in zip(avg, self._target_tols)
-        ])
+        target = np.array(
+            [
+                atol + rtol * mean
+                for mean, (atol, rtol) in zip(avg, self._target_tols)
+            ]
+        )
         target_ntraj = np.max((avg2 - abs(avg) ** 2) / target**2 + 1)
 
         self._estimated_ntraj = min(target_ntraj, self._target_ntraj)
         if (self._estimated_ntraj - self.num_trajectories) <= 0:
-            self.stats['end_condition'] = 'target tolerance reached'
+            self.stats["end_condition"] = "target tolerance reached"
         return self._estimated_ntraj - self.num_trajectories
 
     def _post_init(self):
         self.num_trajectories = 0
         self._target_ntraj = None
 
-        store_states = self.options['store_states']
-        store_final_state = self.options['store_final_state']
-        store_traj = self.options['keep_runs_results']
-
         self.add_processor(self._increment_traj)
-        if store_traj:
+        store_trajectory = self.options["keep_runs_results"]
+        if store_trajectory:
             self.add_processor(self._store_trajectory)
-        if store_states is None:
-            store_states = len(self._raw_ops) == 0
-        if store_states:
+        if self._store_average_density_matricies:
             self.add_processor(self._reduce_states)
-        if store_states or store_final_state:
+        if self._store_final_density_matrix:
             self.add_processor(self._reduce_final_state)
         if self._raw_ops:
             self.add_processor(self._reduce_expect)
 
         self._early_finish_check = self._no_end
-        self.stats['end_condition'] = 'unknown'
+        self.stats["end_condition"] = "unknown"
 
     def add(self, trajectory_info):
         """
@@ -676,7 +733,7 @@ def add_end_condition(self, ntraj, target_tol=None):
             Error estimation is done with jackknife resampling.
         """
         self._target_ntraj = ntraj
-        self.stats['end_condition'] = 'timeout'
+        self.stats["end_condition"] = "timeout"
 
         if target_tol is None:
             self._early_finish_check = self._fixed_end
@@ -691,14 +748,16 @@ def add_end_condition(self, ntraj, target_tol=None):
 
         targets = np.array(target_tol)
         if targets.ndim == 0:
-            self._target_tols = np.array([(target_tol, 0.)] * num_e_ops)
+            self._target_tols = np.array([(target_tol, 0.0)] * num_e_ops)
         elif targets.shape == (2,):
             self._target_tols = np.ones((num_e_ops, 2)) * targets
         elif targets.shape == (num_e_ops, 2):
             self._target_tols = targets
         else:
-            raise ValueError("target_tol must be a number, a pair of (atol, "
-                             "rtol) or a list of (atol, rtol) for each e_ops")
+            raise ValueError(
+                "target_tol must be a number, a pair of (atol, "
+                "rtol) or a list of (atol, rtol) for each e_ops"
+            )
 
         self._early_finish_check = self._target_tolerance_end
 
@@ -718,8 +777,18 @@ def average_states(self):
         States averages as density matrices.
         """
         if self._sum_states is None:
-            return None
-        return [final / self.num_trajectories for final in self._sum_states]
+            if not (self.trajectories and self.trajectories[0].states):
+                return None
+            self._sum_states = [
+                qzero_like(self._to_dm(state))
+                for state in self.trajectories[0].states
+            ]
+            for trajectory in self.trajectories:
+                self._reduce_states(trajectory)
+
+        return [
+            final / self.num_trajectories for final in self._sum_states
+        ]
 
     @property
     def states(self):
@@ -744,6 +813,8 @@ def average_final_state(self):
         Last states of each trajectories averaged into a density matrix.
         """
         if self._sum_final_states is None:
+            if self.average_states is not None:
+                return self.average_states[-1]
             return None
         return self._sum_final_states / self.num_trajectories
 
@@ -770,6 +841,10 @@ def runs_expect(self):
     def expect(self):
         return [np.array(val) for val in self.e_data.values()]
 
+    @property
+    def e_data(self):
+        return self.runs_e_data or self.average_e_data
+
     def steady_state(self, N=0):
         """
         Average the states of the last ``N`` times of every runs as a density
@@ -796,16 +871,13 @@ def __repr__(self):
         ]
         if self.stats:
             lines.append("  Solver stats:")
-            lines.extend(
-                f"    {k}: {v!r}"
-                for k, v in self.stats.items()
-            )
+            lines.extend(f"    {k}: {v!r}" for k, v in self.stats.items())
         if self.times:
             lines.append(
                 f"  Time interval: [{self.times[0]}, {self.times[-1]}]"
                 f" ({len(self.times)} steps)"
             )
-        lines.append(f"  Number of e_ops: {len(self.e_ops)}")
+        lines.append(f"  Number of e_ops: {len(self.e_data)}")
         if self.states:
             lines.append("  States saved.")
         elif self.final_state is not None:
@@ -827,23 +899,30 @@ def __add__(self, other):
             raise ValueError("Shared `e_ops` is required to merge results")
         if self.times != other.times:
             raise ValueError("Shared `times` are is required to merge results")
-        new = self.__class__(self._raw_ops, self.options,
-                             solver=self.solver, stats=self.stats)
+        new = self.__class__(
+            self._raw_ops, self.options, solver=self.solver, stats=self.stats
+        )
         if self.trajectories and other.trajectories:
             new.trajectories = self.trajectories + other.trajectories
         new.num_trajectories = self.num_trajectories + other.num_trajectories
         new.times = self.times
         new.seeds = self.seeds + other.seeds
 
-        if self._sum_states is not None and other._sum_states is not None:
-            new._sum_states = self._sum_states + other._sum_states
+        if (
+            self._sum_states is not None
+            and other._sum_states is not None
+        ):
+            new._sum_states = (
+                self._sum_states + other._sum_states
+            )
 
         if (
             self._sum_final_states is not None
             and other._sum_final_states is not None
         ):
             new._sum_final_states = (
-                self._sum_final_states + other._sum_final_states
+                self._sum_final_states
+                + other._sum_final_states
             )
         new._target_tols = None
 
@@ -854,23 +933,21 @@ def __add__(self, other):
 
         for i, k in enumerate(self._raw_ops):
             new._sum_expect.append(self._sum_expect[i] + other._sum_expect[i])
-            new._sum2_expect.append(self._sum2_expect[i]
-                                    + other._sum2_expect[i])
+            new._sum2_expect.append(
+                self._sum2_expect[i] + other._sum2_expect[i]
+            )
 
             avg = new._sum_expect[i] / new.num_trajectories
             avg2 = new._sum2_expect[i] / new.num_trajectories
 
             new.average_e_data[k] = list(avg)
-            new.e_data = new.average_e_data
-
             new.std_e_data[k] = np.sqrt(np.abs(avg2 - np.abs(avg**2)))
 
-            if new.trajectories:
+            if self.runs_e_data and other.runs_e_data:
                 new.runs_e_data[k] = self.runs_e_data[k] + other.runs_e_data[k]
-                new.e_data = new.runs_e_data
 
         new.stats["run time"] += other.stats["run time"]
-        new.stats['end_condition'] = "Merged results"
+        new.stats["end_condition"] = "Merged results"
 
         return new
 
@@ -881,8 +958,9 @@ class McTrajectoryResult(Result):
     """
 
     def __init__(self, e_ops, options, *args, **kwargs):
-        super().__init__(e_ops, {**options, "normalize_output": False},
-                         *args, **kwargs)
+        super().__init__(
+            e_ops, {**options, "normalize_output": False}, *args, **kwargs
+        )
 
 
 class McResult(MultiTrajResult):
@@ -922,6 +1000,7 @@ class McResult(MultiTrajResult):
         For each runs, a list of every collapse as a tuple of the time it
         happened and the corresponding ``c_ops`` index.
     """
+
     # Collapse are only produced by mcsolve.
 
     def _add_collapse(self, trajectory):
@@ -941,7 +1020,7 @@ def col_times(self):
         out = []
         for col_ in self.collapse:
             col = list(zip(*col_))
-            col = ([] if len(col) == 0 else col[0])
+            col = [] if len(col) == 0 else col[0]
             out.append(col)
         return out
 
@@ -953,7 +1032,7 @@ def col_which(self):
         out = []
         for col_ in self.collapse:
             col = list(zip(*col_))
-            col = ([] if len(col) == 0 else col[1])
+            col = [] if len(col) == 0 else col[1]
             out.append(col)
         return out
 
@@ -968,7 +1047,9 @@ def photocurrent(self):
             for t, which in collapses:
                 cols[which].append(t)
         mesurement = [
-            np.histogram(cols[i], tlist)[0] / np.diff(tlist) / self.num_trajectories
+            np.histogram(cols[i], tlist)[0]
+            / np.diff(tlist)
+            / self.num_trajectories
             for i in range(self.num_c_ops)
         ]
         return mesurement
@@ -984,10 +1065,12 @@ def runs_photocurrent(self):
             cols = [[] for _ in range(self.num_c_ops)]
             for t, which in collapses:
                 cols[which].append(t)
-            measurements.append([
-                np.histogram(cols[i], tlist)[0] / np.diff(tlist)
-                for i in range(self.num_c_ops)
-            ])
+            measurements.append(
+                [
+                    np.histogram(cols[i], tlist)[0] / np.diff(tlist)
+                    for i in range(self.num_c_ops)
+                ]
+            )
         return measurements
 
 
@@ -998,6 +1081,7 @@ class McResultImprovedSampling(McResult, MultiTrajResult):
     using the improved sampling algorithm, which samples the no-jump trajectory
     first and then only samples jump trajectories afterwards.
     """
+
     def __init__(self, e_ops, options, **kw):
         MultiTrajResult.__init__(self, e_ops=e_ops, options=options, **kw)
         self._sum_expect_no_jump = None
@@ -1016,14 +1100,16 @@ def _reduce_states(self, trajectory):
         if self.num_trajectories == 1:
             self._sum_states_no_jump = [
                 accu + self._to_dm(state)
-                for accu, state
-                in zip(self._sum_states_no_jump, trajectory.states)
+                for accu, state in zip(
+                    self._sum_states_no_jump, trajectory.states
+                )
             ]
         else:
             self._sum_states_jump = [
                 accu + self._to_dm(state)
-                for accu, state
-                in zip(self._sum_states_jump, trajectory.states)
+                for accu, state in zip(
+                    self._sum_states_jump, trajectory.states
+                )
             ]
 
     def _reduce_final_state(self, trajectory):
@@ -1040,13 +1126,15 @@ def _average_computer(self):
 
     def _add_first_traj(self, trajectory):
         super()._add_first_traj(trajectory)
-        if trajectory.states:
+        if trajectory.states and self._store_average_density_matricies:
             del self._sum_states
-            self._sum_states_no_jump = [qzero_like(self._to_dm(state))
-                                        for state in trajectory.states]
-            self._sum_states_jump = [qzero_like(self._to_dm(state))
-                                     for state in trajectory.states]
-        if trajectory.final_state:
+            self._sum_states_no_jump = [
+                qzero_like(self._to_dm(state)) for state in trajectory.states
+            ]
+            self._sum_states_jump = [
+                qzero_like(self._to_dm(state)) for state in trajectory.states
+            ]
+        if trajectory.final_state and self._store_final_density_matrix:
             state = trajectory.final_state
             del self._sum_final_states
             self._sum_final_states_no_jump = qzero_like(self._to_dm(state))
@@ -1063,10 +1151,12 @@ def _add_first_traj(self, trajectory):
         self._sum2_expect_no_jump = [
             np.zeros_like(expect) for expect in trajectory.expect
         ]
-        self._sum_expect_jump = [np.zeros_like(expect)
-                                 for expect in trajectory.expect]
-        self._sum2_expect_jump = [np.zeros_like(expect)
-                                  for expect in trajectory.expect]
+        self._sum_expect_jump = [
+            np.zeros_like(expect) for expect in trajectory.expect
+        ]
+        self._sum2_expect_jump = [
+            np.zeros_like(expect) for expect in trajectory.expect
+        ]
         del self._sum_expect
         del self._sum2_expect
 
@@ -1088,12 +1178,12 @@ def _reduce_expect(self, trajectory):
             else:
                 self._sum_expect_jump[i] += expect_traj * (1 - p)
                 self._sum2_expect_jump[i] += expect_traj**2 * (1 - p)
-                avg = (self._sum_expect_no_jump[i]
-                       + self._sum_expect_jump[i]
-                       / (self.num_trajectories - 1))
-                avg2 = (self._sum2_expect_no_jump[i]
-                        + self._sum2_expect_jump[i]
-                        / (self.num_trajectories - 1))
+                avg = self._sum_expect_no_jump[i] + (
+                    self._sum_expect_jump[i] / (self.num_trajectories - 1)
+                )
+                avg2 = self._sum2_expect_no_jump[i] + (
+                    self._sum2_expect_jump[i] / (self.num_trajectories - 1)
+                )
 
             self.average_e_data[k] = list(avg)
 
@@ -1110,12 +1200,28 @@ def average_states(self):
         States averages as density matrices.
         """
         if self._sum_states_no_jump is None:
-            return None
+            if not (self.trajectories and self.trajectories[0].states):
+                return None
+            self._sum_states_no_jump = [
+                qzero_like(self._to_dm(state))
+                for state in self.trajectories[0].states
+            ]
+            self._sum_states_jump = [
+                qzero_like(self._to_dm(state))
+                for state in self.trajectories[0].states
+            ]
+            self.num_trajectories = 0
+            for trajectory in self.trajectories:
+                self.num_trajectories += 1
+                self._reduce_states(trajectory)
         p = self.no_jump_prob
-        return [p * final_no_jump
-                + (1 - p) * final_jump / (self.num_trajectories - 1)
-                for final_no_jump, final_jump in
-                zip(self._sum_states_no_jump, self._sum_states_jump)]
+        return [
+            p * final_no_jump
+            + (1 - p) * final_jump / (self.num_trajectories - 1)
+            for final_no_jump, final_jump in zip(
+                self._sum_states_no_jump, self._sum_states_jump
+            )
+        ]
 
     @property
     def average_final_state(self):
@@ -1123,11 +1229,11 @@ def average_final_state(self):
         Last states of each trajectory averaged into a density matrix.
         """
         if self._sum_final_states_no_jump is None:
-            return None
+            if self.average_states is not None:
+                return self.average_states[-1]
         p = self.no_jump_prob
-        return (
-            p * self._sum_final_states_no_jump
-            + (1 - p) * self._sum_final_states_jump
+        return p * self._sum_final_states_no_jump + (
+            ((1 - p) * self._sum_final_states_jump)
             / (self.num_trajectories - 1)
         )
 
@@ -1145,8 +1251,10 @@ def photocurrent(self):
             for t, which in collapses:
                 cols[which].append(t)
         mesurement = [
-            (1 - self.no_jump_prob) / (self.num_trajectories - 1) *
-            np.histogram(cols[i], tlist)[0] / np.diff(tlist)
+            (1 - self.no_jump_prob)
+            / (self.num_trajectories - 1)
+            * np.histogram(cols[i], tlist)[0]
+            / np.diff(tlist)
             for i in range(self.num_c_ops)
         ]
         return mesurement
@@ -1241,15 +1349,15 @@ def _add_first_traj(self, trajectory):
     def _add_trace(self, trajectory):
         new_trace = np.array(trajectory.trace)
         self._sum_trace += new_trace
-        self._sum2_trace += np.abs(new_trace)**2
+        self._sum2_trace += np.abs(new_trace) ** 2
 
         avg = self._sum_trace / self.num_trajectories
         avg2 = self._sum2_trace / self.num_trajectories
 
         self.average_trace = avg
-        self.std_trace = np.sqrt(np.abs(avg2 - np.abs(avg)**2))
+        self.std_trace = np.sqrt(np.abs(avg2 - np.abs(avg) ** 2))
 
-        if self.options['keep_runs_results']:
+        if self.options["keep_runs_results"]:
             self.runs_trace.append(trajectory.trace)
 
     @property


Test patch: