Store grouping in Hamiltonian (#1515)

* delete arithmetic and rewrite rest with qml.math * fix most errors not pertaining to grouping or ExpvalCost * add grouping to Hamiltonian * change order of arithmetic ops * change order of arithmetic ops 2 * fix bug in hamiltonian_expand * make all tests pass * fix diffability * backup * undo comments * write some crucial tests * Hamiltonians are differentiable * some tweaks to make more tests pass * black * all tests pass * add param shift test * black * port some changes from a child branch * all changes and tests * black * make codefactor happy * make codefactor even more happy * better docstrings * update changelog * black * increase coverage * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Josh Izaac <josh146@gmail.com> * implement comments * add num_params attribute * make pylint happy * implement second round of comments * black * port changes from prototype * move test file * black * add missed line to change * fixes * add test for differentiable simplify * backup * backup * delete lines, start writing test * add test * fix test * rename one test * update changelog * main tests pass * backup * remove test * backup * missed one conflict * test diffable grouping * add another changelog entry * remove check * change how the return type is checked * black * black again * add comment on arithmetics * improve docstring * polish * fix some tests * polish more * Update tests/ops/test_hamiltonian.py * add kron to tf * Update .github/CHANGELOG.md Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update .github/CHANGELOG.md Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * add Hamiltonian test * black * add tests for list/tuple return of grouping * remove tuple conversion again * fix data issue and add test * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * polish * update torch version * bump torchvision too, to be compatible * fix codevec * correct torchvision version again * add some more tests * figure out diff bug * backup * fix tests and dimensions * add changelog * clean changelog * improve docstring * polish * fix test * fix codefactor issue * Update tests/transforms/test_hamiltonian_expand.py * Update pennylane/transforms/hamiltonian_expand.py * Update pennylane/transforms/hamiltonian_expand.py * backup * Update pennylane/vqe/vqe.py * improve docstrings * add test * add a queue test * black * Update .github/CHANGELOG.md Co-authored-by: Josh Izaac <josh146@gmail.com> * Update .github/CHANGELOG.md Co-authored-by: Josh Izaac <josh146@gmail.com> * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * port changes from wrong branch * improve grouping control * Update .github/CHANGELOG.md Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * Update pennylane/vqe/vqe.py Co-authored-by: Josh Izaac <josh146@gmail.com> * implement suggestions * remove get_grouping * black * finish * Update tests/transforms/test_hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/vqe/vqe.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/vqe/vqe.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/vqe/vqe.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/vqe/vqe.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * Update pennylane/transforms/hamiltonian_expand.py Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com> * implement suggestions * black * check docs are fine * add tests for measurement grouping trafo which is not used in hamiltonian_expand any more * change kwarg logic * fix test * adapt changelog * black Co-authored-by: Josh Izaac <josh146@gmail.com> Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>
PennyLaneAI · Aug 16, 2021 · 9ed7ca4 · 9ed7ca4
1 parent 0088704
commit 9ed7ca4
Show file tree

Hide file tree

Showing 7 changed files with 385 additions and 58 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,24 @@
 
 <h3>New features since last release</h3>
 
+* The Hamiltonian can now store grouping information, which can be accessed by a device to 
+  speed up computations of the expectation value of a Hamiltonian. 
+  [(#1515)](https://github.com/PennyLaneAI/pennylane/pull/1515)
+
+  ```python
+  obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
+  coeffs = np.array([1., 2., 3.])
+  H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+  ```
+
+  Initialization with a ``grouping_type`` other than ``None`` stores the indices 
+  required to make groups of commuting observables and their coefficients. 
+
+  ``` pycon
+  >>> H.grouping_indices
+  [[0, 1], [2]]
+  ```
+
 * Hamiltonians are now trainable with respect to their coefficients.
   [(#1483)](https://github.com/PennyLaneAI/pennylane/pull/1483)
 

diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
@@ -26,7 +26,8 @@ def hamiltonian_expand(tape, group=True):
     Args:
         tape (.QuantumTape): the tape used when calculating the expectation value
             of the Hamiltonian
-        group (bool): whether to compute groups of non-commuting Pauli observables, leading to fewer tapes
+        group (bool): Whether to compute disjoint groups of commuting Pauli observables, leading to fewer tapes.
+            If grouping information can be found in the Hamiltonian, it will be used even if group=False.
 
     Returns:
         tuple[list[.QuantumTape], function]: Returns a tuple containing a list of
@@ -67,32 +68,54 @@ def hamiltonian_expand(tape, group=True):
     >>> fn(res)
     -0.5
 
-    .. Warning::
+    Fewer tapes can be constructed by grouping commuting observables. This can be achieved
+    by the ``group`` keyword argument:
 
-         Note that defining Hamiltonians inside of QNodes using arithmetic can lead to errors.
-         See :class:`~pennylane.Hamiltonian` for more information.
+    .. code-block:: python3
+
+        H = qml.Hamiltonian([1., 2., 3.], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliX(wires=2)
+            qml.expval(H)
+
+    With grouping, the Hamiltonian gets split into two groups of observables (here ``[qml.PauliZ(0)]`` and
+    ``[qml.PauliX(1), qml.PauliX(0)]``):
+
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape)
+    >>> len(tapes)
+    2
+
+    Without grouping it gets split into three groups (``[qml.PauliZ(0)]``, ``[qml.PauliX(1)]`` and ``[qml.PauliX(0)]``):
+
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+    >>> len(tapes)
+    3
 
-    The ``group`` keyword argument toggles between the creation of one tape per Pauli observable, or
-    one tape per group of non-commuting Pauli observables computed by the :func:`.measurement_grouping`
-    transform:
+    Alternatively, if the Hamiltonian has already computed groups, they are used even if ``group=False``:
 
     .. code-block:: python3
 
-        H = qml.Hamiltonian([1., 2., 3.], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)])
+        obs = [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)]
+        coeffs = [1., 2., 3.]
+        H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+
+        # the initialisation already computes grouping information and stores it in the Hamiltonian
+        assert H.grouping_indices is not None
 
         with qml.tape.QuantumTape() as tape:
             qml.Hadamard(wires=0)
             qml.CNOT(wires=[0, 1])
             qml.PauliX(wires=2)
             qml.expval(H)
 
-        # split H into observable groups [qml.PauliZ(0)] and [qml.PauliX(1), qml.PauliX(0)]
-        tapes, fn = qml.transforms.hamiltonian_expand(tape)
-        print(len(tapes)) # 2
+    Grouping information has been used to reduce the number of tapes from 3 to 2:
 
-        # split H into observables [qml.PauliZ(0)], [qml.PauliX(1)] and [qml.PauliX(0)]
-        tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
-        print(len(tapes)) # 3
+    >>> tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+    >>> len(tapes)
+    2
     """
 
     hamiltonian = tape.measurements[0].obs
@@ -106,27 +129,48 @@ def hamiltonian_expand(tape, group=True):
             "Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`"
         )
 
-    if group:
-        hamiltonian.simplify()
-        return qml.transforms.measurement_grouping(tape, hamiltonian.ops, hamiltonian.coeffs)
-
-    # create tapes that measure the Pauli-words in the Hamiltonian
-    tapes = []
-    for ob in hamiltonian.ops:
-        # we need to create a new tape here, because
-        # updating metadata of a copied tape is error-prone
-        # when the observables were changed
-        with tape.__class__() as new_tape:
-            for op in tape.operations:
-                qml.apply(op)
-            qml.expval(ob)
-        tapes.append(new_tape)
-
-    # create processing function that performs linear recombination
-    def processing_fn(res):
-        dot_products = [
-            qml.math.dot(qml.math.squeeze(res[i]), hamiltonian.coeffs[i]) for i in range(len(res))
+    if group or hamiltonian.grouping_indices is not None:
+
+        if hamiltonian.grouping_indices is None:
+            hamiltonian.compute_grouping()
+
+        # use groups of observables if available or explicitly requested
+        coeffs = [
+            qml.math.squeeze(qml.math.take(hamiltonian.coeffs, indices, axis=0))
+            for indices in hamiltonian.grouping_indices
+        ]
+        obs_groupings = [
+            [hamiltonian.ops[i] for i in indices] for indices in hamiltonian.grouping_indices
         ]
+
+        tapes = []
+        for obs in obs_groupings:
+
+            with tape.__class__() as new_tape:
+                for op in tape.operations:
+                    op.queue()
+
+                for o in obs:
+                    qml.expval(o)
+
+            new_tape = new_tape.expand(stop_at=lambda obj: True)
+            tapes.append(new_tape)
+    else:
+        coeffs = hamiltonian.coeffs
+
+        tapes = []
+        for o in hamiltonian.ops:
+            with tape.__class__() as new_tape:
+                for op in tape.operations:
+                    op.queue()
+                qml.expval(o)
+
+            tapes.append(new_tape)
+
+    def processing_fn(res):
+        # note: res could have an extra dimension here if a shots_distribution
+        # is used for evaluation
+        dot_products = [qml.math.dot(qml.math.squeeze(r), c) for c, r in zip(coeffs, res)]
         return qml.math.sum(qml.math.stack(dot_products), axis=0)
 
     return tapes, processing_fn
diff --git a/pennylane/vqe/vqe.py b/pennylane/vqe/vqe.py
@@ -30,6 +30,35 @@
 OBS_MAP = {"PauliX": "X", "PauliY": "Y", "PauliZ": "Z", "Hadamard": "H", "Identity": "I"}
 
 
+def _compute_grouping_indices(observables, grouping_type="qwc", method="rlf"):
+
+    # todo: directly compute the
+    # indices, instead of extracting groups of observables first
+    observable_groups = qml.grouping.group_observables(
+        observables, coefficients=None, grouping_type=grouping_type, method=method
+    )
+
+    observables = copy(observables)
+
+    indices = []
+    available_indices = list(range(len(observables)))
+    for partition in observable_groups:
+        indices_this_group = []
+        for pauli_word in partition:
+            # find index of this pauli word in remaining original observables,
+            for observable in observables:
+                if qml.grouping.utils.are_identical_pauli_words(pauli_word, observable):
+                    ind = observables.index(observable)
+                    indices_this_group.append(available_indices[ind])
+                    # delete this observable and its index, so it cannot be found again
+                    observables.pop(ind)
+                    available_indices.pop(ind)
+                    break
+        indices.append(indices_this_group)
+
+    return indices
+
+
 class Hamiltonian(qml.operation.Observable):
     r"""Operator representing a Hamiltonian.
 
@@ -41,6 +70,13 @@ class Hamiltonian(qml.operation.Observable):
         observables (Iterable[Observable]): observables in the Hamiltonian expression, of same length as coeffs
         simplify (bool): Specifies whether the Hamiltonian is simplified upon initialization
                          (like-terms are combined). The default value is `False`.
+        grouping_type (str): If not None, compute and store information on how to group commuting
+            observables upon initialization. This information may be accessed when QNodes containing this
+            Hamiltonian are executed on devices. The string refers to the type of binary relation between Pauli words.
+            Can be ``'qwc'`` (qubit-wise commuting), ``'commuting'``, or ``'anticommuting'``.
+        method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+            can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
+        id (str): name to be assigned to this Hamiltonian instance
 
     **Example:**
 
@@ -93,17 +129,51 @@ class Hamiltonian(qml.operation.Observable):
 
     >>> H1 = qml.Hamiltonian(torch.tensor([1.]), [qml.PauliX(0)])
     >>> H2 = qml.Hamiltonian(torch.tensor([2., 3.]), [qml.PauliY(0), qml.PauliX(1)])
-    >>> H3 = qml.Hamiltonian(torch.tensor([1., 2., 3.]), [qml.PauliX(0), qml.PauliY(0), qml.PauliX(1)])
+    >>> obs3 = [qml.PauliX(0), qml.PauliY(0), qml.PauliX(1)]
+    >>> H3 = qml.Hamiltonian(torch.tensor([1., 2., 3.]), obs3)
     >>> H3.compare(H1 + H2)
     True
+
+    A Hamiltonian can store information on which commuting observables should be measured together in
+    a circuit:
+
+    >>> obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)]
+    >>> coeffs = np.array([1., 2., 3.])
+    >>> H = qml.Hamiltonian(coeffs, obs, grouping_type='qwc')
+    >>> H.grouping_indices
+    [[0, 1], [2]]
+
+    This attribute can be used to compute groups of coefficients and observables:
+
+    >>> grouped_coeffs = [coeffs[indices] for indices in H.grouping_indices]
+    >>> grouped_obs = [[H.ops[i] for i in indices] for indices in H.grouping_indices]
+    >>> grouped_coeffs
+    [tensor([1., 2.], requires_grad=True), tensor([3.], requires_grad=True)]
+    >>> grouped_obs
+    [[qml.PauliX(0), qml.PauliX(1)], [qml.PauliZ(0)]]
+
+    Devices that evaluate a Hamiltonian expectation by splitting it into its local observables can
+    use this information to reduce the number of circuits evaluated.
+
+    Note that one can compute the ``grouping_indices`` for an already initialized Hamiltonian by
+    using the :func:`compute_grouping <pennylane.Hamiltonian.compute_grouping>` method.
     """
 
     num_wires = qml.operation.AnyWires
     num_params = 1
     par_domain = "A"
     grad_method = "A"  # supports analytic gradients
 
-    def __init__(self, coeffs, observables, simplify=False, id=None, do_queue=True):
+    def __init__(
+        self,
+        coeffs,
+        observables,
+        simplify=False,
+        grouping_type=None,
+        method="rlf",
+        id=None,
+        do_queue=True,
+    ):
 
         if qml.math.shape(coeffs)[0] != len(observables):
             raise ValueError(
@@ -123,8 +193,16 @@ def __init__(self, coeffs, observables, simplify=False, id=None, do_queue=True):
 
         self.return_type = None
 
+        # attribute to store indices used to form groups of
+        # commuting observables, since recomputation is costly
+        self._grouping_indices = None
+
         if simplify:
             self.simplify()
+        if grouping_type is not None:
+            self._grouping_indices = qml.transforms.invisible(_compute_grouping_indices)(
+                self.ops, grouping_type=grouping_type, method=method
+            )
 
         coeffs_flat = [self._coeffs[i] for i in range(qml.math.shape(self._coeffs)[0])]
         # overwrite this attribute, now that we have the correct info
@@ -175,6 +253,31 @@ def wires(self):
     def name(self):
         return "Hamiltonian"
 
+    @property
+    def grouping_indices(self):
+        """Return the grouping indices attribute.
+
+        Returns:
+            list[list[int]]: indices needed to form groups of commuting observables
+        """
+        return self._grouping_indices
+
+    def compute_grouping(self, grouping_type="qwc", method="rlf"):
+        """
+        Compute groups of indices corresponding to commuting observables of this
+        Hamiltonian, and store it in the ``grouping_indices`` attribute.
+
+        Args:
+            grouping_type (str): The type of binary relation between Pauli words used to compute the grouping.
+                Can be ``'qwc'``, ``'commuting'``, or ``'anticommuting'``.
+            method (str): The graph coloring heuristic to use in solving minimum clique cover for grouping, which
+                can be ``'lf'`` (Largest First) or ``'rlf'`` (Recursive Largest First).
+        """
+
+        self._grouping_indices = qml.transforms.invisible(_compute_grouping_indices)(
+            self.ops, grouping_type=grouping_type, method=method
+        )
+
     def simplify(self):
         r"""Simplifies the Hamiltonian by combining like-terms.
 
