Merge pull request qutip#2341 from nwlambert/enr_fixes

Making ENR states work with mesolve()

qutip/core/cy/qobjevo.pyx

@@ -613,11 +613,11 @@ cdef class QobjEvo:
 
     def __imatmul__(QobjEvo self, other):
         if isinstance(other, (Qobj, QobjEvo)):
-            if self.dims[1] != other.dims[0]:
+            if self._dims[1] != other._dims[0]:
                 raise TypeError("incompatible dimensions" +
                                 str(self.dims[1]) + ", " +
                                 str(other.dims[0]))
-            self._dims = Dimensions([self.dims[0], other.dims[1]])
+            self._dims = Dimensions([self._dims[0], other._dims[1]])
             self.shape = (self.shape[0], other.shape[1])
             if isinstance(other, Qobj):
                 other = _ConstantElement(other)

qutip/core/qobj.py

@@ -312,7 +312,7 @@ def __init__(self, arg=None, dims=None,
     def copy(self):
         """Create identical copy"""
         return Qobj(arg=self._data,
-                    dims=self.dims,
+                    dims=self._dims,
                     isherm=self._isherm,
                     isunitary=self._isunitary,
                     copy=True)

qutip/core/superoperator.py

@@ -12,6 +12,7 @@
 from . import data as _data
 from .dimensions import Compound, SuperSpace, Space
 
+
 def _map_over_compound_operators(f):
     """
     Convert a function which takes Qobj into one that can also take compound
@@ -49,7 +50,7 @@ def liouvillian(H=None, c_ops=None, data_only=False, chi=None):
         In some systems it is possible to determine the statistical moments
         (mean, variance, etc) of the probability distributions of occupation of
         various states by numerically evaluating the derivatives of the steady
-        state occupation probability as a function of artificial phase 
+        state occupation probability as a function of artificial phase
         parameters ``chi`` which are included in the
         :func:`lindblad_dissipator` for each collapse operator. See the
         documentation of :func:`lindblad_dissipator` for references and further
@@ -95,15 +96,10 @@ def liouvillian(H=None, c_ops=None, data_only=False, chi=None):
         L += sum(lindblad_dissipator(c_op, chi=chi_)
                  for c_op, chi_ in zip(c_ops, chi))
         return L
-
-    op_dims = H.dims
-    op_shape = H.shape
-    sop_dims = [[op_dims[0], op_dims[0]], [op_dims[1], op_dims[1]]]
-    sop_shape = [np.prod(op_dims), np.prod(op_dims)]
-    spI = _data.identity(op_shape[0], dtype=type(H.data))
-
+    spI = _data.identity_like(H.data)
     data = _data.mul(_data.kron(spI, H.data), -1j)
-    data = _data.add(data, _data.kron_transpose(H.data, spI), scale=1j)
+    data = _data.add(data, _data.kron_transpose(H.data, spI),
+                     scale=1j)
 
     for c_op, chi_ in zip(c_ops, chi):
         c = c_op.data
@@ -117,7 +113,7 @@ def liouvillian(H=None, c_ops=None, data_only=False, chi=None):
         return data
     else:
         return Qobj(data,
-                    dims=sop_dims,
+                    dims=[H._dims, H._dims],
                     superrep='super',
                     copy=False)
 
@@ -316,10 +312,9 @@ def spost(A):
     """
     if not A.isoper:
         raise TypeError('Input is not a quantum operator')
-    Id = _data.identity(A.shape[0], dtype=type(A.data))
     data = _data.kron_transpose(A.data, _data.identity_like(A.data))
     return Qobj(data,
-                dims=[A.dims, A.dims],
+                dims=[A._dims, A._dims],
                 superrep='super',
                 isherm=A._isherm,
                 copy=False)

qutip/solver/mcsolve.py

@@ -463,7 +463,7 @@ def _restore_state(self, data, *, copy=True):
         """
         Retore the Qobj state from its data.
         """
-        if self._state_metadata['dims'] == self.rhs.dims[1]:
+        if self._state_metadata['dims'] == self.rhs._dims[1]:
             state = Qobj(unstack_columns(data),
                          **self._state_metadata, copy=False)
         else:

qutip/solver/solver_base.py

@@ -84,7 +84,7 @@ def _prepare_state(self, state):
                             f" and {state.dims}")
 
         self._state_metadata = {
-            'dims': state.dims,
+            'dims': state._dims,
             'isherm': state.isherm and not (self.rhs.dims == state.dims)
         }
         if self.rhs.dims[1] == state.dims:
