diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index c73d3bc09a4..15500bc8a53 100644 --- a/.github/CHANGELOG.md +++ b/.github/CHANGELOG.md @@ -2,6 +2,24 @@

New features since last release

+* 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 index 590c8dd5335..0bd9104a8fb 100644 --- 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,18 +68,42 @@ 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) @@ -86,13 +111,11 @@ def hamiltonian_expand(tape, group=True): 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 index 069e681550a..96197513305 100644 --- 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,9 +129,34 @@ 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 ` method. """ num_wires = qml.operation.AnyWires @@ -103,7 +164,16 @@ class Hamiltonian(qml.operation.Observable): 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.") diff --git a/tests/ops/test_hamiltonian.py b/tests/ops/test_hamiltonian.py index 2d429b9dc7f..5414055874f 100644 --- a/tests/ops/test_hamiltonian.py +++ b/tests/ops/test_hamiltonian.py @@ -404,6 +404,86 @@ def test_hamiltonian_matmul(self): assert H.compare(H1 @ H2) +class TestGrouping: + """Tests for the grouping functionality""" + + def test_grouping_is_correct_kwarg(self): + """Basic test checking that grouping with a kwarg works as expected""" + a = qml.PauliX(0) + b = qml.PauliX(1) + c = qml.PauliZ(0) + obs = [a, b, c] + coeffs = [1.0, 2.0, 3.0] + + H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc") + assert H.grouping_indices == [[0, 1], [2]] + + def test_grouping_is_correct_compute_grouping(self): + """Basic test checking that grouping with compute_grouping works as expected""" + a = qml.PauliX(0) + b = qml.PauliX(1) + c = qml.PauliZ(0) + obs = [a, b, c] + coeffs = [1.0, 2.0, 3.0] + + H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc") + H.compute_grouping() + assert H.grouping_indices == [[0, 1], [2]] + + def test_grouping_for_non_groupable_hamiltonians(self): + """Test that grouping is computed correctly, even if no observables commute""" + a = qml.PauliX(0) + b = qml.PauliY(0) + c = qml.PauliZ(0) + obs = [a, b, c] + coeffs = [1.0, 2.0, 3.0] + + H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc") + assert H.grouping_indices == [[0], [1], [2]] + + def test_grouping_is_reset_when_simplifying(self): + """Tests that calling simplify() resets the grouping""" + obs = [qml.PauliX(0), qml.PauliX(1), qml.PauliZ(0)] + coeffs = [1.0, 2.0, 3.0] + + H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc") + assert H.grouping_indices is not None + + H.simplify() + assert H.grouping_indices is None + + def test_grouping_does_not_alter_queue(self): + """Tests that grouping is invisible to the queue.""" + a = qml.PauliX(0) + b = qml.PauliX(1) + c = qml.PauliZ(0) + obs = [a, b, c] + coeffs = [1.0, 2.0, 3.0] + + with qml.tape.QuantumTape() as tape: + H = qml.Hamiltonian(coeffs, obs, grouping_type="qwc") + + assert tape.queue == [a, b, c, H] + + def test_grouping_method_can_be_set(self): + r"""Tests that the grouping method can be controlled by kwargs. + This is done by changing from default to 'rlf' and checking the result.""" + a = qml.PauliX(0) + b = qml.PauliX(1) + c = qml.PauliZ(0) + obs = [a, b, c] + coeffs = [1.0, 2.0, 3.0] + + # compute grouping during construction + H2 = qml.Hamiltonian(coeffs, obs, grouping_type="qwc", method="lf") + assert H2.grouping_indices == [[2, 1], [0]] + + # compute grouping separately + H3 = qml.Hamiltonian(coeffs, obs, grouping_type=None) + H3.compute_grouping(method="lf") + assert H3.grouping_indices == [[2, 1], [0]] + + class TestHamiltonianEvaluation: """Test the usage of a Hamiltonian as an observable""" @@ -453,11 +533,12 @@ def circuit(): assert pars == [0.1, 3.0] +@pytest.mark.parametrize("simplify", [True, False]) +@pytest.mark.parametrize("group", [None, "qwc"]) class TestHamiltonianDifferentiation: """Test that the Hamiltonian coefficients are differentiable""" - @pytest.mark.parametrize("simplify", [True, False]) - def test_vqe_differentiation_paramshift(self, simplify): + def test_vqe_differentiation_paramshift(self, simplify, group): """Test the parameter-shift method by comparing the differentiation of linearly combined subcircuits with the differentiation of a Hamiltonian expectation""" coeffs = np.array([-0.05, 0.17]) @@ -469,7 +550,12 @@ def circuit(coeffs, param): qml.RX(param, wires=0) qml.RY(param, wires=0) return qml.expval( - qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify) + qml.Hamiltonian( + coeffs, + [qml.PauliX(0), qml.PauliZ(0)], + simplify=simplify, + grouping_type=group, + ) ) grad_fn = qml.grad(circuit) @@ -489,8 +575,7 @@ def combine(coeffs, param): assert np.allclose(grad[0], grad_expected[0]) assert np.allclose(grad[1], grad_expected[1]) - @pytest.mark.parametrize("simplify", [True, False]) - def test_vqe_differentiation_autograd(self, simplify): + def test_vqe_differentiation_autograd(self, simplify, group): """Test the autograd interface by comparing the differentiation of linearly combined subcircuits with the differentiation of a Hamiltonian expectation""" coeffs = pnp.array([-0.05, 0.17], requires_grad=True) @@ -502,7 +587,12 @@ def circuit(coeffs, param): qml.RX(param, wires=0) qml.RY(param, wires=0) return qml.expval( - qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify) + qml.Hamiltonian( + coeffs, + [qml.PauliX(0), qml.PauliZ(0)], + simplify=simplify, + grouping_type=group, + ) ) grad_fn = qml.grad(circuit) @@ -522,8 +612,7 @@ def combine(coeffs, param): assert np.allclose(grad[0], grad_expected[0]) assert np.allclose(grad[1], grad_expected[1]) - @pytest.mark.parametrize("simplify", [True, False]) - def test_vqe_differentiation_jax(self, simplify): + def test_vqe_differentiation_jax(self, simplify, group): """Test the jax interface by comparing the differentiation of linearly combined subcircuits with the differentiation of a Hamiltonian expectation""" @@ -538,7 +627,12 @@ def circuit(coeffs, param): qml.RX(param, wires=0) qml.RY(param, wires=0) return qml.expval( - qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify) + qml.Hamiltonian( + coeffs, + [qml.PauliX(0), qml.PauliZ(0)], + simplify=simplify, + grouping_type=group, + ) ) grad_fn = jax.grad(circuit) @@ -558,8 +652,7 @@ def combine(coeffs, param): assert np.allclose(grad[0], grad_expected[0]) assert np.allclose(grad[1], grad_expected[1]) - @pytest.mark.parametrize("simplify", [True, False]) - def test_vqe_differentiation_torch(self, simplify): + def test_vqe_differentiation_torch(self, simplify, group): """Test the torch interface by comparing the differentiation of linearly combined subcircuits with the differentiation of a Hamiltonian expectation""" @@ -573,7 +666,12 @@ def circuit(coeffs, param): qml.RX(param, wires=0) qml.RY(param, wires=0) return qml.expval( - qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify) + qml.Hamiltonian( + coeffs, + [qml.PauliX(0), qml.PauliZ(0)], + simplify=simplify, + grouping_type=group, + ) ) res = circuit(coeffs, param) @@ -600,8 +698,7 @@ def combine(coeffs, param): assert np.allclose(grad[0], grad_expected[0]) assert np.allclose(grad[1], grad_expected[1]) - @pytest.mark.parametrize("simplify", [True, False]) - def test_vqe_differentiation_tf(self, simplify): + def test_vqe_differentiation_tf(self, simplify, group): """Test the tf interface by comparing the differentiation of linearly combined subcircuits with the differentiation of a Hamiltonian expectation""" @@ -615,7 +712,12 @@ def circuit(coeffs, param): qml.RX(param, wires=0) qml.RY(param, wires=0) return qml.expval( - qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)], simplify=simplify) + qml.Hamiltonian( + coeffs, + [qml.PauliX(0), qml.PauliZ(0)], + simplify=simplify, + grouping_type=group, + ) ) with tf.GradientTape() as tape: diff --git a/tests/test_vqe.py b/tests/test_vqe.py index 8754804ab90..3c74455ccbc 100644 --- a/tests/test_vqe.py +++ b/tests/test_vqe.py @@ -1253,7 +1253,7 @@ def circuit(): res_expected.append(circuit()) - res_expected = np.sum(c * r for c, r in zip(coeffs, res_expected)) + res_expected = np.sum([c * r for c, r in zip(coeffs, res_expected)]) assert np.isclose(res, res_expected, atol=tol) diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py index e1cde7c5dd3..cbfcd758b58 100644 --- a/tests/transforms/test_hamiltonian_expand.py +++ b/tests/transforms/test_hamiltonian_expand.py @@ -98,6 +98,22 @@ def test_hamiltonians_no_grouping(self, tape, output): assert np.isclose(output, expval) + def test_grouping_is_used(self): + """Test that the grouping in a Hamiltonian is used""" + H = qml.Hamiltonian( + [1.0, 2.0, 3.0], [qml.PauliZ(0), qml.PauliX(1), qml.PauliX(0)], grouping_type="qwc" + ) + 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) + + tapes, fn = qml.transforms.hamiltonian_expand(tape, group=False) + assert len(tapes) == 2 + def test_number_of_tapes(self): """Tests that the the correct number of tapes is produced""" diff --git a/tests/transforms/test_measurement_grouping.py b/tests/transforms/test_measurement_grouping.py new file mode 100644 index 00000000000..527b1695e49 --- /dev/null +++ b/tests/transforms/test_measurement_grouping.py @@ -0,0 +1,37 @@ +# Copyright 2018-2021 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. + +import pytest +import numpy as np +import pennylane as qml + + +def test_measurement_grouping(): + """Test that measurement grouping works as expected.""" + + with qml.tape.QuantumTape() as tape: + qml.RX(0.1, wires=0) + qml.RX(0.2, wires=1) + qml.CNOT(wires=[0, 1]) + qml.CNOT(wires=[1, 2]) + + obs = [qml.PauliZ(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliX(2)] + coeffs = [2.0, -0.54, 0.1] + + tapes, fn = qml.transforms.measurement_grouping(tape, obs, coeffs) + assert len(tapes) == 2 + + dev = qml.device("default.qubit", wires=3) + res = fn(dev.batch_execute(tapes)) + assert np.isclose(res, 2.0007186)