refactor: Simplify CustomMatrices

src/qibojit/backends/gpu.py

@@ -5,11 +5,7 @@
 from qibo.config import log, raise_error
 
 from qibojit.backends.cpu import NumbaBackend
-from qibojit.backends.matrices import (
-    CupyMatrices,
-    CuQuantumMatrices,
-    CustomCupyMatrices,
-)
+from qibojit.backends.matrices import CupyMatrices, CuQuantumMatrices
 
 
 class CupyBackend(NumbaBackend):  # pragma: no cover
@@ -47,8 +43,7 @@ def __init__(self):
         self.sparse = cp.sparse
         self.device = "/GPU:0"
         self.kernel_type = "double"
-        self.matrices = CupyMatrices(self.dtype)
-        self.custom_matrices = CustomCupyMatrices(self.dtype)
+        self.custom_matrices = CupyMatrices(self.dtype)
         try:
             if not cp.cuda.runtime.getDeviceCount():  # pragma: no cover
                 raise RuntimeError("Cannot use cupy backend if GPU is not available.")

src/qibojit/backends/matrices.py

@@ -4,67 +4,14 @@
 from qibo.backends.npmatrices import NumpyMatrices
 
 
-class CupyMatrices(NumpyMatrices):  # pragma: no cover
+class CustomMatrices(NumpyMatrices):
+    """
+    TODO: explain why this class is neeeded.
+    """
 
-    def __init__(self, dtype):
-        super().__init__(dtype)
-        import cupy as cp  # pylint: disable=E0401
-
-        self.cp = cp
-
-    def _cast(self, x, dtype):
-        is_cupy = [
-            isinstance(item, self.cp.ndarray) for sublist in x for item in sublist
-        ]
-        if any(is_cupy) and not all(is_cupy):
-            # for parametrized gates x is a mixed list of cp.arrays and floats
-            # thus a simple cp.array(x) fails
-            # first convert the cp.arrays to numpy, then build the numpy array and move it
-            # back to GPU
-            dim = len(x)
-            return self.cp.array(
-                np.array(
-                    [
-                        item.get() if isinstance(item, self.cp.ndarray) else item
-                        for sublist in x
-                        for item in sublist
-                    ]
-                ).reshape(dim, dim),
-                dtype=dtype,
-            )
-        return self.cp.array(x, dtype=dtype)
-
-    # Necessary to avoid https://github.com/qiboteam/qibo/issues/928
-    def Unitary(self, u):
-        if isinstance(u, self.cp.ndarray):
-            u = u.get()
-        return super().Unitary(u)
-
-
-class CuQuantumMatrices(NumpyMatrices):
     # These matrices are used by the custom operators and may
     # not correspond to the mathematical representation of each gate
 
-    @cached_property
-    def CNOT(self):
-        return self.X
-
-    @cached_property
-    def CY(self):
-        return self.Y
-
-    @cached_property
-    def CZ(self):
-        return self.Z
-
-    @cached_property
-    def CSX(self):
-        return self.SX
-
-    @cached_property
-    def CSXDG(self):
-        return self.SXDG
-
     def CRX(self, theta):
         return self.RX(theta)
 
@@ -74,35 +21,30 @@ def CRY(self, theta):
     def CRZ(self, theta):
         return self.RZ(theta)
 
-    def CU1(self, theta):
-        return self.U1(theta)
-
     def CU2(self, phi, lam):
         return self.U2(phi, lam)
 
     def CU3(self, theta, phi, lam):
         return self.U3(theta, phi, lam)
 
+    def U1(self, theta):
+        dtype = getattr(np, self.dtype)
+        return dtype(np.exp(1j * theta))
+
+    def CU1(self, theta):
+        return self.U1(theta)
+
     @cached_property
-    def TOFFOLI(self):
-        return self.X
+    def CSX(self):
+        return self.SX
 
     @cached_property
-    def CCZ(self):
-        return self.Z
+    def CSXDG(self):
+        return self.SXDG
 
     def DEUTSCH(self, theta):
         return 1j * self.RX(2 * theta)
 
-
-class CustomMatrices(CuQuantumMatrices):
-    # These matrices are used by the custom operators and may
-    # not correspond to the mathematical representation of each gate
-
-    def U1(self, theta):
-        dtype = getattr(np, self.dtype)
-        return dtype(np.exp(1j * theta))
-
     def fSim(self, theta, phi):
         cost = np.cos(theta) + 0j
         isint = -1j * np.sin(theta)
@@ -114,7 +56,10 @@ def GeneralizedfSim(self, u, phi):
         return np.array([u[0, 0], u[0, 1], u[1, 0], u[1, 1], phase], dtype=self.dtype)
 
 
-class CustomCupyMatrices(CustomMatrices):  # pragma: no cover
+class CupyMatrices(CustomMatrices):  # pragma: no cover
+    """
+    TODO: explain why this class is needed
+    """
 
     def __init__(self, dtype):
         super().__init__(dtype)
@@ -127,6 +72,10 @@ def _cast(self, x, dtype):
             isinstance(item, self.cp.ndarray) for sublist in x for item in sublist
         ]
         if any(is_cupy) and not all(is_cupy):
+            # for parametrized gates x is a mixed list of cp.arrays and floats
+            # thus a simple cp.array(x) fails
+            # first convert the cp.arrays to numpy, then build the numpy array and move it
+            # back to GPU
             dim = len(x)
             return self.cp.array(
                 np.array(
@@ -140,11 +89,39 @@ def _cast(self, x, dtype):
             )
         return self.cp.array(x, dtype=dtype)
 
+    # Necessary to avoid https://github.com/qiboteam/qibo/issues/928
+    def Unitary(self, u):
+        if isinstance(u, self.cp.ndarray):
+            u = u.get()
+        return super().Unitary(u)
+
+
+class CuQuantumMatrices(CupyMatrices):
+    """
+    TODO: explain why this class is needed
+    """
+
+    @cached_property
+    def CNOT(self):
+        return self.X
+
+    @cached_property
+    def TOFFOLI(self):
+        return self.X
+
+    @cached_property
+    def CY(self):
+        return self.Y
+
+    @cached_property
+    def CZ(self):
+        return self.Z
+
     def U1(self, theta):
-        return self.cp.array(super().U1(theta), dtype=self.dtype)
+        return self._cast(NumpyMatrices.U1(self, theta), self.dtype)
 
     def fSim(self, theta, phi):
-        return self.cp.array(super().fSim(theta, phi), dtype=self.dtype)
+        return self._cast(NumpyMatrices.fSim(self, theta, phi), self.dtype)
 
     def GeneralizedfSim(self, u, phi):
-        return self.cp.array(super().GeneralizedfSim(u, phi), dtype=self.dtype)
+        return self._cast(NumpyMatrices.GeneralizedfSim(self, u, phi), self.dtype)