Add resources() to qml.TrotterProduct class

doc/releases/changelog-dev.md

@@ -50,6 +50,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `qml.TrotterProduct` is now compatible with resource tracking by inheriting from `ResourcesOperation`. 
+   [(#5680)](https://github.com/PennyLaneAI/pennylane/pull/5680)
+
 * The wires for the `default.tensor` device are selected at runtime if they are not provided by user.
   [(#5744)](https://github.com/PennyLaneAI/pennylane/pull/5744)
 
@@ -380,5 +383,6 @@ Vincent Michaud-Rioux,
 Lee James O'Riordan,
 Mudit Pandey,
 Kenya Sakka,
+Jay Soni,
 Haochen Paul Wang,
 David Wierichs.

pennylane/templates/subroutines/trotter.py

@@ -15,10 +15,12 @@
 Contains templates for Suzuki-Trotter approximation based subroutines.
 """
 import copy
+from collections import defaultdict
 
 import pennylane as qml
 from pennylane.ops import Sum
 from pennylane.ops.op_math import SProd
+from pennylane.resource import Resources, ResourcesOperation
 from pennylane.resource.error import ErrorOperation, SpectralNormError
 from pennylane.wires import Wires
 
@@ -64,7 +66,7 @@ def _recursive_expression(x, order, ops):
     return (2 * ops_lst_1) + ops_lst_2 + (2 * ops_lst_1)
 
 
-class TrotterProduct(ErrorOperation):
+class TrotterProduct(ErrorOperation, ResourcesOperation):
     r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
     exponential of a given Hamiltonian.
 
@@ -251,6 +253,29 @@ def queue(self, context=qml.QueuingManager):
         context.append(self)
         return self
 
+    def resources(self) -> Resources:
+        """The resource requirements for a given instance of the Suzuki-Trotter product.
+
+        Returns:
+            Resources: The resources for an instance of ``TrotterProduct``.
+        """
+        with qml.QueuingManager.stop_recording():
+            decomp = self.compute_decomposition(*self.parameters, **self.hyperparameters)
+
+        num_wires = len(self.wires)
+        num_gates = len(decomp)
+
+        depth = qml.tape.QuantumTape(ops=decomp).graph.get_depth()
+
+        gate_types = defaultdict(int)
+        gate_sizes = defaultdict(int)
+
+        for op in decomp:
+            gate_types[op.name] += 1
+            gate_sizes[len(op.wires)] += 1
+
+        return Resources(num_wires, num_gates, gate_types, gate_sizes, depth)
+
     def error(
         self, method: str = "commutator-bound", fast: bool = True
     ):  # pylint: disable=arguments-differ
@@ -413,9 +438,10 @@ def compute_decomposition(*args, **kwargs):
         ops = kwargs["base"].operands
 
         decomp = _recursive_expression(time / n, order, ops)[::-1] * n
+        unique_decomp = [copy.copy(op) for op in decomp]
 
         if qml.QueuingManager.recording():
-            for op in decomp:  # apply operators in reverse order of expression
+            for op in unique_decomp:  # apply operators in reverse order of expression
                 qml.apply(op)
 
-        return decomp
+        return unique_decomp

tests/templates/test_subroutines/test_trotter.py

@@ -16,13 +16,15 @@
 """
 # pylint: disable=private-access, protected-access
 import copy
+from collections import defaultdict
 from functools import reduce
 
 import pytest
 
 import pennylane as qml
 from pennylane import numpy as qnp
 from pennylane.math import allclose, get_interface
+from pennylane.resource import Resources
 from pennylane.resource.error import SpectralNormError
 from pennylane.templates.subroutines.trotter import _recursive_expression, _scalar
 
@@ -211,6 +213,94 @@
     }
 )
 
+test_resources_data = {
+    # (hamiltonian_index, order): Resources computed by hand
+    (0, 1): Resources(
+        num_wires=2,
+        num_gates=3,
+        gate_types=defaultdict(int, {"Exp": 3}),
+        gate_sizes=defaultdict(int, {1: 3}),
+        depth=2,
+    ),
+    (0, 2): Resources(
+        num_wires=2,
+        num_gates=6,
+        gate_types=defaultdict(int, {"Exp": 6}),
+        gate_sizes=defaultdict(int, {1: 6}),
+        depth=4,
+    ),
+    (0, 4): Resources(
+        num_wires=2,
+        num_gates=30,
+        gate_types=defaultdict(int, {"Exp": 30}),
+        gate_sizes=defaultdict(int, {1: 30}),
+        depth=20,
+    ),
+    (1, 1): Resources(
+        num_wires=2,
+        num_gates=2,
+        gate_types=defaultdict(int, {"Exp": 2}),
+        gate_sizes=defaultdict(int, {1: 1, 2: 1}),
+        depth=2,
+    ),
+    (1, 2): Resources(
+        num_wires=2,
+        num_gates=4,
+        gate_types=defaultdict(int, {"Exp": 4}),
+        gate_sizes=defaultdict(int, {1: 2, 2: 2}),
+        depth=4,
+    ),
+    (1, 4): Resources(
+        num_wires=2,
+        num_gates=20,
+        gate_types=defaultdict(int, {"Exp": 20}),
+        gate_sizes=defaultdict(int, {1: 10, 2: 10}),
+        depth=20,
+    ),
+    (2, 1): Resources(
+        num_wires=2,
+        num_gates=3,
+        gate_types=defaultdict(int, {"Exp": 3}),
+        gate_sizes=defaultdict(int, {1: 2, 2: 1}),
+        depth=3,
+    ),
+    (2, 2): Resources(
+        num_wires=2,
+        num_gates=6,
+        gate_types=defaultdict(int, {"Exp": 6}),
+        gate_sizes=defaultdict(int, {1: 4, 2: 2}),
+        depth=6,
+    ),
+    (2, 4): Resources(
+        num_wires=2,
+        num_gates=30,
+        gate_types=defaultdict(int, {"Exp": 30}),
+        gate_sizes=defaultdict(int, {1: 20, 2: 10}),
+        depth=30,
+    ),
+    (3, 1): Resources(
+        num_wires=2,
+        num_gates=3,
+        gate_types=defaultdict(int, {"Exp": 3}),
+        gate_sizes=defaultdict(int, {1: 3}),
+        depth=2,
+    ),
+    (3, 2): Resources(
+        num_wires=2,
+        num_gates=6,
+        gate_types=defaultdict(int, {"Exp": 6}),
+        gate_sizes=defaultdict(int, {1: 6}),
+        depth=4,
+    ),
+    (3, 4): Resources(
+        num_wires=2,
+        num_gates=30,
+        gate_types=defaultdict(int, {"Exp": 30}),
+        gate_sizes=defaultdict(int, {1: 30}),
+        depth=20,
+    ),
+}
+
 
 def _generate_simple_decomp(coeffs, ops, time, order, n):
     """Given coeffs, ops and a time argument in a given framework, generate the
@@ -554,6 +644,113 @@ def test_tensorflow_interface(self):
             _ = op.error()
 
 
+class TestResources:
+    """Test the resources method of the TrotterProduct class"""
+
+    def test_resources_no_queuing(self):
+        """Test that no operations are queued when computing resources."""
+        time = 0.5
+        hamiltonian = qml.sum(qml.PauliX(0), qml.PauliZ(0))
+        op = qml.TrotterProduct(hamiltonian, time, n=5, order=2)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            _ = op.resources()
+
+        assert len(q.queue) == 0
+
+    @pytest.mark.parametrize("order", (1, 2, 4))
+    @pytest.mark.parametrize("hamiltonian_index, hamiltonian", enumerate(test_hamiltonians))
+    def test_resources(self, hamiltonian, hamiltonian_index, order):
+        """Test that the resources are tracked accurately."""
+        op = qml.TrotterProduct(hamiltonian, 4.2, order=order)
+
+        tracked_resources = op.resources()
+        expected_resources = test_resources_data[(hamiltonian_index, order)]
+
+        assert expected_resources == tracked_resources
+
+    @pytest.mark.parametrize("n", (1, 5, 10))
+    def test_resources_with_trotter_steps(self, n):
+        """Test that the resources are tracked accurately with number of steps."""
+        order = 2
+        hamiltonian_index = 0
+
+        op = qml.TrotterProduct(test_hamiltonians[hamiltonian_index], 0.5, order=order, n=n)
+        tracked_resources = op.resources()
+
+        expected_resources = Resources(
+            num_wires=2,
+            num_gates=6 * n,
+            gate_types=defaultdict(int, {"Exp": 6 * n}),
+            gate_sizes=defaultdict(int, {1: 6 * n}),
+            depth=4 * n,
+        )
+
+        assert expected_resources == tracked_resources
+
+    def test_resources_integration(self):
+        """Test that the resources integrate well with qml.tracker and qml.specs
+        resource tracking."""
+        time = 0.5
+        hamiltonian = qml.sum(qml.X(0), qml.Y(0), qml.Z(1))
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circ():
+            qml.TrotterProduct(hamiltonian, time, n=5, order=2)
+            return qml.expval(qml.Z(0))
+
+        expected_resources = Resources(
+            num_wires=2,
+            num_gates=30,
+            gate_types=defaultdict(int, {"Exp": 30}),
+            gate_sizes=defaultdict(int, {1: 30}),
+            depth=20,
+        )
+
+        with qml.Tracker(dev) as tracker:
+            circ()
+
+        spec_resources = qml.specs(circ)()["resources"]
+        tracker_resources = tracker.history["resources"][0]
+
+        assert expected_resources == spec_resources
+        assert expected_resources == tracker_resources
+
+    def test_resources_and_error(self):
+        """Test that we can compute the resources and error together"""
+        time = 0.1
+        coeffs = qml.math.array([1.0, 0.5])
+        hamiltonian = qml.dot(coeffs, [qml.X(0), qml.Y(0)])
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circ():
+            qml.TrotterProduct(hamiltonian, time, n=2, order=2)
+            return qml.expval(qml.Z(0))
+
+        specs = qml.specs(circ)()
+
+        computed_error = (specs["errors"])["SpectralNormError"]
+        computed_resources = specs["resources"]
+
+        # Expected resources and errors (computed by hand)
+        expected_resources = Resources(
+            num_wires=1,
+            num_gates=8,
+            gate_types=defaultdict(int, {"Exp": 8}),
+            gate_sizes=defaultdict(int, {1: 8}),
+            depth=8,
+        )
+        expected_error = 0.001
+
+        assert computed_resources == expected_resources
+        assert isinstance(computed_error, SpectralNormError)
+        assert qnp.isclose(computed_error.error, qml.math.array(expected_error))
+
+
 class TestDecomposition:
     """Test the decomposition of the TrotterProduct class."""