Merge pull request #19 from Ericgig/prepare0.1.2

Prepare 0.1.2

Author: Ericgig

README.md

@@ -30,7 +30,7 @@ see [nvidia's documentation](https://docs.nvidia.com/cuda/cuquantum/latest/getti
 
 ## Usage
 
-In simple case, simply calling `set_as_default` before a qutip script should be sufficient to use the backend common solver:
+In simple case, simply calling `set_as_default` before a qutip script should be sufficient to use the backend in common solver:
 
 ```
 import qutip_cuquantum
@@ -39,5 +39,12 @@ from cuquantum.densitymat import WorkStream
 qutip_cuquantum.set_as_default(WorkStream())
 ```
 
+It can also be used as a context:
+
+```
+with CuQuantumBackend(ctx):
+    ...
+```
+
 qutip-cuquantum work well to speed-up large simulation using `mesolve` or `sesolve`.
 However this backend is not compatible with advanced qutip solvers (brmesolve, HEOM) and other various feature.

VERSION

@@ -1 +1 @@
-0.1.1
+0.1.2

doc/release_distribution.md

@@ -0,0 +1,24 @@
+# Release Procedure for qutip-cuquantum
+
+#### Major and Minor Releases (e.g., v0.1.0, v0.2.0)
+
+- Create a new branch named qutip-cuquantum-A.B.X, where A.B is the major/minor version number.
+
+- Update the VERSION file for Release. Remove the .dev suffix from the version number. `alpha`, `beta`, `post` tag are acceptable.
+
+- Increment Version on main. Edit the VERSION file and increment the version number, resetting the dev number. For example, if you are releasing `v0.1.0`, the version on main should become `v0.2.0.dev1`.
+
+
+#### Micro/Patch Releases (e.g., v0.1.1)
+
+- Make a branch off the release branch.
+- Cherry-Pick changes to include from main.
+- Increment the VERSION file
+- Merge into the release branch with a pull request.
+
+#### Deployment
+
+- Go to the "Releases" page on the GitHub repository and draft a new release.
+- Create a new tag with the scheme vA.B.C (e.g., v0.1.0) targeting your release branch (qutip-cuquantum-0.1.X).
+- Write the changelog into the release description.
+- Publishing the release will trigger the GitHub Action to build and publish the package to PyPI.

doc/source/solver.rst

@@ -32,9 +32,23 @@ This is done by calling the ``set_as_default`` function and providing it with a
 The ``set_as_default`` function changes several QuTiP defaults to route computations through the cuQuantum library.
 This includes setting the default data format for quantum objects (``Qobj``) to ``CuOperator`` and configuring the solvers to use GPU-compatible integrators.
 
-.. warning::
-    This operation is **not reversible** within the same Python session.
-    Once the cuQuantum backend is set, all subsequent compatible operations will be dispatched to the GPU.
+This operation can be reversed with:
+
+.. code-block:: python
+
+    qutip_cuquantum.set_as_default(reverse=True)
+
+
+The backend can also be enabled with a context:
+
+.. code-block:: python
+
+    with CuQuantumBackend(ctx):
+        ...
+
+However be careful when mixing core Qutip object and Qutip-cuQuantum's one.
+Qutip's Qobj do not keep all the internal structure needed for cuQuantum's optimizations.
+Qutip-cuQuantum's states can be distributed in multiple processes and unusable for many qutip's core features. 
 
 ==================================
 Usage with Solvers

setup.py

@@ -9,6 +9,7 @@
 import setuptools
 from Cython.Build import cythonize
 import qutip
+import qutip.core.cy as qutip_cc
 import numpy
 
 
@@ -94,7 +95,9 @@ def get_ext_modules(options):
     pyx_file = os.path.join("src", "qutip_cuquantum", "qobjevo.pyx")
     include_dirs = [
         numpy.get_include(),
-        os.path.abspath(os.path.join(qutip.core.data.__file__, os.pardir))
+        os.path.abspath(os.path.join(qutip.core.data.__file__, os.pardir)),
+        os.path.abspath(os.path.join(qutip_cc.__file__, os.pardir)),
+        os.path.abspath(os.path.join(qutip.__file__, os.pardir))
     ]
     print("*********************************************************************************")
     print(include_dirs)
@@ -107,7 +110,7 @@ def get_ext_modules(options):
         language="c++",
     )
 
