Repository: qiboteam/qibo

Issue:
U1.matrix returns a single value
When trying to get the matrix form of U1, qibo returns a single value of the exponential instead of the full 2x2 matrix.

Discussion:

Patch:
diff --git a/src/qibo/abstractions/abstract_gates.py b/src/qibo/abstractions/abstract_gates.py
index 8ec07d9ab..09ecc28ba 100644
--- a/src/qibo/abstractions/abstract_gates.py
+++ b/src/qibo/abstractions/abstract_gates.py
@@ -404,8 +404,9 @@ class ParametrizedGate(Gate):
         # ``circuit.set_parameters`` method works properly.
         # pylint: disable=E1101
         if isinstance(self, BaseBackendGate):
-            self._unitary = None
             self._matrix = None
+            self._native_op_matrix = None
+            self._custom_op_matrix = None
             for devgate in self.device_gates:
                 devgate.parameters = x
 
@@ -440,7 +441,6 @@ class BaseBackendGate(Gate, ABC):
     def __init__(self):
         Gate.__init__(self)
         self._matrix = None
-        self._unitary = None
         self._cache = None
         # Cast gate matrices to the proper device
         self.device = get_device()
@@ -455,11 +455,16 @@ class BaseBackendGate(Gate, ABC):
         if len(self.qubits) > 2:
             raise_error(NotImplementedError, "Cannot calculate unitary matrix for "
                                              "gates that target more than two qubits.")
-        if self._unitary is None:
-            self._unitary = self.construct_unitary()
-        if self.is_controlled_by and tuple(self._unitary.shape) == (2, 2):
-            self._unitary = self.control_unitary(self._unitary)
-        return self._unitary
+        if self._matrix is None:
+            self._matrix = self.construct_unitary()
+        if self.is_controlled_by and tuple(self._matrix.shape) == (2, 2):
+            self._matrix = self.control_unitary(self._matrix)
+        return self._matrix
+
+    @property
+    def matrix(self):
+        """Unitary matrix representing the gate in the computational basis."""
+        return self.unitary
 
     def __matmul__(self, other: "Gate") -> "Gate":
         """Gate multiplication."""
@@ -530,7 +535,7 @@ class BaseBackendGate(Gate, ABC):
             U\\rho
 
         This is useful for :class:`qibo.abstractions.hamiltonians.SymbolicHamiltonian`
