Merge pull request #17 from Ericgig/better_oper

Rework to operterm

Author: Ericgig

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/operator.py

@@ -26,17 +26,32 @@
 # Input fixed in tests, but did not fail early
 
 
-def _transpose_cu_operator(oper):
-    if isinstance(oper, MultidiagonalOperator):
-        out = MultidiagonalOperator(
-            oper.data,
-            [-offset for offset in oper.offsets],
-            callback=oper.callback,
-        )
+def _transpose_cu_operator(oper, transpose: list[bool]):
+    """
+    Transpose modes where ``transpose`` is True.
+    """
+    if isinstance(oper, MultidiagonalOperator) and len(transpose) == 1:
+        if transpose[0]:
+            out = MultidiagonalOperator(
+                oper.data,
+                [-offset for offset in oper.offsets],
+                callback=oper.callback,
+            )
     elif isinstance(oper, DenseOperator):
         N = oper.num_modes
         batch_dims_oper = len(oper.data.shape) % 2
-        perm = tuple(range(N, 2*N)) + tuple(range(N))
+
+        perm_l = []
+        perm_r = []
+        for i, trans in enumerate(transpose):
+            if trans:
+                perm_l.append(i + N)
+                perm_r.append(i)
+            else:
+                perm_l.append(i)
+                perm_r.append(i+ N)
+        perm = tuple(perm_l + perm_r)
+
         new_callback = None
 
         if oper.callback is not None:
@@ -102,7 +117,8 @@ def _oper_to_ElementaryOperator(
     oper,
     hilbert_idx,
     hilbert_dims,
-    transform,
+    transpose,
+    dag,
     copy=False
 ):
     N = len(hilbert_idx)
