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

[Akhils777]

This PR begins the process of integrating GRAPE and CRAB from QTRL into QOC.

✔️ What’s done in this PR:

• Removed QTRL’s logging system (qutip_qtrl.logging_utils) from src/qutip_qoc/_grape.py and src/qutip_qoc/_crab.py
• Deleted the following lines:
  • import qutip_qtrl.logging_utils as logging
  • logger = logging.get_logger()
  • All logger.debug(...) statements
• Retained all self._qtrl.log_level conditions (no functional logic changed)

Setup:

• As a staging step, I copied the original QTRL files into a new folder: q2/, inside the QOC repository.
  This allows us to reference the original code during refactor without losing context.

Cherry-picked from PR #47:

A previously discovered issue in the JAX-compatible infidelity computation for open-system TRACEDIFF fidelity is also addressed here.

Specifically, the original implementation used:

Qobj(diff.data.adjoint(), dims=...)

which failed under JAX tracing due to dimension mismatches and non-JAX-compatible operations like .tr() and .full().

Fix:
I replaced this with a fully JAX-compatible expression:

diff_dag = diff.dag()
g = 0.5 * jnp.trace(diff_dag.data._jxa @ diff.data._jxa)

This avoids any QuTiP methods that break under JAX tracing and accesses the underlying JAX array safely.

Additional Fix (Test Failures):
After raising the PR, the tests failed due to a breaking change in jax==0.6.2, which removed access to jaxlib.xla_extension.PjitFunction used in result.py for JIT type checks.

To fix this, I updated the type-checking logic using:

_jitfun_type = type(jax.jit(lambda x: x))

This restored compatibility across JAX versions and fixed the test failures.

Now, both JOPT and GRAPE run correctly for open-system state transfer problems using TRACEDIFF fidelity.

Let me know if you have suggestions or if you’d like this broken down further. Would love to hear your thoughts!

[BoxiLi]

Let's also remove the qutip_qtrl used in

qutip-qoc/src/qutip_qoc/pulse_optim.py

Lines 8 to 10 in d5ff028

	import qutip_qtrl.logging_utils as logging
	import qutip_qtrl.pulseoptim as cpo

And import cpo from the files you copied. Then it really tests the imported modules.

[BoxiLi]

What is src/qutip_qoc/q2/test_grape.py?

Is this a new test? I think it belongs to the tests folder not src.

[rochisha0]

Hello @Akhils777 , good start.

The idea here is to take inspiration from qutip-qtrl but follow the code standards of qoc. The structure right now is that, you have to update the file _grape.py in src folder that is callable by pulse_optim.py. The previous tests should still work as there should be no change in the existing API.

[BoxiLi]

Hi Rochisha, thanks! Maybe it is easier to first merge this "moving" PR after we confirm that qutip_qtrl is no longer a requirement as an external dependence, and the tests pass. Further changes could be included in separate PR(s). This will make it easier to track the work.

Note that this is merging to a development branch, not the main branch; this won't break anything in main.