Add TrotterProduct template (#4661)

**Context:**
First of 3 PRs adding the new TrotterProduct template to allow for
advanced Trotter methods in Pennylane

**Description of the Change:**
- Implement the template
- Add basic tests

---------

Co-authored-by: Tom Bromley <[email protected]>
Co-authored-by: soranjh <[email protected]>
Co-authored-by: soranjh <[email protected]>

doc/introduction/templates.rst

@@ -227,6 +227,10 @@ Other useful templates which do not belong to the previous categories can be fou
   :description: :doc:`ApproxTimeEvolution <../code/api/pennylane.ApproxTimeEvolution>`
   :figure: _static/templates/subroutines/approx_time_evolution.png
 
+.. gallery-item::
+  :description: :doc:`TrotterProduct <../code/api/pennylane.TrotterProduct>`
+  :figure: _static/templates/subroutines/trotter_product.png
+
 .. gallery-item::
   :description: :doc:`Permute <../code/api/pennylane.Permute>`
   :figure: _static/templates/subroutines/permute.png

doc/releases/changelog-dev.md

@@ -4,6 +4,41 @@
 
 <h3>New features since last release</h3>
 
+<h4>Exponentiate Hamiltonians with flexible Trotter products 🐖</h4>
+
+* Higher-order Trotter-Suzuki methods are now easily accessible through a new operation
+  called `TrotterProduct`.
+  [(#4661)](https://github.com/PennyLaneAI/pennylane/pull/4661)
+
+  Trotterization techniques are an affective route towards accurate and efficient
+  Hamiltonian simulation. The Suzuki-Trotter product formula allows for the ability
+  to express higher-order approximations to the matrix exponential of a Hamiltonian, 
+  and it is now available to use in PennyLane via the `TrotterProduct` operation. 
+  Simply specify the `order` of the approximation and the evolution `time`.
+
+  ```python
+  coeffs = [0.25, 0.75]
+  ops = [qml.PauliX(0), qml.PauliZ(0)]
+  H = qml.dot(coeffs, ops)
+
+  dev = qml.device("default.qubit", wires=2)
+
+  @qml.qnode(dev)
+  def circuit():
+      qml.Hadamard(0)
+      qml.TrotterProduct(H, time=2.4, order=2)
+      return qml.state()
+  ```
+
+  ```pycon
+  >>> circuit()
+  [-0.13259524+0.59790098j  0.        +0.j         -0.13259524-0.77932754j  0.        +0.j        ]
+  ```
+
+  The already-available `ApproxTimeEvolution` operation represents the special case of `order=1`.
+  It is recommended to switch over to use of `TrotterProduct` because `ApproxTimeEvolution` will be
+  deprecated and removed in upcoming releases.
+
 * Support drawing QJIT QNode from Catalyst.
   [(#4609)](https://github.com/PennyLaneAI/pennylane/pull/4609)
 

pennylane/ops/functions/equal.py

@@ -264,7 +264,9 @@ def _equal_adjoint(op1: Adjoint, op2: Adjoint, **kwargs):
 # pylint: disable=unused-argument
 def _equal_exp(op1: Exp, op2: Exp, **kwargs):
     """Determine whether two Exp objects are equal"""
-    if op1.coeff != op2.coeff:
+    rtol, atol = (kwargs["rtol"], kwargs["atol"])
+
+    if not qml.math.allclose(op1.coeff, op2.coeff, rtol=rtol, atol=atol):
         return False
     return qml.equal(op1.base, op2.base)
 
@@ -273,7 +275,9 @@ def _equal_exp(op1: Exp, op2: Exp, **kwargs):
 # pylint: disable=unused-argument
 def _equal_sprod(op1: SProd, op2: SProd, **kwargs):
     """Determine whether two SProd objects are equal"""
-    if op1.scalar != op2.scalar:
+    rtol, atol = (kwargs["rtol"], kwargs["atol"])
+
+    if not qml.math.allclose(op1.scalar, op2.scalar, rtol=rtol, atol=atol):
         return False
     return qml.equal(op1.base, op2.base)
 

pennylane/templates/subroutines/__init__.py

@@ -35,3 +35,4 @@
 from .basis_rotation import BasisRotation
 from .qsvt import QSVT, qsvt
 from .select import Select
+from .trotter import TrotterProduct

pennylane/templates/subroutines/trotter.py

@@ -0,0 +1,292 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Contains templates for Suzuki-Trotter approximation based subroutines.
+"""
+import pennylane as qml
+from pennylane.operation import Operation
+from pennylane.ops import Sum
+
+
+def _scalar(order):
+    """Compute the scalar used in the recursive expression.
+
+    Args:
+        order (int): order of Trotter product (assume order is an even integer > 2).
+
+    Returns:
+        float: scalar to be used in the recursive expression.
+    """
+    root = 1 / (order - 1)
+    return (4 - 4**root) ** -1
+
+
+@qml.QueuingManager.stop_recording()
+def _recursive_expression(x, order, ops):
+    """Generate a list of operations using the
+    recursive expression which defines the Trotter product.
+
+    Args:
+        x (complex): the evolution 'time'
+        order (int): the order of the Trotter expansion
+        ops (Iterable(~.Operators)): a list of terms in the Hamiltonian
+
+    Returns:
+        list: the approximation as product of exponentials of the Hamiltonian terms
+    """
+    if order == 1:
+        return [qml.exp(op, x * 1j) for op in ops]
+
+    if order == 2:
+        return [qml.exp(op, x * 0.5j) for op in ops + ops[::-1]]
+
+    scalar_1 = _scalar(order)
+    scalar_2 = 1 - 4 * scalar_1
+
+    ops_lst_1 = _recursive_expression(scalar_1 * x, order - 2, ops)
+    ops_lst_2 = _recursive_expression(scalar_2 * x, order - 2, ops)
+
+    return (2 * ops_lst_1) + ops_lst_2 + (2 * ops_lst_1)
+
+
+class TrotterProduct(Operation):
+    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
+    exponential of a given Hamiltonian.
+
+    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
+    Hamiltonian expressed as a linear combination of terms which in general do not commute. Consider
+    the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`, the product formula is constructed using
+    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
+    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:
+
+    .. math::
+
+        \begin{align}
+            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
+            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
+            &\vdots \\
+            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
+        \end{align}
+
+    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m`th order,
+    :math:`n`-step Suzuki-Trotter approximation is then defined as:
+
+    .. math:: e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.
+
+    For more details see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.
+
+    Args:
+        hamiltonian (Union[.Hamiltonian, .Sum]): The Hamiltonian written as a linear combination
+            of operators with known matrix exponentials.
+        time (float): The time of evolution, namely the parameter :math:`t` in :math:`e^{iHt}`
+        n (int): An integer representing the number of Trotter steps to perform
+        order (int): An integer (:math:`m`) representing the order of the approximation (must be 1 or even)
+        check_hermitian (bool): A flag to enable the validation check to ensure this is a valid unitary operator
+
+    Raises:
+        TypeError: The ``hamiltonian`` is not of type :class:`~.Hamiltonian`, or :class:`~.Sum`.
+        ValueError: The ``hamiltonian`` must have atleast two terms.
+        ValueError: One or more of the terms in ``hamiltonian`` are not Hermitian.
+        ValueError: The ``order`` is not one or a positive even integer.
+
+    **Example**
+
+    .. code-block:: python3
+
+        coeffs = [0.25, 0.75]
+        ops = [qml.PauliX(0), qml.PauliZ(0)]
+        H = qml.dot(coeffs, ops)
+
+        dev = qml.device("default.qubit", wires=2)
+        @qml.qnode(dev)
+        def my_circ():
+            # Prepare some state
+            qml.Hadamard(0)
+
+            # Evolve according to H
+            qml.TrotterProduct(H, time=2.4, order=2)
+
+            # Measure some quantity
+            return qml.state()
+
+    >>> my_circ()
+    [-0.13259524+0.59790098j  0.        +0.j         -0.13259524-0.77932754j  0.        +0.j        ]
+
+    .. details::
+        :title: Usage Details
+
+        One can recover the behaviour of :class:`~.ApproxTimeEvolution` by setting :code:`order=1`.
+        We can also compute the gradient with respect to the coefficients of the Hamiltonian and the
+        evolution time:
+
+        .. code-block:: python3
+
+            @qml.qnode(dev)
+            def my_circ(c1, c2, time):
+                # Prepare H:
+                H = qml.dot([c1, c2], [qml.PauliX(0), qml.PauliZ(0)])
+
+                # Prepare some state
+                qml.Hadamard(0)
+
+                # Evolve according to H
+                qml.TrotterProduct(H, time, order=2)
+
+                # Measure some quantity
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        >>> args = np.array([1.23, 4.5, 0.1])
+        >>> qml.grad(my_circ)(*tuple(args))
+        (tensor(0.00961064, requires_grad=True), tensor(-0.12338274, requires_grad=True), tensor(-5.43401259, requires_grad=True))
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self, hamiltonian, time, n=1, order=1, check_hermitian=True, id=None
+    ):
+        r"""Initialize the TrotterProduct class"""
+
+        if order <= 0 or order != 1 and order % 2 != 0:
+            raise ValueError(
+                f"The order of a TrotterProduct must be 1 or a positive even integer, got {order}."
+            )
+
+        if isinstance(hamiltonian, qml.Hamiltonian):
+            coeffs, ops = hamiltonian.terms()
+            if len(coeffs) < 2:
+                raise ValueError(
+                    "There should be atleast 2 terms in the Hamiltonian. Otherwise use `qml.exp`"
+                )
+
+            hamiltonian = qml.dot(coeffs, ops)
+
+        if not isinstance(hamiltonian, Sum):
+            raise TypeError(
+                f"The given operator must be a PennyLane ~.Hamiltonian or ~.Sum got {hamiltonian}"
+            )
+
+        if check_hermitian:
+            for op in hamiltonian.operands:
+                if not op.is_hermitian:
+                    raise ValueError(
+                        "One or more of the terms in the Hamiltonian may not be Hermitian"
+                    )
+
+        self._hyperparameters = {
+            "n": n,
+            "order": order,
+            "base": hamiltonian,
+            "check_hermitian": check_hermitian,
+        }
+        super().__init__(time, wires=hamiltonian.wires, id=id)
+
+    def _flatten(self):
+        """Serialize the operation into trainable and non-trainable components.
+
+        Returns:
+            data, metadata: The trainable and non-trainable components.
+
+        See ``Operator._unflatten``.
+
+        The data component can be recursive and include other operations. For example, the trainable component of ``Adjoint(RX(1, wires=0))``
+        will be the operator ``RX(1, wires=0)``.
+
+        The metadata **must** be hashable.  If the hyperparameters contain a non-hashable component, then this
+        method and ``Operator._unflatten`` should be overridden to provide a hashable version of the hyperparameters.
+
+        **Example:**
+
+        >>> op = qml.Rot(1.2, 2.3, 3.4, wires=0)
+        >>> qml.Rot._unflatten(*op._flatten())
+        Rot(1.2, 2.3, 3.4, wires=[0])
+        >>> op = qml.PauliRot(1.2, "XY", wires=(0,1))
+        >>> qml.PauliRot._unflatten(*op._flatten())
+        PauliRot(1.2, XY, wires=[0, 1])
+
+        Operators that have trainable components that differ from their ``Operator.data`` must implement their own
+        ``_flatten`` methods.
+
+        >>> op = qml.ctrl(qml.U2(3.4, 4.5, wires="a"), ("b", "c") )
+        >>> op._flatten()
+        ((U2(3.4, 4.5, wires=['a']),),
+        (<Wires = ['b', 'c']>, (True, True), <Wires = []>))
+        """
+        hamiltonian = self.hyperparameters["base"]
+        time = self.parameters[0]
+
+        hashable_hyperparameters = tuple(
+            (key, value) for key, value in self.hyperparameters.items() if key != "base"
+        )
+        return (hamiltonian, time), hashable_hyperparameters
+
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        """Recreate an operation from its serialized format.
+
+        Args:
+            data: the trainable component of the operation
+            metadata: the non-trainable component of the operation.
+
+        The output of ``Operator._flatten`` and the class type must be sufficient to reconstruct the original
+        operation with ``Operator._unflatten``.
+
+        **Example:**
+
+        >>> op = qml.Rot(1.2, 2.3, 3.4, wires=0)
+        >>> op._flatten()
+        ((1.2, 2.3, 3.4), (<Wires = [0]>, ()))
+        >>> qml.Rot._unflatten(*op._flatten())
+        >>> op = qml.PauliRot(1.2, "XY", wires=(0,1))
+        >>> op._flatten()
+        ((1.2,), (<Wires = [0, 1]>, (('pauli_word', 'XY'),)))
+        >>> op = qml.ctrl(qml.U2(3.4, 4.5, wires="a"), ("b", "c") )
+        >>> type(op)._unflatten(*op._flatten())
+        Controlled(U2(3.4, 4.5, wires=['a']), control_wires=['b', 'c'])
+
+        """
+        hyperparameters_dict = dict(metadata)
+        return cls(*data, **hyperparameters_dict)
+
+    @staticmethod
+    def compute_decomposition(*args, **kwargs):
+        r"""Representation of the operator as a product of other operators (static method).
+
+        .. math:: O = O_1 O_2 \dots O_n.
+
+        .. note::
+
+            Operations making up the decomposition should be queued within the
+            ``compute_decomposition`` method.
+
+        .. seealso:: :meth:`~.Operator.decomposition`.
+
+        Args:
+            *params (list): trainable parameters of the operator, as stored in the ``parameters`` attribute
+            wires (Iterable[Any], Wires): wires that the operator acts on
+            **hyperparams (dict): non-trainable hyperparameters of the operator, as stored in the ``hyperparameters`` attribute
+
+        Returns:
+            list[Operator]: decomposition of the operator
+        """
+        time = args[0]
+        n = kwargs["n"]
+        order = kwargs["order"]
+        ops = kwargs["base"].operands
+
+        decomp = _recursive_expression(time / n, order, ops)[::-1] * n
+
+        if qml.QueuingManager.recording():
+            for op in decomp:  # apply operators in reverse order of expression
+                qml.apply(op)
+
+        return decomp