Merge pull request #2427 from Ericgig/misc.solver.state0_norm

Only normalize if initial state is normalized.

doc/changes/2427.misc

@@ -0,0 +1 @@
+Only normalize solver states when the initial state is already normalized.

qutip/solver/brmesolve.py

@@ -102,7 +102,8 @@ def brmesolve(H, psi0, tlist, a_ops=(), e_ops=(), c_ops=(),
             On `None` the states will be saved if no expectation operators are
             given.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/floquet.py

@@ -535,7 +535,8 @@ def fsesolve(H, psi0, tlist, e_ops=None, T=0.0, args=None, options=None):
             On `None` the states will be saved if no expectation operators are
             given.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
 
     Returns
     -------
@@ -638,7 +639,8 @@ def fmmesolve(
           | Whether or not to store the density matrices in the floquet basis
             in ``result.floquet_states``.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/heom/bofin_solvers.py

@@ -498,7 +498,8 @@ def heomsolve(
         - | store_ados : bool
           | Whether or not to store the HEOM ADOs.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/krylovsolve.py

@@ -70,7 +70,8 @@ def krylovsolve(
             On `None` the states will be saved if no expectation operators are
             given.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/mesolve.py

@@ -107,7 +107,8 @@ def mesolve(
             On `None` the states will be saved if no expectation operators are
             given.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/sesolve.py

@@ -82,7 +82,8 @@ def sesolve(
             On `None` the states will be saved if no expectation operators are
             given.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/solver/solver_base.py

@@ -6,12 +6,14 @@
 from .. import Qobj, QobjEvo, ket2dm
 from .options import _SolverOptions
 from ..core import stack_columns, unstack_columns
+from .. import settings
 from .result import Result
 from .integrator import Integrator
 from ..ui.progressbar import progress_bars
 from ._feedback import _ExpectFeedback
 from time import time
 import warnings
+import numpy as np
 
 
 class Solver:
@@ -93,6 +95,13 @@ def _prepare_state(self, state):
             # anything other than dimensions.
             'isherm': state.isherm and not (self.rhs.dims == state.dims)
         }
+        if state.isoper:
+            norm = state.tr()
+        else:
+            norm = state.norm()
+        # Use the settings atol instead of the solver one since the second
+        # refer to the ODE tolerance and some integrator do not use it.
+        self._normalized = np.abs(norm - 1) <= settings.core["atol"]
         if self.rhs.dims[1] == state.dims:
             return stack_columns(state.data)
         return state.data
@@ -107,7 +116,11 @@ def _restore_state(self, data, *, copy=True):
         else:
             state = Qobj(data, **self._state_metadata, copy=copy)
 
-        if data.shape[1] == 1 and self._options['normalize_output']:
+        if (
+            data.shape[1] == 1
+            and self._options['normalize_output']
+            and self._normalized
+        ):
             if state.isoper:
                 state = state * (1 / state.tr())
             else:

qutip/solver/stochastic.py

@@ -359,7 +359,8 @@ def smesolve(
           | Whether to store results from all trajectories or just store the
             averages.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error
@@ -483,7 +484,8 @@ def ssesolve(
           | Whether to store results from all trajectories or just store the
             averages.
         - | normalize_output : bool
-          | Normalize output state to hide ODE numerical errors.
+          | Normalize output state to hide ODE numerical errors. Only normalize
+            the state if the initial state is already normalized.
         - | progress_bar : str {'text', 'enhanced', 'tqdm', ''}
           | How to present the solver progress.
             'tqdm' uses the python module of the same name and raise an error

qutip/tests/solver/test_mesolve.py

@@ -698,3 +698,15 @@ def f(t, A):
     solver = qutip.MESolver(H, c_ops=[a])
     result = solver.run(psi0, np.linspace(0, 30, 301), e_ops=[qutip.num(N)])
     assert np.all(result.expect[0] > 4. - tol)
+
+
+@pytest.mark.parametrize(
+    'rho0',
+    [qutip.sigmax(), qutip.sigmaz(), qutip.qeye(2)],
+    ids=["sigmax", "sigmaz", "tr=2"]
+)
+def test_non_normalized_dm(rho0):
+    H = qutip.QobjEvo(qutip.num(2))
+    solver = qutip.MESolver(H, c_ops=[qutip.sigmaz()])
+    result = solver.run(rho0, np.linspace(0, 1, 10), e_ops=[qutip.qeye(2)])
+    np.testing.assert_allclose(result.expect[0], rho0.tr(), atol=1e-7)