@@ -186,6 +289,11 @@ def simplify(self):
         >>> print(H)
           (-1) [X0]
         + (1) [Y2]
+
+        .. warning::
+
+            Calling this method will reset ``grouping_indices`` to None, since
+            the observables it refers to are updated.
         """
         data = []
         ops = []
@@ -213,6 +321,8 @@ def simplify(self):
         self._coeffs = qml.math.stack(data) if data else []
         self.data = data
         self._ops = ops
+        # reset grouping, since the indices refer to the old observables and coefficients
+        self._grouping_indices = None
 
     def __str__(self):
         # Lambda function that formats the wires
@@ -277,7 +387,7 @@ def _obs_data(self):
 
         return data
 
-    def compare(self, H):
+    def compare(self, other):
         r"""Compares with another :class:`~Hamiltonian`, :class:`~.Observable`, or :class:`~.Tensor`,
         to determine if they are equivalent.
 
@@ -317,15 +427,15 @@ def compare(self, H):
         >>> ob1.compare(ob2)
         False
         """
-        if isinstance(H, Hamiltonian):
+        if isinstance(other, Hamiltonian):
             self.simplify()
-            H.simplify()
-            return self._obs_data() == H._obs_data()  # pylint: disable=protected-access
+            other.simplify()
+            return self._obs_data() == other._obs_data()  # pylint: disable=protected-access
 
-        if isinstance(H, (Tensor, Observable)):
+        if isinstance(other, (Tensor, Observable)):
             self.simplify()
             return self._obs_data() == {
-                (1, frozenset(H._obs_data()))  # pylint: disable=protected-access
+                (1, frozenset(other._obs_data()))  # pylint: disable=protected-access
             }
 
         raise ValueError("Can only compare a Hamiltonian, and a Hamiltonian/Observable/Tensor.")