@@ -111,32 +127,27 @@ def _oper_to_ElementaryOperator(
     if isinstance(oper, (DenseOperator, MultidiagonalOperator)):
         if N != 1 and isinstance(oper, MultidiagonalOperator):
             raise ValueError(
-                "MultidiagonalOperator on multiple hilbert spaces"
+                "MultidiagonalOperator on multiple hilbert spaces are not supported."
             )
         if list(oper.shape[:len(oper.shape) // 2]) != list(shape):
             raise ValueError(
                 f"Operator shape does not match hilbert spaces: "
                 f"{list(oper.shape[:len(oper.shape) // 2])}, {shape}"
             )
 
-        if transform == Transform.DIRECT:
-            out = oper
-        elif transform == Transform.ADJOINT:
-            out = oper.dag()
-        elif transform == Transform.CONJ:
-            out = _transpose_cu_operator(oper).dag()
-        elif transform == Transform.TRANSPOSE:
-            out = _transpose_cu_operator(oper)
+        cu_oper = oper
 
     else:
-        if transform == Transform.DIRECT:
-            pass
-        elif transform == Transform.ADJOINT:
-            oper = oper.adjoint()
-        elif transform == Transform.CONJ:
+        if all(transpose) and dag:
             oper = oper.conj()
-        elif transform == Transform.TRANSPOSE:
+            transpose = [False]
+            dag = False
+        elif dag:
+            oper = oper.adjoint()
+            dag = False
+        elif all(transpose):
             oper = oper.transpose()
+            transpose = [False]
 
         if isinstance(oper, _data.Dia) and N == 1:
             dia_matrix = oper.as_scipy()
@@ -145,12 +156,17 @@ def _oper_to_ElementaryOperator(
             for i, offset in enumerate(offsets):
                 end = None if offset == 0 else -abs(offset)
                 data[:end, i] = dia_matrix.diagonal(offset)
-            out = MultidiagonalOperator(data, offsets)
+            cu_oper = MultidiagonalOperator(data, offsets)
 
         else:
-            out = DenseOperator(oper.to_array().reshape(shape + shape))
+            cu_oper = DenseOperator(oper.to_array().reshape(shape + shape))
 
-    return out
+    if any(transpose):
+        cu_oper = _transpose_cu_operator(cu_oper, transpose)
+    if dag:
+        cu_oper = cu_oper.dag()
+
+    return cu_oper
 
 
 ###############################################################################
@@ -538,11 +554,14 @@ def to_OperatorTerm(self, dual=False, copy=True, hilbert_dims=None):
             for term in self.terms:
                 cuterm = tensor_product(dtype="complex128")
                 for pterm in term.prod_terms:
+                    transpose = pterm.transform in [Transform.TRANSPOSE, Transform.CONJ]
+                    dag = pterm.transform in [Transform.ADJOINT, Transform.CONJ]
                     oper = _oper_to_ElementaryOperator(
                         pterm.operator,
                         pterm.hilbert,
                         self.hilbert_space_dims,
-                        pterm.transform,
+                        [transpose] * len(pterm.hilbert),
+                        dag,
                         copy
                     )
                     # Inverted order confirmed by nvidia
@@ -552,40 +571,29 @@ def to_OperatorTerm(self, dual=False, copy=True, hilbert_dims=None):
         else:
             N_hilbert = len(self.hilbert_dims) // 2
             # TODO: make this tests weak compare?
-            assert self.hilbert_dims[:N_hilbert] == self.hilbert_dims[N_hilbert:]
+            if self.hilbert_dims[:N_hilbert] != self.hilbert_dims[N_hilbert:]:
+                raise ValueError(
+                    f"Hilbert space inconsistent with square superoperator: {self.hilbert_dims}"
+                )
             for term in self.terms:
                 cuterm = tensor_product(dtype="complex128")
                 for pterm in term.prod_terms:
-                    if all(i < N_hilbert for i in pterm.hilbert):
-                        oper = _oper_to_ElementaryOperator(
-                            pterm.operator,
-                            pterm.hilbert,
-                            self.hilbert_space_dims,
-                            trans_transform[pterm.transform],
-                            copy
-                        )
-                        # Inverted order confirmed by nvidia
-                        cuterm = cuterm * tensor_product(
-                            (oper, pterm.hilbert, (True,))
-                        )
-
-                    elif any(i < N_hilbert for i in pterm.hilbert):
-                        raise NotImplementedError(
-                            "Operators acting on both original and "
-                            "dual spaces are not supported."
-                        )
+                    modes = tuple(mode % N_hilbert for mode in pterm.hilbert)
+                    duals = tuple(mode < N_hilbert for mode in pterm.hilbert)
+                    transpose = tuple(duals)
+                    dag = pterm.transform in [Transform.ADJOINT, Transform.CONJ]
+                    if pterm.transform in [Transform.CONJ, Transform.TRANSPOSE]:
+                        transpose = tuple(not trans for trans in transpose)
 
-                    else:
-                        oper = _oper_to_ElementaryOperator(
-                            pterm.operator,
-                            pterm.hilbert,
-                            self.hilbert_space_dims,
-                            pterm.transform,
-                            copy
-                        )
-                        cuterm = cuterm * tensor_product(
-                            (oper, tuple(i - N_hilbert for i in pterm.hilbert))
-                        )
+                    oper = _oper_to_ElementaryOperator(
+                        pterm.operator,
+                        pterm.hilbert,
+                        self.hilbert_space_dims,
+                        transpose,
+                        dag,
+                        copy
+                    )
+                    cuterm = cuterm * tensor_product((oper, modes, duals,))
 
                 out = out + (cuterm * term.factor)
 
@@ -731,6 +739,22 @@ def isherm(operator, tol=-1):
     if tol < 0:
         tol = settings.core["atol"]
     return cp.allclose(oper, oper.T.conj(), atol=tol)
+
+
+@_data.identity_like.register(CuOperator)
+def identity_like(data, /):
+    """
+    Create an identity matrix of the same type and shape.
+    """
+    if not data.shape[0] == data.shape[1]:
+        raise ValueError(
+            "Can't create an identity matrix like a non square matrix."
+        )
+
+    new = CuOperator(hilbert_dims=data.hilbert_dims)
+    new.terms.append(Term([], 1.))
+    return new
+
 ###############################################################################
 ###############################################################################