@@ -95,7 +95,7 @@ def _restore_state(self, data, *, copy=True):
         """
         Retore the Qobj state from its data.
         """
-        if self._state_metadata['dims'] == self.rhs.dims[1]:
+        if self._state_metadata['dims'] == self.rhs._dims[1]:
             state = Qobj(unstack_columns(data),
                          **self._state_metadata, copy=False)
         else:

qutip/solver/stochastic.py

@@ -4,6 +4,7 @@
 from .result import MultiTrajResult, Result, ExpectOp
 from .multitraj import MultiTrajSolver
 from .. import Qobj, QobjEvo
+from ..core.dimensions import Dimensions
 import numpy as np
 from functools import partial
 from .solver_base import _solver_deprecation
@@ -210,8 +211,10 @@ def __init__(self, issuper, H, sc_ops, c_ops, heterodyne):
 
         if self.issuper and not self.H.issuper:
             self.dims = [self.H.dims, self.H.dims]
+            self._dims = Dimensions([self.H._dims, self.H._dims])
         else:
             self.dims = self.H.dims
+            self._dims = self.H._dims
 
     def __call__(self, options):
         if self.issuper:

qutip/tests/test_enr_state_operator.py

@@ -114,3 +114,41 @@ def test_thermal_dm(dimensions, n_excitations, nbar_type):
     test_dm = qutip.enr_thermal_dm(dimensions, n_excitations, nbars)
     expect_dm = _reference_dm(dimensions, n_excitations, nbars)
     np.testing.assert_allclose(test_dm.full(), expect_dm.full(), atol=1e-12)
+
+
+def test_mesolve_ENR():
+    # Ensure ENR states work with mesolve
+    # We compare the output to an exact truncation of the
+    # single-excitation Jaynes-Cummings model
+    eps = 2 * np.pi
+    omega_c = 2 * np.pi
+    g = 0.1 * omega_c
+    gam = 0.01 * omega_c
+    tlist = np.linspace(0, 20, 100)
+    N_cut = 2
+
+    sz = qutip.sigmaz() & qutip.qeye(N_cut)
+    sm = qutip.destroy(2).dag() & qutip.qeye(N_cut)
+    a = qutip.qeye(2) & qutip.destroy(N_cut)
+    H_JC = (0.5 * eps * sz + omega_c * a.dag()*a +
+            g * (a * sm.dag() + a.dag() * sm))
+    psi0 = qutip.basis(2, 0) & qutip.basis(N_cut, 0)
+    c_ops = [np.sqrt(gam) * a]
+
+    result_JC = qutip.mesolve(H_JC, psi0, tlist, c_ops, e_ops=[sz])
+
+    N_exc = 1
+    dims = [2, N_cut]
+    d = qutip.enr_destroy(dims, N_exc)
+    sz = 2*d[0].dag()*d[0]-1
+    b = d[0]
+    a = d[1]
+    psi0 = qutip.enr_fock(dims, N_exc, [1, 0])
+    H_enr = (eps * b.dag()*b + omega_c * a.dag() * a +
+             g * (b.dag() * a + a.dag() * b))
+    c_ops = [np.sqrt(gam) * a]
+
+    result_enr = qutip.mesolve(H_enr, psi0, tlist, c_ops, e_ops=[sz])
+
+    np.testing.assert_allclose(result_JC.expect[0],
+                               result_enr.expect[0], atol=1e-2)