Fix: Enable JOPT to support open-system optimization with TRACEDIFF fidelity (qutip#49)

Author: Akhils777

src/qutip_qoc/_jopt.py

@@ -150,8 +150,8 @@ def _infid(self, params):
         if self._fid_type == "TRACEDIFF":
             diff = X - self._target
             # to prevent if/else in qobj.dag() and qobj.tr()
-            diff_dag = Qobj(diff.data.adjoint(), dims=diff.dims)
-            g = 1 / 2 * (diff_dag * diff).data.trace()
+            diff_dag = diff.dag() # direct access to JAX array, no fallback!
+            g = 1 / 2 * jnp.trace(diff_dag.data._jxa @ diff.data._jxa)
             infid = jnp.real(self._norm_fac * g)
         else:
             g = self._norm_fac * self._target.overlap(X)
@@ -160,4 +160,4 @@ def _infid(self, params):
             elif self._fid_type == "SU":  # f_SU (incl global phase)
                 infid = 1 - jnp.real(g)
 
-        return infid
+        return infid

src/qutip_qoc/result.py

@@ -13,7 +13,7 @@
 try:
     import jax
     import jaxlib
-    _jitfun_type = jaxlib.xla_extension.PjitFunction
+    _jitfun_type = type(jax.jit(lambda x: x))
 except ImportError:
     _jitfun_type = None
 

tests/test_jopt_open_system_bug.py

@@ -0,0 +1,39 @@
+import numpy as np
+import qutip as qt
+from qutip_qoc import Objective, optimize_pulses
+
+from jax import jit, numpy
+
+def test_open_system_jopt_runs_without_error():
+    Hd = qt.Qobj(np.diag([1, 2]))
+    c_ops = [np.sqrt(0.1) * qt.sigmam()]
+    Hc = qt.sigmax()
+
+    Ld = qt.liouvillian(H=Hd, c_ops=c_ops)
+    Lc = qt.liouvillian(Hc)
+
+    initial_state = qt.fock_dm(2, 0)
+    target_state = qt.fock_dm(2, 1)
+
+    times = np.linspace(0, 2 * np.pi, 250)
+
+    @jit
+    def sin_x(t, c, **kwargs):
+        return c[0] * numpy.sin(c[1] * t)
+    L = [Ld, [Lc, sin_x]]
+
+    guess_params = [1, 0.5]
+
+    res_jopt = optimize_pulses(
+        objectives = Objective(initial_state, L, target_state),
+        control_parameters = {
+            "ctrl_x": {"guess": guess_params, "bounds": [(-1, 1), (0, 2 * np.pi)]}
+        },
+        tlist = times,
+        algorithm_kwargs = {
+            "alg": "JOPT",
+            "fid_err_targ": 0.001,
+        },
+    )
+
+    assert res_jopt.infidelity < 0.25, f"Fidelity error too high: {res_jopt.infidelity}"