-    return cythonize(ext)
+    return cythonize(ext, include_path=include_dirs)
 
 
 if __name__ == "__main__":

src/qutip_cuquantum/__init__.py

@@ -27,6 +27,7 @@
 from qutip.core.options import QutipOptions
 from .operator import CuOperator
 from .state import CuState
+import numpy
 
 
 # TODO: The split per density is not great
@@ -62,6 +63,9 @@ def CuState_from_CuPyDense(mat):
 
 from .qobjevo import CuQobjEvo
 from .ode import Result, CuMCIntegrator
+from qutip import settings
+from qutip.solver import SESolver, MESolver, MCSolver, Result as BaseResult
+from qutip.solver.mcsolve import MCIntegrator
 
 
 class cuDensityOption(QutipOptions):
@@ -74,7 +78,7 @@ class cuDensityOption(QutipOptions):
 cuDensityOption_instance._set_as_global_default()
 
 
-def set_as_default(ctx: cuquantum.densitymat.WorkStream):
+def set_as_default(ctx: cuquantum.densitymat.WorkStream=None, reverse=False):
     """
     Update qutip's default to use cuQuantum as a backend.
 
@@ -83,22 +87,97 @@ def set_as_default(ctx: cuquantum.densitymat.WorkStream):
     ctx: WorkStream
         A WorkStream instance from cuquantum.density.
         It can be set with mpi support for multi-gpu simulations.
+        Can be ignored when ``reverse=True``.
+
+    reverse: bool, default: False
+        Undo the change of default backend to qutip core defaults.
     """
-    qutip.settings.cuDensity["ctx"] = ctx
-    qutip.settings.core["default_dtype"] = "cuDensity"
-    qutip.settings.core['numpy_backend'] = cupy
+    if not reverse:
+        settings.cuDensity["ctx"] = ctx
+        settings.core["default_dtype"] = "cuDensity"
+        settings.core['numpy_backend'] = cupy
+
+        if True:  # if mpi, how to check from ctx?
+            settings.core["auto_real_casting"] = False
+
+        SESolver.solver_options['method'] = "CuVern7"
+        MESolver.solver_options['method'] = "CuVern7"
+        MCSolver.solver_options['method'] = "CuVern7"
+
+        SESolver._resultclass = Result
+        MESolver._resultclass = Result
+        MCSolver._trajectory_resultclass = Result
+        MCSolver._mc_integrator_class = CuMCIntegrator
+
+    else:
+        settings.core["default_dtype"] = "core"
+        settings.core['numpy_backend'] = numpy
+        settings.core["auto_real_casting"] = True
+
+        SESolver.solver_options['method'] = "adams"
+        MESolver.solver_options['method'] = "adams"
+        MCSolver.solver_options['method'] = "vern7"
+
+        SESolver._resultclass = BaseResult
+        MESolver._resultclass = BaseResult
+        MCSolver._trajectory_resultclass = BaseResult
+        MCSolver._mc_integrator_class = MCIntegrator
 
-    if True:  # if mpi, how to check from ctx?
-        qutip.settings.core["auto_real_casting"] = False
 
-    qutip.SESolver.solver_options['method'] = "CuVern7"
-    qutip.MESolver.solver_options['method'] = "CuVern7"
-    qutip.MCSolver.solver_options['method'] = "CuVern7"
 
