Repository: qutip/qutip

Issue:
average_states option not functioning with ssesolve
### Bug Description

The documentation of the [`StochasticSolverOptions`](https://qutip.org/docs/latest/apidoc/classes.html#qutip.stochastic.StochasticSolverOptions) class states that the `average_states` option can be passed through the `options` keyword and that it will be used:

> options:
> Generic solver options. Only `options.average_states` and `options.store_states` are used.

However, it does not seem that this option is working when passed to `ssesolve`, i.e., the states are never averaged. I attach an example code to reproduce the behaviour.


### Code to Reproduce the Bug

```shell
## IMPORTS

import matplotlib.pyplot as plt
import numpy as np
import qutip as qt

# Inline plotting
%matplotlib inline


## PARAMETERS

# Maximum number of photons
n_max = 50

# Resonator frequency
w_r = 10
# Resonator leakage rate
kappa = 1

# Drive parameters
w_d = 12
t_cut = 10/kappa
t_ramp = 0.05/kappa
drive_pars = {"w_d": w_d, "t_cut": t_cut, "t_ramp": t_ramp}

# Integration parameters
dt = (2*np.pi/w_d)/30
T = 3/kappa
times = np.arange(0, T, dt)


## FREE HAMILTONIAN

# Annihilation operator
a = qt.destroy(n_max)

# Number operator
n = a.dag() * a

# Free Hamiltonian
H_0 = w_r * n


## DRIVE HAMILTONIAN

# Define drive pulse
def smooth_square_pulse(t, t0, tf, w):
    pulse = (1 + np.tanh((tf - t)/w))/2
    pulse = pulse * (1 + np.tanh((t-t0)/w))/2
    return pulse
def pulse(t, args):
    w_d = args["w_d"]
    t_cut = args["t_cut"]
    t_ramp = args["t_ramp"]
    drive = smooth_square_pulse(t, 0, t_cut, t_ramp)
    drive = drive * np.sin(w_d * t)
    return drive

# Generate of the drive
V_d_gen = -1j * (a-a.dag())


## STOCHASTIC COLLAPSE OPERATORS
sc_ops = [np.sqrt(kappa) * a]


## HAMILTONIAN TO FEED TO THE SOLVER
H = [H_0, [V_d_gen, pulse]]


## INITIAL STATE
psi_0 = qt.fock(n_max, 0)


## STOCHASTIC SOLVER KWARGS AND OPTIONS
options = qt.Options()
options.average_states = True
kwargs = {"sc_ops": sc_ops, "args": drive_pars, "ntraj": 500, "method": "heterodyne", "store_measurement": True, "options": options}


## SOLVE THE STOCHASTIC MASTER EQUATIONS
solution = qt.ssesolve(H, psi_0, times, **kwargs)


## OUTPUT
print("==========================================")
print("==========================================")
print("\nValue of the average_states option: {}.".format(options.average_states))
print("\nAvailable solver attributes:")
print(solution.__dict__.keys())
print("\nCheck the type of the avg_states attribute")
print(type(solution.avg_states))
print("\nThe states have not been averaged over. The shape of the output is {}, with the first dimension labelling the {} trajectories.".format(np.array(solution.states).shape, solution.ntraj))
```


### Code Output

```shell
10.0%. Run time:   1.02s. Est. time left: 00:00:00:09
20.0%. Run time:   2.04s. Est. time left: 00:00:00:08
30.0%. Run time:   3.10s. Est. time left: 00:00:00:07
40.0%. Run time:   4.10s. Est. time left: 00:00:00:06
50.0%. Run time:   5.11s. Est. time left: 00:00:00:05
60.0%. Run time:   6.11s. Est. time left: 00:00:00:04
70.0%. Run time:   7.19s. Est. time left: 00:00:00:03
80.0%. Run time:   8.20s. Est. time left: 00:00:00:02
90.0%. Run time:   9.21s. Est. time left: 00:00:00:01
Total run time:  10.22s
==========================================
==========================================

Value of the average_states option: True.

Available solver attributes:
dict_keys(['solver', 'times', 'states', 'expect', 'num_expect', 'num_collapse', 'ntraj', 'seeds', 'col_times', 'col_which', 'ss', 'measurement', 'noise', 'traj_states', 'avg_states', 'se'])

Check the type of the avg_states attribute
<class 'NoneType'>

The states have not been averaged over. The shape of the output is (500, 172, 50, 1), with the first dimension labelling the 500 trajectories.
```


### Expected Behaviour

I expect the states to be averaged over when the average_states option is passed.

### Your Environment

```shell
QuTiP Version:      4.7.2
Numpy Version:      1.24.3
Scipy Version:      1.10.1
Cython Version:     0.29.35
Matplotlib Version: 3.7.1
Python Version:     3.10.6
Number of CPUs:     12
BLAS Info:          OPENBLAS
OPENMP Installed:   False
INTEL MKL Ext:      False
Platform Info:      Darwin (arm64)
```


### Additional Context

_No response_


Discussion:

Resolving patch:
diff --git a/qutip/stochastic.py b/qutip/stochastic.py
--- a/qutip/stochastic.py
+++ b/qutip/stochastic.py
@@ -1303,7 +1303,8 @@ def _sesolve_generic(sso, options, progress_bar):
     noise = []
     for result in results:
         states_list, dW, m, expect = result
-        res.states.append(states_list)
+        if options.average_states or options.store_states:
+            res.states.append(states_list)
         noise.append(dW)
         res.measurement.append(m)
         res.expect += expect
@@ -1320,15 +1321,24 @@ def _sesolve_generic(sso, options, progress_bar):
     # ajgpitch 2019-10-25: np.any(res.states) seems to error
     # I guess there may be a potential exception if there are no states?
     # store individual trajectory states
-    res.traj_states = res.states
+    if options.store_states:
+        res.traj_states = res.states
+    else:
+        res.traj_states = None
     res.avg_states = None
-    if options.average_states and options.store_states:
+    if options.average_states:
         avg_states_list = []
         for n in range(len(res.times)):
-            tslot_states = [res.states[mm][n].data for mm in range(nt)]
+            if res.states[0][n].shape[1] == 1:
+                tslot_states = [
+                    res.states[mm][n].proj().data
+                    for mm in range(nt)
+                ]
+            else:
+                tslot_states = [res.states[mm][n].data for mm in range(nt)]
             if len(tslot_states) > 0:
                 state = Qobj(np.sum(tslot_states),
-                             dims=res.states[0][n].dims).unit()
+                             dims=[res.states[0][n].dims[0]] * 2).unit()
                 avg_states_list.append(state)
         # store average states
         res.states = res.avg_states = avg_states_list


