Merge pull request #12 from Ericgig/add_context

Add a contextmanager and reverse default option

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.

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

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"]