-    qutip.SESolver._resultclass = Result
-    qutip.MESolver._resultclass = Result
-    qutip.MCSolver._trajectory_resultclass = Result
-    qutip.MCSolver._mc_integrator_class = CuMCIntegrator
+class CuQuantumBackend:
+    """
+    A context manager class to temporarily set cuQuantum as the default
+    backend.
+
+    Parameters
+    ----------
+    ctx : cuquantum.densitymat.WorkStream
+        A WorkStream instance from cuquantum.density.
+        It can be set with mpi support for multi-gpu simulations.
+    """
+    def __init__(self, ctx):
+        self.ctx = ctx
+        self.previous_values = {}
+
+    def __enter__(self):
+        settings.cuDensity["ctx"] = self.ctx
+        self.previous_values["default_dtype"] = qutip.settings.core["default_dtype"]
+        settings.core["default_dtype"] = "cuDensity"
+        self.previous_values["numpy_backend"] = qutip.settings.core["numpy_backend"]
+        settings.core['numpy_backend'] = cupy
+
+        self.previous_values["auto_real"] = settings.core["auto_real_casting"]
+        if True:  # if mpi, how to check from ctx?
+            settings.core["auto_real_casting"] = False
+
+        self.previous_values["SESolverM"] = SESolver.solver_options['method']
+        self.previous_values["MESolverM"] = MESolver.solver_options['method']
+        self.previous_values["MCSolverM"] = MCSolver.solver_options['method']
+        SESolver.solver_options['method'] = "CuVern7"
+        MESolver.solver_options['method'] = "CuVern7"
+        MCSolver.solver_options['method'] = "CuVern7"
+
+        self.previous_values["SESolverR"] = SESolver._resultclass
+        self.previous_values["MESolverR"] = MESolver._resultclass
+        self.previous_values["MCSolverR"] = MCSolver._trajectory_resultclass
+        self.previous_values["MCSolverI"] = MCSolver._mc_integrator_class
+        SESolver._resultclass = Result
+        MESolver._resultclass = Result
+        MCSolver._trajectory_resultclass = Result
+        MCSolver._mc_integrator_class = CuMCIntegrator
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        settings.core["default_dtype"] = self.previous_values["default_dtype"]
+        settings.core['numpy_backend'] = self.previous_values["numpy_backend"]
+        settings.core["auto_real_casting"] = self.previous_values["auto_real"]
+        SESolver.solver_options['method'] = self.previous_values["SESolverM"]
+        MESolver.solver_options['method'] = self.previous_values["MESolverM"]
+        MCSolver.solver_options['method'] = self.previous_values["MCSolverM"]
+        SESolver._resultclass = self.previous_values["SESolverR"]
+        MESolver._resultclass = self.previous_values["MESolverR"]
+        MCSolver._trajectory_resultclass = self.previous_values["MCSolverR"]
+        MCSolver._mc_integrator_class = self.previous_values["MCSolverI"]
 
 
 

src/qutip_cuquantum/mixed_dispatch.py

@@ -0,0 +1,30 @@
+import qutip.core.data as _data
+from qutip import settings
+from .state import zeros_like_cuState, CuState
+from .operator import CuOperator
+from .utils import _compare_hilbert
+
+import cuquantum.densitymat as cudense
+from cuquantum.densitymat import Operator
+
+@_data.matmul.register(CuOperator, CuState, CuState)
+def matmul_cuoperator_custate_custate(left, right, scale=1., out=None):
+    merged_hilbert = _compare_hilbert(left.hilbert_dims, right.base.hilbert_space_dims)
+    if not merged_hilbert:
+        raise ValueError("Hilbert space missmatch")
+    if left.shape[1] == right.shape[0]:
+        dual = False
+    elif left.shape[1] == right.shape[0] * right.shape[1]:
+        dual = True
+    else:
+        raise ValueError("Shape missmatch")
+
+    oper = Operator(merged_hilbert, [left.to_OperatorTerm(dual=dual, hilbert_dims=merged_hilbert)])
+
+    oper.prepare_action(settings.cuDensity["ctx"], right.base)
+    if out is None:
+        out = zeros_like_cuState(right)
+
+    oper.compute_action(0, [], state_in=right.base, state_out=out.base)
+
+    return out