Test patch:
diff --git a/qutip/tests/test_stochastic_me.py b/qutip/tests/test_stochastic_me.py
--- a/qutip/tests/test_stochastic_me.py
+++ b/qutip/tests/test_stochastic_me.py
@@ -5,7 +5,7 @@
 from qutip import (smesolve, mesolve, photocurrent_mesolve, liouvillian,
                    QobjEvo, spre, spost, destroy, coherent, parallel_map,
                    qeye, fock_dm, general_stochastic, ket2dm, num,
-                   basis, sigmax, sigmay, sigmaz, sigmam)
+                   basis, sigmax, sigmay, sigmaz, sigmam, Options)
 
 def f(t, args):
     return args["a"] * t
@@ -144,6 +144,7 @@ def test_smesolve_homodyne():
     psi0 = coherent(N, 0.5)
     sc_ops = [np.sqrt(gamma) * a, np.sqrt(gamma) * a * 0.5]
     e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]
+    options = Options(average_states=True, store_states=True)
 
     times = np.linspace(0, 1.0, 21)
     res_ref = mesolve(H, psi0, times, sc_ops, e_ops, args={"a":2})
@@ -161,12 +162,17 @@ def test_smesolve_homodyne():
         res = smesolve(H, psi0, times, [], sc_ops, e_ops,
                        ntraj=ntraj, nsubsteps=nsubsteps, args={"a":2},
                        method='homodyne', store_measurement=True,
-                       solver=solver, map_func=parallel_map)
+                       solver=solver, map_func=parallel_map,
+                       options=options)
         assert_(all([np.mean(abs(res.expect[idx] - res_ref.expect[idx])) < tol
                      for idx in range(len(e_ops))]))
         assert_(len(res.measurement) == ntraj)
         assert_(all([m.shape == (len(times), len(sc_ops))
                      for m in res.measurement]))
+        assert_(
+            np.array([state.full() for state in res.states]).shape
+            == (len(times), N, N)
+        )
 
 
 @pytest.mark.slow
diff --git a/qutip/tests/test_stochastic_se.py b/qutip/tests/test_stochastic_se.py
--- a/qutip/tests/test_stochastic_se.py
+++ b/qutip/tests/test_stochastic_se.py
@@ -3,7 +3,7 @@
 from numpy.testing import assert_
 
 from qutip import (ssesolve, destroy, coherent, mesolve, fock, qeye,
-                   parallel_map, photocurrent_sesolve, num)
+                   parallel_map, photocurrent_sesolve, num, Options)
 
 def f(t, args):
     return args["a"] * t
@@ -140,13 +140,15 @@ def test_ssesolve_homodyne():
     psi0 = coherent(N, 0.5)
     sc_ops = [np.sqrt(gamma) * a, np.sqrt(gamma) * a*0.5]
     e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]
+    options = Options(average_states=True, store_states=True)
 
     times = np.linspace(0, 2.5, 50)
     res_ref = mesolve(H, psi0, times, sc_ops, e_ops, args={"a":2})
     res = ssesolve(H, psi0, times, sc_ops, e_ops,
                    ntraj=ntraj, nsubsteps=nsubsteps,
                    method='homodyne', store_measurement=True,
-                   map_func=parallel_map, args={"a":2})
+                   map_func=parallel_map, args={"a":2},
+                   options=options)
 
     assert_(all([np.mean(abs(res.expect[idx] - res_ref.expect[idx])) < tol
                  for idx in range(len(e_ops))]))
@@ -154,6 +156,10 @@ def test_ssesolve_homodyne():
     assert_(len(res.measurement) == ntraj)
     assert_(all([m.shape == (len(times), len(sc_ops))
                  for m in res.measurement]))
+    assert_(
+        np.array([state.full() for state in res.states]).shape
+        == (len(times), N, N)
+    )
 
 
 @pytest.mark.slow