From a57f63adefc70000f1cfe3326b33615cb51d0c96 Mon Sep 17 00:00:00 2001 From: Eric Giguere Date: Fri, 10 Oct 2025 10:58:52 -0400 Subject: [PATCH 1/4] Add context manager --- src/qutip_cuquantum/__init__.py | 105 ++++++++++++++++++++++++++++---- 1 file changed, 92 insertions(+), 13 deletions(-) diff --git a/src/qutip_cuquantum/__init__.py b/src/qutip_cuquantum/__init__.py index 2a5e5cf..34a1696 100644 --- a/src/qutip_cuquantum/__init__.py +++ b/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 + previous_values = {} + + def __enter__(self, ctx): + settings.cuDensity["ctx"] = self.ctx + previous_values{"default_dtype"} = qutip.settings.core["default_dtype"] + settings.core["default_dtype"] = "cuDensity" + previous_values{"numpy_backend"} = qutip.settings.core["numpy_backend"] + settings.core['numpy_backend'] = cupy + + previous_values{"auto_real"} = settings.core["auto_real_casting"] + if True: # if mpi, how to check from ctx? + settings.core["auto_real_casting"] = False + + previous_values{"SESolverM"} = SESolver.solver_options['method'] + previous_values{"MESolverM"} = MESolver.solver_options['method'] + previous_values{"MCSolverM"} = MCSolver.solver_options['method'] + SESolver.solver_options['method'] = "CuVern7" + MESolver.solver_options['method'] = "CuVern7" + MCSolver.solver_options['method'] = "CuVern7" + + previous_values{"SESolverR"} = MCSolver._resultclass + previous_values{"MESolverR"} = MCSolver._resultclass + previous_values{"MCSolverR"} = MCSolver._trajectory_resultclass + 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"] = previous_values{"default_dtype"} + settings.core['numpy_backend'] = previous_values{"numpy_backend"} + settings.core["auto_real_casting"] = previous_values{"auto_real"} + SESolver.solver_options['method'] = previous_values{"SESolverM"} + MESolver.solver_options['method'] = previous_values{"MESolverM"} + MCSolver.solver_options['method'] = previous_values{"MCSolverM"} + SESolver._resultclass = previous_values{"SESolverR"} + MESolver._resultclass = previous_values{"MESolverR"} + MCSolver._trajectory_resultclass = previous_values{"MCSolverR"} + MCSolver._mc_integrator_class = previous_values{"MCSolverI"} From c68bd578d079e50d6851ffb519e8c2c11229fe01 Mon Sep 17 00:00:00 2001 From: Eric Giguere Date: Fri, 10 Oct 2025 13:01:49 -0400 Subject: [PATCH 2/4] Fix typo --- src/qutip_cuquantum/__init__.py | 46 ++++++++++++++++----------------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/src/qutip_cuquantum/__init__.py b/src/qutip_cuquantum/__init__.py index 34a1696..edee656 100644 --- a/src/qutip_cuquantum/__init__.py +++ b/src/qutip_cuquantum/__init__.py @@ -64,7 +64,7 @@ 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 import SESolver, MESolver, MCSolver, Result as BaseResult from qutip.solver.mcsolve import MCIntegrator @@ -138,46 +138,46 @@ class CuQuantumBackend: """ def __init__(self, ctx): self.ctx = ctx - previous_values = {} + self.previous_values = {} - def __enter__(self, ctx): + def __enter__(self): settings.cuDensity["ctx"] = self.ctx - previous_values{"default_dtype"} = qutip.settings.core["default_dtype"] + self.previous_values["default_dtype"] = qutip.settings.core["default_dtype"] settings.core["default_dtype"] = "cuDensity" - previous_values{"numpy_backend"} = qutip.settings.core["numpy_backend"] + self.previous_values["numpy_backend"] = qutip.settings.core["numpy_backend"] settings.core['numpy_backend'] = cupy - previous_values{"auto_real"} = settings.core["auto_real_casting"] + 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 - previous_values{"SESolverM"} = SESolver.solver_options['method'] - previous_values{"MESolverM"} = MESolver.solver_options['method'] - previous_values{"MCSolverM"} = MCSolver.solver_options['method'] + 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" - previous_values{"SESolverR"} = MCSolver._resultclass - previous_values{"MESolverR"} = MCSolver._resultclass - previous_values{"MCSolverR"} = MCSolver._trajectory_resultclass - previous_values{"MCSolverI"} = MCSolver._mc_integrator_class + self.previous_values["SESolverR"] = MCSolver._resultclass + self.previous_values["MESolverR"] = MCSolver._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"] = previous_values{"default_dtype"} - settings.core['numpy_backend'] = previous_values{"numpy_backend"} - settings.core["auto_real_casting"] = previous_values{"auto_real"} - SESolver.solver_options['method'] = previous_values{"SESolverM"} - MESolver.solver_options['method'] = previous_values{"MESolverM"} - MCSolver.solver_options['method'] = previous_values{"MCSolverM"} - SESolver._resultclass = previous_values{"SESolverR"} - MESolver._resultclass = previous_values{"MESolverR"} - MCSolver._trajectory_resultclass = previous_values{"MCSolverR"} - MCSolver._mc_integrator_class = previous_values{"MCSolverI"} + 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"] From 93226c01a6c52b4f2581f7fb907f049334541cff Mon Sep 17 00:00:00 2001 From: Eric Giguere Date: Fri, 10 Oct 2025 13:51:50 -0400 Subject: [PATCH 3/4] Fix typo --- src/qutip_cuquantum/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/qutip_cuquantum/__init__.py b/src/qutip_cuquantum/__init__.py index edee656..b6009da 100644 --- a/src/qutip_cuquantum/__init__.py +++ b/src/qutip_cuquantum/__init__.py @@ -158,8 +158,8 @@ def __enter__(self): MESolver.solver_options['method'] = "CuVern7" MCSolver.solver_options['method'] = "CuVern7" - self.previous_values["SESolverR"] = MCSolver._resultclass - self.previous_values["MESolverR"] = MCSolver._resultclass + 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 From 1badaad6fb4af57a4831e24aed6367b5f26c5c22 Mon Sep 17 00:00:00 2001 From: Eric Giguere Date: Thu, 16 Oct 2025 09:51:38 -0400 Subject: [PATCH 4/4] Add documentation --- README.md | 9 ++++++++- doc/source/solver.rst | 20 +++++++++++++++++--- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 312fe2f..1d29a4f 100644 --- a/README.md +++ b/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. diff --git a/doc/source/solver.rst b/doc/source/solver.rst index e068544..757b467 100644 --- a/doc/source/solver.rst +++ b/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