[WIP] Refactor GRAPE: Merged QTRL Code into QuTiP-QOC

src/qutip_qoc/_crab.py

@@ -5,10 +5,10 @@
 with respect to the control parameters, according to the CRAB algorithm.
 """
 
-import qutip_qtrl.logging_utils as logging
+
 import copy
 
-logger = logging.get_logger()
+
 
 
 class _CRAB:
@@ -43,12 +43,7 @@ def infidelity(self, *args):
         # *** update stats ***
         if self._qtrl.stats is not None:
             self._qtrl.stats.num_fidelity_func_calls = self._qtrl.num_fid_func_calls
-            if self._qtrl.log_level <= logging.DEBUG:
-                logger.debug(
-                    "fidelity error call {}".format(
-                        self._qtrl.stats.num_fidelity_func_calls
-                    )
-                )
+
 
         amps = self._qtrl._get_ctrl_amps(args[0].copy())
         self._qtrl.dynamics.update_ctrl_amps(amps)

src/qutip_qoc/_grape.py

@@ -5,11 +5,11 @@
 with respect to the control parameters, according to the CRAB algorithm.
 """
 
-import qutip_qtrl.logging_utils as logging
+
 import copy
 
 
-logger = logging.get_logger()
+
 
 
 class _GRAPE:
@@ -42,12 +42,7 @@ def infidelity(self, *args):
         # *** update stats ***
         if self._qtrl.stats is not None:
             self._qtrl.stats.num_fidelity_func_calls = self._qtrl.num_fid_func_calls
-            if self._qtrl.log_level <= logging.DEBUG:
-                logger.debug(
-                    "fidelity error call {}".format(
-                        self._qtrl.stats.num_fidelity_func_calls
-                    )
-                )
+
 
         amps = self._qtrl._get_ctrl_amps(args[0].copy())
         self._qtrl.dynamics.update_ctrl_amps(amps)
@@ -81,10 +76,7 @@ def gradient(self, *args):
         self._qtrl.num_grad_func_calls += 1
         if self._qtrl.stats is not None:
             self._qtrl.stats.num_grad_func_calls = self._qtrl.num_grad_func_calls
-            if self._qtrl.log_level <= logging.DEBUG:
-                logger.debug(
-                    "gradient call {}".format(self._qtrl.stats.num_grad_func_calls)
-                )
+
         amps = self._qtrl._get_ctrl_amps(args[0].copy())
         self._qtrl.dynamics.update_ctrl_amps(amps)
         fid_comp = self._qtrl.dynamics.fid_computer

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/pulse_optim.py

@@ -5,8 +5,8 @@
 """
 import numpy as np
 
-import qutip_qtrl.logging_utils as logging
-import qutip_qtrl.pulseoptim as cpo
+import logging
+from qutip_qoc.q2 import pulseoptim as cpo
 
 from qutip_qoc._optimizer import _global_local_optimization
 from qutip_qoc._time import _TimeInterval

src/qutip_qoc/q2/cy_grape.py

@@ -0,0 +1,135 @@
+"""Fast routines for grape implemented on top of the data layer."""
+
+import qutip.core.data as _data
+
+
+def cy_overlap(op1, op2):
+    """
+    Return the overlap of op1 and op2.
+
+    Parameters
+    ----------
+    op1 : :class:`qutip.data.Data`
+        Data layer representation of first operator.
+    op2 : :class:`qutip.data.Data`
+        Data layer representation of second operator.
+
+    Result
+    ------
+    overlap : float
+        The value of the overlap.
+    """
+    return _data.trace(_data.adjoint(op1) @ op2) / op1.shape[0]
+
+
+def cy_grape_inner(
+    U,
+    u,
+    r,
+    J,
+    M,
+    U_b_list,
+    U_f_list,
+    H_ops,
+    dt,
+    eps,
+    alpha,
+    beta,
+    phase_sensitive,
+    use_u_limits,
+    u_min,
+    u_max,
+):
+    """
+    Perform one iteration of GRAPE control pulse
+    updates.
+
+    Parameters
+    ----------
+    U : :class:`qutip.data.Data`
+        The target unitary.
+
+    u : np.ndarray
+        The generated control pulses. It's shape
+        is (iterations, controls, times), i.e.
+        (R, J, M). The result of this iteration
+        is stored in u[r, :, :].
+
+    r : int
+        The number of this GRAPE iteration.
+
+    J : int
+        The number of controls in the Hamiltonian.
+
+    M : int
+        The number of times.
+
+    U_b_list : list of :class:`qutip.data.Data`
+        The backward propagators for each time.
+        The list has length M.
+
+    U_f_list : list of :class:`qutip.data.Data`
+        The forward propagators for each time.
+        The list has length M.
+
+    H_ops : list of :class:`qutip.data.Data`
+        The control operators from the Hamiltonian.
+        The list has length J.
+
+    dt : float
+        The time step.
+
+    eps : float
+        The distance to move along the gradient when updating
+        the controls.
+
+    alpha : float
+        The penalty to apply to higher power control signals.
+
+    beta : float
+        The penalty to apply to later control signals.
+
+    phase_sensitive : bool
+        Whether the overlap is phase sensitive.
+
+    use_u_limits : bool
+        Whether to apply limits to the control amplitudes.
+
+    u_min : float
+        Minimum control amplitude.
+
+    u_max : float
+        Maximum control amplitude.
+
+    Result
+    ------
+    The results are stored in u[r + 1, : , :].
+    """
+    for m in range(M - 1):
+        P = U_b_list[m] @ U
+        for j in range(J):
+            Q = 1j * dt * H_ops[j] @ U_f_list[m]
+
+            if phase_sensitive:
+                du = -cy_overlap(P, Q)
+            else:
+                du = -2 * cy_overlap(P, Q) * cy_overlap(U_f_list[m], P)
+
+            if alpha > 0.0:
+                # penalty term for high power control signals u
+                du += -2 * alpha * u[r, j, m] * dt
+
+            if beta:
+                # penalty term for late control signals u
+                du += -2 * beta * m**2 * u[r, j, m] * dt
+
+            u[r + 1, j, m] = u[r, j, m] + eps * du.real
+
+            if use_u_limits:
+                if u[r + 1, j, m] < u_min:
+                    u[r + 1, j, m] = u_min
+                elif u[r + 1, j, m] > u_max:
+                    u[r + 1, j, m] = u_max
+
+    for j in range(J):
+        u[r + 1, j, M - 1] = u[r + 1, j, M - 2]