add pow method to operators

doc/releases/changelog-dev.md

@@ -39,6 +39,10 @@
 
 <h3>Improvements</h3>
 
+* The developer-facing `pow` method has been added to `Operator` with concrete implementations
+  for many classes.
+  [(#2225)](https://github.com/PennyLaneAI/pennylane/pull/2225)
+
 * Test classes are created in qchem test modules to group the integrals and matrices unittests.
   [(#2545)](https://github.com/PennyLaneAI/pennylane/pull/2545)
 

pennylane/operation.py

@@ -255,6 +255,10 @@ class AdjointUndefinedError(OperatorPropertyUndefined):
     """Raised when an Operator's adjoint version is undefined."""
 
 
+class PowUndefinedError(OperatorPropertyUndefined):
+    """Raised when an Operator's power is undefined."""
+
+
 class GeneratorUndefinedError(OperatorPropertyUndefined):
     """Exception used to indicate that an operator
     does not have a generator"""
@@ -1025,6 +1029,25 @@ def generator(self):  # pylint: disable=no-self-use
         """
         raise GeneratorUndefinedError(f"Operation {self.name} does not have a generator")
 
+    def pow(self, z):
+        """A list of new operators equal to this one raised to the given power.
+
+        Args:
+            z (float): exponent for the operator
+
+        Returns:
+            list[:class:`~.operation.Operator`]
+
+        """
+        # Child methods may call super().pow(z%period) where op**period = I
+        # For example, PauliX**2 = I, SX**4 = I
+        # Hence we define 0 and 1 special cases here.
+        if z == 0:
+            return []
+        if z == 1:
+            return [self.__copy__()]
+        raise PowUndefinedError
+
     def queue(self, context=qml.QueuingContext):
         """Append the operator to the Operator queue."""
         context.append(self)

pennylane/ops/identity.py

@@ -148,3 +148,6 @@ def identity_op(*params):
 
     def adjoint(self):  # pylint:disable=arguments-differ
         return Identity(wires=self.wires)
+
+    def pow(self, n):
+        return [Identity(wires=self.wires)]

pennylane/ops/qubit/matrix_ops.py

@@ -153,6 +153,11 @@ def compute_decomposition(U, wires):
     def adjoint(self):
         return QubitUnitary(qml.math.T(qml.math.conj(self.matrix())), wires=self.wires)
 
+    def pow(self, z):
+        if isinstance(z, int):
+            return [QubitUnitary(qml.math.linalg.matrix_power(self.matrix(), z), wires=self.wires)]
+        return super().pow(z)
+
     def _controlled(self, wire):
         ControlledQubitUnitary(*self.parameters, control_wires=wire, wires=self.wires)
 
@@ -317,6 +322,17 @@ def compute_matrix(
     def control_wires(self):
         return self.hyperparameters["control_wires"]
 
+    def pow(self, z):
+        if isinstance(z, int):
+            return [
+                ControlledQubitUnitary(
+                    qml.math.linalg.matrix_power(self.data[0], z),
+                    control_wires=self.control_wires,
+                    wires=self.hyperparameters["u_wires"],
+                )
+            ]
+        return super().pow(z)
+
     def _controlled(self, wire):
         ctrl_wires = sorted(self.control_wires + wire)
         ControlledQubitUnitary(
@@ -439,6 +455,12 @@ def compute_decomposition(D, wires):
     def adjoint(self):
         return DiagonalQubitUnitary(qml.math.conj(self.parameters[0]), wires=self.wires)
 
+    def pow(self, z):
+        if isinstance(self.data[0], list):
+            return [DiagonalQubitUnitary([(x + 0.0j) ** z for x in self.data[0]], wires=self.wires)]
+        casted_data = qml.math.cast(self.data[0], np.complex128)
+        return [DiagonalQubitUnitary(casted_data**z, wires=self.wires)]
+
     def _controlled(self, control):
         DiagonalQubitUnitary(
             qml.math.concatenate([np.ones_like(self.parameters[0]), self.parameters[0]]),

pennylane/ops/qubit/non_parametric_ops.py

@@ -158,6 +158,9 @@ def single_qubit_rot_angles(self):
         # H = RZ(\pi) RY(\pi/2) RZ(0)
         return [np.pi, np.pi / 2, 0.0]
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
 
 class PauliX(Observable, Operation):
     r"""PauliX(wires)
@@ -288,6 +291,12 @@ def compute_decomposition(wires):
     def adjoint(self):
         return PauliX(wires=self.wires)
 
+    def pow(self, z):
+        z_mod2 = z % 2
+        if abs(z_mod2 - 0.5) < 1e-6:
+            return [SX(wires=self.wires)]
+        return super().pow(z_mod2)
+
     def _controlled(self, wire):
         CNOT(wires=Wires(wire) + self.wires)
 
@@ -427,6 +436,9 @@ def compute_decomposition(wires):
     def adjoint(self):
         return PauliY(wires=self.wires)
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     def _controlled(self, wire):
         CY(wires=Wires(wire) + self.wires)
 
@@ -554,6 +566,20 @@ def compute_decomposition(wires):
     def adjoint(self):
         return PauliZ(wires=self.wires)
 
+    def pow(self, z):
+        z_mod2 = z % 2
+        if z_mod2 == 0:
+            return []
+        if z_mod2 == 1:
+            return [self.__copy__()]
+
+        if abs(z_mod2 - 0.5) < 1e-6:
+            return [S(wires=self.wires)]
+        if abs(z_mod2 - 0.25) < 1e-6:
+            return [T(wires=self.wires)]
+
+        return [qml.PhaseShift(np.pi * z_mod2, wires=self.wires)]
+
     def _controlled(self, wire):
         CZ(wires=Wires(wire) + self.wires)
 
@@ -658,6 +684,18 @@ def adjoint(self):
         op.inverse = not self.inverse
         return op
 
+    def pow(self, z):
+        z_mod4 = z % 4
+        pow_map = {
+            0: lambda op: [],
+            0.5: lambda op: [T(wires=op.wires)],
+            1: lambda op: [op.__copy__()],
+            2: lambda op: [PauliZ(wires=op.wires)],
+        }
+        return pow_map.get(z_mod4, lambda op: [qml.PhaseShift(np.pi * z_mod4 / 2, wires=op.wires)])(
+            self
+        )
+
     def single_qubit_rot_angles(self):
         # S = RZ(\pi/2) RY(0) RZ(0)
         return [np.pi / 2, 0.0, 0.0]
@@ -754,6 +792,18 @@ def compute_decomposition(wires):
         """
         return [qml.PhaseShift(np.pi / 4, wires=wires)]
 
+    def pow(self, z):
+        z_mod8 = z % 8
+        pow_map = {
+            0: lambda op: [],
+            1: lambda op: [op.__copy__()],
+            2: lambda op: [S(wires=op.wires)],
+            4: lambda op: [PauliZ(wires=op.wires)],
+        }
+        return pow_map.get(z_mod8, lambda op: [qml.PhaseShift(np.pi * z_mod8 / 4, wires=op.wires)])(
+            self
+        )
+
     def adjoint(self):
         op = T(wires=self.wires)
         op.inverse = not self.inverse
@@ -865,6 +915,12 @@ def compute_decomposition(wires):
         ]
         return decomp_ops
 
+    def pow(self, z):
+        z_mod4 = z % 4
+        if z_mod4 == 2:
+            return [PauliX(wires=self.wires)]
+        return super().pow(z_mod4)
+
     def adjoint(self):
         op = SX(wires=self.wires)
         op.inverse = not self.inverse
@@ -931,6 +987,9 @@ def compute_matrix():  # pylint: disable=arguments-differ
     def adjoint(self):
         return CNOT(wires=self.wires)
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     def _controlled(self, wire):
         Toffoli(wires=Wires(wire) + self.wires)
 
@@ -1020,6 +1079,9 @@ def compute_eigvals():  # pylint: disable=arguments-differ
     def adjoint(self):
         return CZ(wires=self.wires)
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     @property
     def control_wires(self):
         return Wires(self.wires[0])
@@ -1112,6 +1174,9 @@ def compute_decomposition(wires):
     def adjoint(self):
         return CY(wires=self.wires)
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     @property
     def control_wires(self):
         return Wires(self.wires[0])
@@ -1190,6 +1255,9 @@ def compute_decomposition(wires):
         ]
         return decomp_ops
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     def adjoint(self):
         return SWAP(wires=self.wires)
 
@@ -1309,6 +1377,12 @@ def adjoint(self):
         op.inverse = not self.inverse
         return op
 
+    def pow(self, z):
+        z_mod2 = z % 2
+        if abs(z_mod2 - 0.5) < 1e-6:
+            return [SISWAP(wires=self.wires)]
+        return super().pow(z_mod2)
+
 
 class SISWAP(Operation):
     r"""SISWAP(wires)
@@ -1437,6 +1511,10 @@ def compute_decomposition(wires):
         ]
         return decomp_ops
 
+    def pow(self, z):
+        z_mod4 = z % 4
+        return [ISWAP(wires=self.wires)] if z_mod4 == 2 else super().pow(z_mod4)
+
     def adjoint(self):
         op = SISWAP(wires=self.wires)
         op.inverse = not self.inverse
@@ -1544,6 +1622,9 @@ def compute_decomposition(wires):
         ]
         return decomp_ops
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     def adjoint(self):
         return CSWAP(wires=self.wires)
 
@@ -1682,6 +1763,9 @@ def compute_decomposition(wires):
     def adjoint(self):
         return Toffoli(wires=self.wires)
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     @property
     def control_wires(self):
         return Wires(self.wires[:2])
@@ -1854,6 +1938,9 @@ def adjoint(self):
             control_values=self.hyperparameters["control_values"],
         )
 
+    def pow(self, z):
+        return super().pow(z % 2)
+
     @staticmethod
     def compute_decomposition(wires=None, work_wires=None, control_values=None, **kwargs):
         r"""Representation of the operator as a product of other operators (static method).
@@ -2056,6 +2143,9 @@ def _controlled(self, _):
     def adjoint(self, do_queue=True):
         return Barrier(wires=self.wires, do_queue=do_queue)
 
+    def pow(self, z):
+        return [self.__copy__()]
+
 
 class WireCut(Operation):
     r"""WireCut(wires)
@@ -2103,3 +2193,6 @@ def label(self, decimals=None, base_label=None, cache=None):
 
     def adjoint(self):
         return WireCut(wires=self.wires)
+
+    def pow(self, z):
+        return [self.__copy__()]

pennylane/ops/qubit/observables.py

@@ -441,3 +441,6 @@ def compute_diagonalizing_gates(
         []
         """
         return []
+
+    def pow(self, z):
+        return [self.__copy__()] if (isinstance(z, int) and z > 0) else super().pow(z)