-        multiplication to density matrices. 
+        multiplication to density matrices.
         """
         raise_error(NotImplementedError)
 
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 24f707ff5..0eae47430 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -1,3 +1,4 @@
+import math
 from qibo.abstractions.states import AbstractState
 from qibo.backends import abstract, einsum_utils
 from qibo.config import raise_error, log
@@ -274,8 +275,17 @@ class NumpyBackend(abstract.AbstractBackend):
             cache.calculation_cache = self.create_einsum_cache(gate.qubits, gate.nqubits)
         return cache
 
+    def reshape_matrix(self, matrix):
+        """Reshapes the gate matrix to the tensor shape.
+
+        Helper method for gate calls.
+        """
+        rank = int(math.log2(int(matrix.shape[0])))
+        return self.reshape(matrix, 2 * rank * (2,))
+
     def state_vector_call(self, gate, state):
         state = self.reshape(state, gate.cache.tensor_shape)
+        matrix = self.reshape_matrix(gate.native_op_matrix)
         if gate.is_controlled_by:
             ncontrol = len(gate.control_qubits)
             nactive = gate.nqubits - ncontrol
@@ -283,8 +293,7 @@ class NumpyBackend(abstract.AbstractBackend):
             # Apply `einsum` only to the part of the state where all controls
             # are active. This should be `state[-1]`
             state = self.reshape(state, (2 ** ncontrol,) + nactive * (2,))
-            updates = self.einsum_call(gate.cache.calculation_cache.vector, state[-1],
-                                     gate.matrix)
+            updates = self.einsum_call(gate.cache.calculation_cache.vector, state[-1], matrix)
             # Concatenate the updated part of the state `updates` with the
             # part of of the state that remained unaffected `state[:-1]`.
             state = self.concatenate([state[:-1], updates[self.newaxis]], axis=0)
@@ -293,7 +302,7 @@ class NumpyBackend(abstract.AbstractBackend):
             state = self.transpose(state, gate.cache.control_cache.reverse(False))
         else:
             einsum_str = gate.cache.calculation_cache.vector
-            state = self.einsum_call(einsum_str, state, gate.matrix)
+            state = self.einsum_call(einsum_str, state, matrix)
         return self.reshape(state, gate.cache.flat_shape)
 
     def state_vector_matrix_call(self, gate, state):
@@ -301,6 +310,8 @@ class NumpyBackend(abstract.AbstractBackend):
 
     def density_matrix_call(self, gate, state):
         state = self.reshape(state, gate.cache.tensor_shape)
+        matrix = self.reshape_matrix(gate.native_op_matrix)
+        matrixc = self.conj(matrix)
         if gate.is_controlled_by:
             ncontrol = len(gate.control_qubits)
             nactive = gate.nqubits - ncontrol
@@ -309,17 +320,16 @@ class NumpyBackend(abstract.AbstractBackend):
             state = self.reshape(state, 2 * (n,) + 2 * nactive * (2,))
             state01 = self.gather(state, indices=range(n - 1), axis=0)
             state01 = self.squeeze(self.gather(state01, indices=[n - 1], axis=1), axis=1)
-            state01 = self.einsum_call(gate.cache.calculation_cache.right0,
-                                     state01, self.conj(gate.matrix))
+            state01 = self.einsum_call(gate.cache.calculation_cache.right0, state01, matrixc)
             state10 = self.gather(state, indices=range(n - 1), axis=1)
             state10 = self.squeeze(self.gather(state10, indices=[n - 1], axis=0), axis=0)
             state10 = self.einsum_call(gate.cache.calculation_cache.left0,
-                                       state10, gate.matrix)
+                                       state10, matrix)
 
             state11 = self.squeeze(self.gather(state, indices=[n - 1], axis=0), axis=0)
             state11 = self.squeeze(self.gather(state11, indices=[n - 1], axis=0), axis=0)
-            state11 = self.einsum_call(gate.cache.calculation_cache.right, state11, self.conj(gate.matrix))
-            state11 = self.einsum_call(gate.cache.calculation_cache.left, state11, gate.matrix)
+            state11 = self.einsum_call(gate.cache.calculation_cache.right, state11, matrixc)
+            state11 = self.einsum_call(gate.cache.calculation_cache.left, state11, matrix)
 
             state00 = self.gather(state, indices=range(n - 1), axis=0)
             state00 = self.gather(state00, indices=range(n - 1), axis=1)
@@ -329,9 +339,8 @@ class NumpyBackend(abstract.AbstractBackend):
             state = self.reshape(state, 2 * gate.nqubits * (2,))
             state = self.transpose(state, gate.cache.control_cache.reverse(True))
         else:
-            state = self.einsum_call(gate.cache.calculation_cache.right, state,
-                                   self.conj(gate.matrix))
-            state = self.einsum_call(gate.cache.calculation_cache.left, state, gate.matrix)
+            state = self.einsum_call(gate.cache.calculation_cache.right, state, matrixc)
+            state = self.einsum_call(gate.cache.calculation_cache.left, state, matrix)
         return self.reshape(state, gate.cache.flat_shape)
 
     def density_matrix_matrix_call(self, gate, state):
@@ -342,8 +351,9 @@ class NumpyBackend(abstract.AbstractBackend):
             raise_error(NotImplementedError, "Gate density matrix half call is "
                                              "not implemented for ``controlled_by``"
                                              "gates.")
+        matrix = self.reshape_matrix(gate.native_op_matrix)
         state = self.reshape(state, gate.cache.tensor_shape)
-        state = self.einsum_call(gate.cache.calculation_cache.left, state, gate.matrix)
+        state = self.einsum_call(gate.cache.calculation_cache.left, state, matrix)
         return self.reshape(state, gate.cache.flat_shape)
 
     def density_matrix_half_matrix_call(self, gate, state):
@@ -560,7 +570,7 @@ class JITCustomBackend(NumpyBackend): # pragma: no cover
                             gate.cache.qubits_tensor)
 
     def state_vector_matrix_call(self, gate, state):
-        return gate.gate_op(state, gate.matrix, gate.nqubits, *gate.target_qubits,
+        return gate.gate_op(state, gate.custom_op_matrix, gate.nqubits, *gate.target_qubits,
                             gate.cache.qubits_tensor)
 
     def density_matrix_call(self, gate, state):
@@ -575,9 +585,9 @@ class JITCustomBackend(NumpyBackend): # pragma: no cover
     def density_matrix_matrix_call(self, gate, state):
         qubits = tuple(x + gate.nqubits for x in gate.cache.qubits_tensor)
         shape = state.shape
-        state = gate.gate_op(state.flatten(), gate.matrix, 2 * gate.nqubits,
+        state = gate.gate_op(state.flatten(), gate.custom_op_matrix, 2 * gate.nqubits,
                              *gate.target_qubits, qubits)
-        adjmatrix = self.conj(gate.matrix)
+        adjmatrix = self.conj(gate.custom_op_matrix)
         state = gate.gate_op(state, adjmatrix, 2 * gate.nqubits, *gate.cache.target_qubits_dm,
                              gate.cache.qubits_tensor)
         return self.reshape(state, shape)
@@ -592,7 +602,7 @@ class JITCustomBackend(NumpyBackend): # pragma: no cover
     def density_matrix_half_matrix_call(self, gate, state):
         qubits = tuple(x + gate.nqubits for x in gate.cache.qubits_tensor)
         shape = state.shape
-        state = gate.gate_op(state.flatten(), gate.matrix, 2 * gate.nqubits,
+        state = gate.gate_op(state.flatten(), gate.custom_op_matrix, 2 * gate.nqubits,
                              *gate.target_qubits, qubits)
         return self.reshape(state, shape)
 
diff --git a/src/qibo/backends/tensorflow.py b/src/qibo/backends/tensorflow.py
index c7af417fb..eb50c487e 100644
--- a/src/qibo/backends/tensorflow.py
+++ b/src/qibo/backends/tensorflow.py
@@ -272,7 +272,7 @@ class TensorflowCustomBackend(TensorflowBackend):
                             *gate.target_qubits, self.nthreads)
 
     def state_vector_matrix_call(self, gate, state):
-        return gate.gate_op(state, gate.matrix, gate.cache.qubits_tensor, # pylint: disable=E1121
+        return gate.gate_op(state, gate.custom_op_matrix, gate.cache.qubits_tensor, # pylint: disable=E1121
                             gate.nqubits, *gate.target_qubits,
                             self.nthreads)
 
@@ -285,10 +285,10 @@ class TensorflowCustomBackend(TensorflowBackend):
         return state
 
     def density_matrix_matrix_call(self, gate, state):
-        state = gate.gate_op(state, gate.matrix, gate.cache.qubits_tensor + gate.nqubits, # pylint: disable=E1121
+        state = gate.gate_op(state, gate.custom_op_matrix, gate.cache.qubits_tensor + gate.nqubits, # pylint: disable=E1121
                              2 * gate.nqubits, *gate.target_qubits,
                              self.nthreads)
-        adjmatrix = self.conj(gate.matrix)
+        adjmatrix = self.conj(gate.custom_op_matrix)
         state = gate.gate_op(state, adjmatrix, gate.cache.qubits_tensor,
                              2 * gate.nqubits, *gate.cache.target_qubits_dm,
                              self.nthreads)
@@ -300,7 +300,7 @@ class TensorflowCustomBackend(TensorflowBackend):
                             self.nthreads)
 
     def density_matrix_half_matrix_call(self, gate, state):
-        return gate.gate_op(state, gate.matrix, gate.cache.qubits_tensor + gate.nqubits, # pylint: disable=E1121
+        return gate.gate_op(state, gate.custom_op_matrix, gate.cache.qubits_tensor + gate.nqubits, # pylint: disable=E1121
                             2 * gate.nqubits, *gate.target_qubits,
                             self.nthreads)
 
diff --git a/src/qibo/core/gates.py b/src/qibo/core/gates.py
index a68da7ec0..033a8c1d5 100644
--- a/src/qibo/core/gates.py
+++ b/src/qibo/core/gates.py
@@ -28,6 +28,8 @@ class BackendGate(BaseBackendGate):
         else:
             self.gate_op = None
         super().__init__()
+        self._native_op_matrix = None
+        self._custom_op_matrix = None
 
     @staticmethod
     def control_unitary(unitary):
@@ -49,17 +51,16 @@ class BackendGate(BaseBackendGate):
         return self._cache
 
     @property
-    def matrix(self):
-        if self._matrix is None:
-            self._matrix = self.calculate_matrix()
-        return self._matrix
+    def native_op_matrix(self):
+        if self._native_op_matrix is None:
+            self._native_op_matrix = self.construct_unitary()
+        return self._native_op_matrix
 
-    def calculate_matrix(self):
-        matrix = self.construct_unitary()
-        if K.op is None:
-            rank = int(math.log2(int(matrix.shape[0])))
-            matrix = K.reshape(matrix, 2 * rank * (2,))
-        return matrix
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
+            self._custom_op_matrix = self.construct_unitary()
+        return self._custom_op_matrix
 
     def set_nqubits(self, state):
         if self._nqubits is None:
@@ -133,9 +134,9 @@ class Y(BackendGate, abstract_gates.Y):
         self._target_qubits = self.cache.target_qubits_dm
         self._nqubits *= 2
         self.gate_op = K.op.apply_gate
-        self._matrix = K.conj(K.matrices.Y)
+        self._custom_op_matrix = K.conj(K.matrices.Y)
         state = K.state_vector_matrix_call(self, state)
-        self._matrix = K.matrices.Y
+        self._custom_op_matrix = K.matrices.Y
         self.gate_op = K.op.apply_y
         self._nqubits //= 2
         self._target_qubits = original_targets
@@ -363,11 +364,11 @@ class U1(MatrixGate, abstract_gates.U1):
         if self.gate_op:
             self.gate_op = K.op.apply_z_pow
 
-    def calculate_matrix(self):
-        if K.op is not None:
-            return K.cast(K.qnp.exp(1j * self.parameters))
-        else:
-            return MatrixGate.calculate_matrix(self)
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
+            self._custom_op_matrix = K.cast(K.qnp.exp(1j * self.parameters))
+        return self._custom_op_matrix
 
     def construct_unitary(self):
         if isinstance(self.parameters, K.native_types): # pragma: no cover
@@ -447,15 +448,15 @@ class _CUn_(MatrixGate):
         cbase = "C{}".format(self.base.__name__)
         getattr(abstract_gates, cbase).__init__(self, q0, q1, **params)
 
-    def calculate_matrix(self):
-        matrix = self.base.construct_unitary(self)
-        if K.op is None:
-            matrix = K.reshape(self.control_unitary(matrix), 4 * (2,))
-        return matrix
-
     def construct_unitary(self):
         return MatrixGate.control_unitary(self.base.construct_unitary(self))
 
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
+            self._custom_op_matrix = self.base.construct_unitary(self)
+        return self._custom_op_matrix
+
 
 class CRX(_CUn_, abstract_gates.CRX):
     base = RX
@@ -486,8 +487,11 @@ class CU1(_CUn_, abstract_gates.CU1):
         if self.gate_op:
             self.gate_op = K.op.apply_z_pow
 
-    def calculate_matrix(self):
-        return U1.calculate_matrix(self)
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
+            self._custom_op_matrix = K.cast(K.qnp.exp(1j * self.parameters))
+        return self._custom_op_matrix
 
 
 class CU2(_CUn_, abstract_gates.CU2):
@@ -525,13 +529,14 @@ class fSim(MatrixGate, abstract_gates.fSim):
         if self.gate_op:
             self.gate_op = K.op.apply_fsim
 
-    def calculate_matrix(self):
-        if K.op is not None:
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
             theta, phi = self.parameters
             cos, isin = K.qnp.cos(theta) + 0j, -1j * K.qnp.sin(theta)
             phase = K.qnp.exp(-1j * phi)
-            return K.cast([cos, isin, isin, cos, phase])
-        return MatrixGate.calculate_matrix(self)
+            self._custom_op_matrix = K.cast([cos, isin, isin, cos, phase])
+        return self._custom_op_matrix
 
     def construct_unitary(self):
         theta, phi = self.parameters
@@ -555,14 +560,15 @@ class GeneralizedfSim(MatrixGate, abstract_gates.GeneralizedfSim):
         if self.gate_op:
             self.gate_op = K.op.apply_fsim
 
-    def calculate_matrix(self):
-        if K.op is not None:
+    @property
+    def custom_op_matrix(self):
+        if self._custom_op_matrix is None:
             unitary, phi = self.parameters
             matrix = K.qnp.zeros(5)
             matrix[:4] = K.qnp.reshape(unitary, (4,))
             matrix[4] = K.qnp.exp(-1j * phi)
-            return K.cast(matrix)
-        return MatrixGate.calculate_matrix(self)
+            self._custom_op_matrix = K.cast(matrix)
+        return self._custom_op_matrix
 
     def construct_unitary(self):
         unitary, phi = self.parameters
@@ -598,9 +604,9 @@ class TOFFOLI(BackendGate, abstract_gates.TOFFOLI):
 
     @property
     def unitary(self):
-        if self._unitary is None:
-            self._unitary = self.construct_unitary()
-        return self._unitary
+        if self._matrix is None:
+            self._matrix = self.construct_unitary()
+        return self._matrix
 
 
 class Unitary(MatrixGate, abstract_gates.Unitary):


Test patch:
diff --git a/src/qibo/tests/test_core_gates_features.py b/src/qibo/tests/test_core_gates_features.py
index 470f65f0f..579c57a64 100644
--- a/src/qibo/tests/test_core_gates_features.py
+++ b/src/qibo/tests/test_core_gates_features.py
@@ -52,6 +52,7 @@ def test_construct_unitary_rotations(backend, gate, target_matrix):
     else:
         gate = getattr(gates, gate)(0, theta)
     K.assert_allclose(gate.unitary, target_matrix(theta))
+    K.assert_allclose(gate.matrix, target_matrix(theta))
 
 
 def test_construct_unitary_controlled(backend):