average_states option not functioning with ssesolve

[BenjaminDAnjou]

Bug Description

The documentation of the 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

## 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

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

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

[Ericgig]

Thank you for reporting.
Right now, the average states is only computed when both options.average_states and options.store_states are set.
ssesovle average states as ket, not dm...
The fix is in #2217.

[Ericgig]

Fixed in 4.7.3