Repository: qutip/qutip

Issue:
UnboundLocalError when `mesolve` is used with callable `e_ops`
### Bug Description

`UnboundLocalError: local variable 'rho_t' referenced before assignment` is thrown when `mesolve` is used with callable `e_ops`. Note, it seems to be fixed on the 5.x branch (using `pip install --pre qutip`).

### Code to Reproduce the Bug

```shell
from qutip import * 
import numpy as np

s_00 = tensor(basis(2,0), basis(2,0))
P_00 = s_00 * s_00.dag()
e_ops = [ 
    lambda t, state: (state.ptrace([0,1]) * P_00).tr()   
    ]

tlist = np.linspace(0, 2*np.pi, 50)
res = mesolve(tensor(identity(2), identity(2), identity(3)),
              tensor(P_00, thermal_dm(3,0.1)), 
              tlist, 
              [],
              e_ops)
```


### Code Output

```shell
---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
<ipython-input-1-49ed87444849> in <module>
      9 
     10 tlist = np.linspace(0, 2*np.pi, 50)
---> 11 res = mesolve(tensor(identity(2), identity(2), identity(3)),
     12               tensor(P_00, thermal_dm(3,0.1)),
     13               tlist,

1 frames
/usr/local/lib/python3.9/dist-packages/qutip/mesolve.py in mesolve(H, rho0, tlist, c_ops, e_ops, args, options, progress_bar, _safe_mode)
    242         func(0., v, *ode_args) + v
    243 
--> 244     res = _generic_ode_solve(func, ode_args, rho0, tlist, e_ops, options,
    245                              progress_bar, dims=rho0.dims)
    246     res.num_collapse = len(c_ops)

/usr/local/lib/python3.9/dist-packages/qutip/mesolve.py in _generic_ode_solve(func, ode_args, rho0, tlist, e_ops, opt, progress_bar, dims)
    517         for m in range(n_expt_op):
    518             if not isinstance(e_ops[m], Qobj) and callable(e_ops[m]):
--> 519                 output.expect[m][t_idx] = e_ops[m](t, rho_t)
    520                 continue
    521             output.expect[m][t_idx] = expect_rho_vec(e_ops_data[m], r.y,

UnboundLocalError: local variable 'rho_t' referenced before assignment
```


### Expected Behaviour

no error.

### Your Environment

```shell
QuTiP Version:      4.7.1
Numpy Version:      1.22.4
Scipy Version:      1.10.1
Cython Version:     0.29.33
Matplotlib Version: 3.5.3
Python Version:     3.9.16
Number of CPUs:     2
BLAS Info:          OPENBLAS
OPENMP Installed:   False
INTEL MKL Ext:      False
Platform Info:      Linux (x86_64)
Installation path:  /usr/local/lib/python3.9/dist-packages/qutip
```


### Additional Context

_No response_


Discussion:
Thank you for reporting.

`e_ops=callable` should work and allow you to obtain the result you need for now.
`e_ops=[callable]` support should have been added, but it seems it was only properly done for `sesolve`.

The error is in the qutip-4.7.X branch at: https://github.com/qutip/qutip/blob/15880398a222ebfd62fd52361cc3852a685b1ac7/qutip/mesolve.py#L458
where is should be set to `True` if at least one of the `e_ops` is a callable but not a `Qobj`.
Hello! I would like to fix this issue.
@manulpatel
Thank you for the interest.
We would be grateful if you fixed the issue.


Resolving patch:
diff --git a/qutip/mesolve.py b/qutip/mesolve.py
--- a/qutip/mesolve.py
+++ b/qutip/mesolve.py
@@ -449,10 +449,12 @@ def _generic_ode_solve(func, ode_args, rho0, tlist, e_ops, opt,
 
     e_ops_data = []
     output.expect = []
+    need_qobj_state = opt.store_states
     if callable(e_ops):
         n_expt_op = 0
         expt_callback = True
         output.num_expect = 1
+        need_qobj_state = True
     elif isinstance(e_ops, list):
         n_expt_op = len(e_ops)
         expt_callback = False
@@ -460,10 +462,13 @@ def _generic_ode_solve(func, ode_args, rho0, tlist, e_ops, opt,
         if n_expt_op == 0:
             # fall back on storing states
             opt.store_states = True
+            need_qobj_state = True
         else:
             for op in e_ops:
                 if not isinstance(op, Qobj) and callable(op):
                     output.expect.append(np.zeros(n_tsteps, dtype=complex))
+                    need_qobj_state = True
+                    e_ops_data.append(None)
                     continue
                 if op.dims != rho0.dims:
                     raise TypeError(f"e_ops dims ({op.dims}) are not "
@@ -497,7 +502,7 @@ def get_curr_state_data(r):
                             "the allowed number of substeps by increasing "
                             "the nsteps parameter in the Options class.")
 
-        if opt.store_states or expt_callback:
+        if need_qobj_state:
             cdata = get_curr_state_data(r)
             fdata = dense2D_to_fastcsr_fmode(cdata, size, size)
 
@@ -517,10 +522,10 @@ def get_curr_state_data(r):
         for m in range(n_expt_op):
             if not isinstance(e_ops[m], Qobj) and callable(e_ops[m]):
                 output.expect[m][t_idx] = e_ops[m](t, rho_t)
-                continue
-            output.expect[m][t_idx] = expect_rho_vec(e_ops_data[m], r.y,
-                                                     e_ops[m].isherm
-                                                     and rho0.isherm)
+            else:
+                output.expect[m][t_idx] = expect_rho_vec(e_ops_data[m], r.y,
+                                                         e_ops[m].isherm
+                                                         and rho0.isherm)
 
         if t_idx < n_tsteps - 1:
             r.integrate(r.t + dt[t_idx])


Test patch:
diff --git a/qutip/tests/test_mesolve.py b/qutip/tests/test_mesolve.py
--- a/qutip/tests/test_mesolve.py
+++ b/qutip/tests/test_mesolve.py
@@ -957,6 +957,22 @@ def test_tlist_h_with_constant_c_ops():
     assert len(result.states) == len(few_times)
 
 
+def test_mixed_e_ops():
+    """
+    Test callable and Qobj e_ops can mix.
+
+    See gh-2118.
+    """
+    state = basis(2, 0)
+    hamiltonian = QobjEvo(sigmax())
+    collapse = create(2)
+    e_ops = [qeye(2), lambda t, qobj: qobj.norm()]
+    result = mesolve(
+        hamiltonian, state, [0, 1, 2], c_ops=[collapse], e_ops=e_ops
+    )
+    assert result.num_expect == 2
+
+
 def test_tlist_h_with_other_tlist_c_ops_raises():
     state = basis(2, 0)
     all_times = np.linspace(0, 1, 11)