Merge branch 'main' into cga

rochisha0 · web-flow · commit d11841e96975 · 2025-07-15T14:28:56.000+09:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -33,7 +33,7 @@ jobs:
       - name: Build a binary wheel and a source tarball
         run: python3 -m cibuildwheel --output-dir dist
       - name: Store the distribution packages
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: python-package-distributions
           path: dist/
@@ -52,7 +52,7 @@ jobs:
 
     steps:
       - name: Download all the dists
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
         with:
           name: python-package-distributions
           path: dist/
diff --git a/.github/workflows/build_documentation.yml b/.github/workflows/build_documentation.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
       - uses: actions/setup-python@v4
         name: Install Python
@@ -42,7 +42,7 @@ jobs:
           #   -T : display a full traceback if a Python exception occurs
 
       - name: Upload built files
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: qutip_qoc_html_docs
           path: doc/_build/html/*
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -47,7 +47,7 @@ jobs:
             python-version: "3.12"
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
@@ -92,7 +92,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
diff --git a/README.md b/README.md
@@ -26,9 +26,12 @@ To install the package, use
 pip install qutip-qoc
 ```
 
+By default, the dependencies required for JOPT and for the RL (reinforcement learning) algorithm are omitted.
+They can be included by using the targets `qutip-qoc[jopt]` and `qutip-qoc[rl]`, respectively (or `qutip-qoc[full]` for all dependencies).
+
 ## Documentation and tutorials
 
-The documentation of `qutip-qoc` updated to the latest development version is hosted at [qutip-qoc.readthedocs.io](https://qutip-qoc.readthedocs.io/en/latest/).
+The documentation of `qutip-qoc` updated to the latest development version is hosted at [qutip-qoc.readthedocs.io](https://qutip-qoc.readthedocs.io/latest/).
 Tutorials related to using quantum optimal control in `qutip-qoc` can be found [_here_](https://qutip.org/qutip-tutorials/#optimal-control).
 
 ## Installation from source
@@ -40,7 +43,7 @@ pip install --upgrade pip
 pip install -e .
 ```
 
-which makes sure that you are up to date with the latest `pip` version. Contribution guidelines are available [_here_](https://qutip-qoc.readthedocs.io/en/latest/contribution-code.html).
+which makes sure that you are up to date with the latest `pip` version. Contribution guidelines are available [_here_](https://qutip-qoc.readthedocs.io/latest/contribution/code.html).
 
 To build and test the documentation, additional packages need to be installed:
 
diff --git a/doc/conf.py b/doc/conf.py
@@ -62,6 +62,10 @@ def qutip_qoc_version():
 
 html_theme = "sphinx_rtd_theme"
 html_static_path = []
+html_js_files = [
+    'https://code.jquery.com/jquery-3.6.0.min.js',
+]
+
 
 # -- Options for numpydoc ---------------------------------------
 
diff --git a/doc/contribution/code.rst b/doc/contribution/code.rst
@@ -7,7 +7,7 @@ Contributing to the source code
 Build up an development environment
 ===================================
 
-Please follow the instruction on the `QuTiP contribution guide <https://qutip.org/docs/latest/development/contributing.html#building>`_ to
+Please follow the instruction on the `QuTiP contribution guide <https://qutip.readthedocs.io/en/latest/development/contributing.html>`_ to
 build a conda environment.
 
 You don't need to build ``qutip`` in the editable mode unless you also want to contribute to `qutip`.
diff --git a/doc/installation.rst b/doc/installation.rst
@@ -23,7 +23,19 @@ In particular, the following packages are necessary for running ``qutip-qoc``:
 
 .. code-block:: bash
 
-    numpy scipy jax jaxlib cython qutip qutip-jax qutip-qtrl
+    numpy scipy cython qutip qutip-qtrl
+
+The following packages are required for using the JOPT algorithm:
+
+.. code-block:: bash
+
+    jax jaxlib qutip-jax
+
+The following packages are required for the RL (reinforcement learning) algorithm:
+
+.. code-block:: bash
+
+    gymnasium stable-baselines3
 
 The following package is used for testing:
 
diff --git a/doc/requirements.txt b/doc/requirements.txt
@@ -9,18 +9,18 @@ matplotlib>=3.6.1
 docutils==0.18.1
 alabaster==0.7.12
 Babel==2.9.1
-certifi==2020.12.5
+certifi==2024.7.4
 chardet==4.0.0
 colorama==0.4.4
-idna==2.10
+idna==3.7
 imagesize==1.4.1
-Jinja2==3.0.1
+Jinja2==3.1.6
 MarkupSafe==2.0.1
 packaging==23.2
 Pygments==2.17.2
 pyparsing==2.4.7
 pytz==2021.1
-requests==2.25.1
+requests==2.32.2
 snowballstemmer==2.1.0
 Sphinx==6.1.3
 sphinx-gallery==0.12.2
@@ -32,4 +32,4 @@ sphinxcontrib-htmlhelp==2.0.1
 sphinxcontrib-jsmath==1.0.1
 sphinxcontrib-qthelp==1.0.3
 sphinxcontrib-serializinghtml==1.1.5
-urllib3==1.26.4
+urllib3==1.26.19
diff --git a/requirements.txt b/requirements.txt
@@ -1,11 +1,6 @@
 cython>=1.0
 numpy>=1.16.6,<2.0
 scipy>=1.10.1
-jax==0.4.28
-jaxlib==0.4.28
 qutip>=5.0.1
 qutip-qtrl
-qutip-jax
 pre-commit
-gymnasium>=0.29.1
-stable-baselines3>=2.3.2
diff --git a/setup.cfg b/setup.cfg
@@ -28,15 +28,10 @@ package_dir=
 packages = find:
 include_package_data = True
 install_requires =
-    jax
-    jaxlib
     packaging
     qutip
     qutip-qtrl
-    qutip-jax
     numpy>=1.16.6,<2.0
-    gymnasium>=0.29.1
-    stable-baselines3>=2.3.2
 setup_requires =
     cython>=1.0
     packaging
@@ -45,7 +40,16 @@ setup_requires =
 where = src
 
 [options.extras_require]
+jopt =
+    jax
+    jaxlib
+    qutip-jax
+rl =
+    gymnasium>=0.29.1
+    stable-baselines3>=2.3.2
 tests =
     pytest>=6.0
 full =
+    %(jopt)s
+    %(rl)s
     %(tests)s
diff --git a/src/qutip_qoc/_jopt.py b/src/qutip_qoc/_jopt.py
@@ -5,38 +5,47 @@
 import qutip as qt
 from qutip import Qobj, QobjEvo
 
-from diffrax import Dopri5, PIDController
+try:
+    import jax
+    from jax import custom_jvp
+    import jax.numpy as jnp
+    import qutip_jax  # noqa: F401
 
-import jax
-from jax import custom_jvp
-import jax.numpy as jnp
-import qutip_jax  # noqa: F401
+    import jaxlib  # noqa: F401
 
+    from diffrax import Dopri5, PIDController
 
-@custom_jvp
-def _abs(x):
-    return jnp.abs(x)
+    _jax_available = True
+except ImportError:
+    _jax_available = False
 
 
-def _abs_jvp(primals, tangents):
-    """
-    Custom jvp for absolute value of complex functions
-    """
-    (x,) = primals
-    (t,) = tangents
+if _jax_available:
+
+    @custom_jvp
+    def _abs(x):
+        return jnp.abs(x)
+
 
-    abs_x = _abs(x)
-    res = jnp.where(
-        abs_x == 0,
-        0.0,  # prevent division by zero
-        jnp.real(jnp.multiply(jnp.conj(x), t)) / abs_x,
-    )
+    def _abs_jvp(primals, tangents):
+        """
+        Custom jvp for absolute value of complex functions
+        """
+        (x,) = primals
+        (t,) = tangents
+
+        abs_x = _abs(x)
+        res = jnp.where(
+            abs_x == 0,
+            0.0,  # prevent division by zero
+            jnp.real(jnp.multiply(jnp.conj(x), t)) / abs_x,
+        )
 
-    return abs_x, res
+        return abs_x, res
 
 
-# register custom jvp for absolut value of complex functions
-_abs.defjvp(_abs_jvp)
+    # register custom jvp for absolut value of complex functions
+    _abs.defjvp(_abs_jvp)
 
 
 class _JOPT:
@@ -55,6 +64,10 @@ def __init__(
         guess_params,
         **integrator_kwargs,
     ):
+        if not _jax_available:
+            raise ImportError("The JOPT algorithm requires the modules jax, "
+                              "jaxlib, and qutip_jax to be installed.")
+
         self._Hd = objective.H[0]
         self._Hc_lst = objective.H[1:]
 
@@ -137,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)
@@ -147,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
diff --git a/src/qutip_qoc/objective.py b/src/qutip_qoc/objective.py
@@ -80,12 +80,6 @@ def __init__(self, initial, H, target, weight=1):
         self.target = target
         self.weight = weight
 
-        # Check if any Hamiltonian in H is a superoperator
-        if any(qt.issuper(H_i) for H_i in (H if isinstance(H, list) else [H])):
-            # Convert initial and target accordingly
-            self.initial = qt.to_super(self.initial)
-            self.target = qt.to_super(self.target)
-
     def __getstate__(self):
         """
         Extract picklable information from the objective.
diff --git a/src/qutip_qoc/result.py b/src/qutip_qoc/result.py
@@ -2,7 +2,6 @@
 This module contains the Result class for storing and
 reporting the results of a full pulse control optimization run.
 """
-import jaxlib
 import pickle
 import textwrap
 import numpy as np
@@ -11,6 +10,13 @@
 
 import qutip as qt
 
+try:
+    import jax
+    import jaxlib
+    _jitfun_type = type(jax.jit(lambda x: x))
+except ImportError:
+    _jitfun_type = None
+
 __all__ = ["Result"]
 
 
@@ -331,8 +337,10 @@ def final_states(self):
                 evo_time = self.time_interval.evo_time
 
             # choose solver method based on type of control function
-            if isinstance(
-                self.objectives[0].H[1][1], jaxlib.xla_extension.PjitFunction
+            # if jax is installed, _jitfun_type is set to
+            # jaxlib.xla_extension.PjitFunction, otherwise it is None
+            if _jitfun_type is not None and isinstance(
+                self.objectives[0].H[1][1], _jitfun_type
             ):
                 method = "diffrax"  # for JAX defined contols
             else:
diff --git a/tests/test_analytical_pulses.py b/tests/test_analytical_pulses.py
diff --git a/tests/test_fidelity.py b/tests/test_fidelity.py
diff --git a/tests/test_jopt_open_system_bug.py b/tests/test_jopt_open_system_bug.py
diff --git a/tests/test_result.py b/tests/test_result.py
