mesolve e_ops list can mix callable and Qobj

doc/changes/2118.bugfix

@@ -0,0 +1 @@
+mesolve can support mixed callable and Qobj e_ops.

qutip/mesolve.py

@@ -449,21 +449,26 @@ 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
         output.num_expect = n_expt_op
         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])

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)