Add resources to TrotterProduct class

pennylane/templates/subroutines/trotter.py

@@ -15,9 +15,13 @@
 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
 
 
@@ -62,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.
 
@@ -236,6 +240,30 @@ 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)
+
+ unique_operation_decomp = (copy.deepcopy(op) for op in decomp)
+ depth = qml.tape.QuantumTape(ops=unique_operation_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

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
@@ -565,6 +655,74 @@ 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 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))
+
+ tracked_resources = qml.specs(circ)()["resources"]
+ expected_resources = Resources(
+ num_wires=2,
+ num_gates=30,
+ gate_types=defaultdict(int, {"Exp": 30}),
+ gate_sizes=defaultdict(int, {1: 30}),
+ depth=20,
+ )
+
+ assert expected_resources == tracked_resources
+
+
 class TestDecomposition:
  """Test the decomposition of the TrotterProduct class."""