[Pauli6] Support for PauliWord and PauliSentence in qml.dot (#5027)

Work in progress, some concepts like how best to treat identities need some fine tuning. see also [Pauli1](#4989), [Pauli2](#5001), [Pauli3](#5003), [Pauli4](#5017), on top of which this PR builds; as well as [Pauli5](#5018) --------- Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com> Co-authored-by: Christina Lee <christina@xanadu.ai>
PennyLaneAI · Feb 9, 2024 · 81d948e · 81d948e
1 parent e69be08
commit 81d948e
Show file tree

Hide file tree

Showing 5 changed files with 226 additions and 19 deletions.
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
@@ -68,10 +68,12 @@
   sentence `ps1 = PauliSentence({pw1: 3.})`.
   You can now subtract `PauliWord` and `PauliSentence` instances, as well as scalars, from each other. For example `ps1 - pw1 - 1`.
   Overall, you can now intuitively construct `PauliSentence` operators like `0.5 * pw1 - 1.5 * ps1 + 2`.
+  You can now also use `qml.dot` with `PauliWord`, `PauliSentence` and operators, e.g. `qml.dot([0.5, -1.5, 2], [pw1, ps1, id_word])` with `id_word = PauliWord({})`.
   [(#4989)](https://github.com/PennyLaneAI/pennylane/pull/4989)
   [(#5001)](https://github.com/PennyLaneAI/pennylane/pull/5001)
   [(#5003)](https://github.com/PennyLaneAI/pennylane/pull/5003)
   [(#5017)](https://github.com/PennyLaneAI/pennylane/pull/5017)
+  [(#5027)](https://github.com/PennyLaneAI/pennylane/pull/5027)
 
 * `qml.matrix` now accepts `PauliWord` and `PauliSentence` instances, `qml.matrix(PauliWord({0:"X"}))`.
   [(#5018)](https://github.com/PennyLaneAI/pennylane/pull/5018)

diff --git a/pennylane/ops/functions/dot.py b/pennylane/ops/functions/dot.py
@@ -15,17 +15,21 @@
 This file contains the definition of the dot function, which computes the dot product between
 a vector and a list of operators.
 """
+# pylint: disable=too-many-branches
 from collections import defaultdict
 from typing import Sequence, Union, Callable
 
 import pennylane as qml
 from pennylane.operation import Operator, Tensor
 from pennylane.pulse import ParametrizedHamiltonian
+from pennylane.pauli import PauliWord, PauliSentence
 
 
 def dot(
-    coeffs: Sequence[Union[float, Callable]], ops: Sequence[Operator], pauli=False
-) -> Union[Operator, ParametrizedHamiltonian]:
+    coeffs: Sequence[Union[float, Callable]],
+    ops: Sequence[Union[Operator, PauliWord, PauliSentence]],
+    pauli=False,
+) -> Union[Operator, ParametrizedHamiltonian, PauliSentence]:
     r"""Returns the dot product between the ``coeffs`` vector and the ``ops`` list of operators.
 
     This function returns the following linear combination: :math:`\sum_{k} c_k O_k`, where
@@ -36,7 +40,8 @@ def dot(
         ops (Sequence[Operator]): sequence containing the operators of the linear combination
         pauli (bool, optional): If ``True``, a :class:`~.PauliSentence`
             operator is used to represent the linear combination. If False, a :class:`Sum` operator
-            is returned. Defaults to ``False``.
+            is returned. Defaults to ``False``. Note that when ``ops`` consists solely of ``PauliWord``
+            and ``PauliSentence`` instances, the function still returns a PennyLane operator when ``pauli=False``.
 
     Raises:
         ValueError: if the number of coefficients and operators does not match or if they are empty
@@ -54,6 +59,16 @@ def dot(
     1.1 * X(0)
     + 2.2 * Y(0)
 
+    Note that additions of the same operator are not executed by default.
+
+    >>> qml.dot([1., 1.], [qml.PauliX(0), qml.PauliX(0)])
+    PauliX(wires=[0]) + PauliX(wires=[0])
+
+    You can obtain a cleaner version by simplifying the resulting expression.
+
+    >>> qml.dot([1., 1.], [qml.PauliX(0), qml.PauliX(0)]).simplify()
+    2.0*(PauliX(wires=[0]))
+
     ``pauli=True`` can be used to construct a more efficient, simplified version of the operator.
     Note that it returns a :class:`~.PauliSentence`, which is not an :class:`~.Operator`. This
     specialized representation can be converted to an operator:
@@ -83,8 +98,17 @@ def dot(
     if any(callable(c) for c in coeffs):
         return ParametrizedHamiltonian(coeffs, ops)
 
+    # User-specified Pauli route
     if pauli:
-        return _pauli_dot(coeffs, ops)
+        if all(isinstance(pauli, (PauliWord, PauliSentence)) for pauli in ops):
+            # Use pauli arithmetic when ops are just PauliWord and PauliSentence instances
+            return _dot_pure_paulis(coeffs, ops)
+
+        # Else, transform all ops to pauli sentences
+        return _dot_with_ops_and_paulis(coeffs, ops)
+
+    # Convert possible PauliWord and PauliSentence instances to operation
+    ops = [op.operation() if isinstance(op, (PauliWord, PauliSentence)) else op for op in ops]
 
     # When casting a Hamiltonian to a Sum, we also cast its inner Tensors to Prods
     ops = [qml.prod(*op.obs) if isinstance(op, Tensor) else op for op in ops]
@@ -104,11 +128,18 @@ def dot(
     return operands[0] if len(operands) == 1 else qml.sum(*operands)
 
 
-def _pauli_dot(coeffs: Sequence[float], ops: Sequence[Operator]):
+def _dot_with_ops_and_paulis(coeffs: Sequence[float], ops: Sequence[Operator]):
+    """Compute dot when operators are a mix of pennylane operators, PauliWord and PauliSentence by turning them all into a PauliSentence instance.
+    Returns a PauliSentence instance"""
     pauli_words = defaultdict(lambda: 0)
     for coeff, op in zip(coeffs, ops):
         sentence = qml.pauli.pauli_sentence(op)
         for pw in sentence:
             pauli_words[pw] += sentence[pw] * coeff
 
     return qml.pauli.PauliSentence(pauli_words)
+
+
+def _dot_pure_paulis(coeffs: Sequence[float], ops: Sequence[Union[PauliWord, PauliSentence]]):
+    """Faster computation of dot when all ops are PauliSentences or PauliWords"""
+    return sum((c * op for c, op in zip(coeffs[1:], ops[1:])), start=coeffs[0] * ops[0])
diff --git a/pennylane/pauli/pauli_arithmetic.py b/pennylane/pauli/pauli_arithmetic.py
@@ -539,7 +539,7 @@ class PauliSentence(dict):
     .. note::
 
         An empty :class:`~.PauliSentence` will be treated as the additive
-        identity (i.e 0 * Identity on all wires).
+        identity (i.e ``0 * Identity()``).
 
     **Examples**
 
@@ -567,7 +567,7 @@ class PauliSentence(dict):
     >>> PauliSentence({})
     0 * I
 
-    We can compute commutators using the `PauliSentence.commutator()` method
+    We can compute commutators using the ``PauliSentence.commutator()`` method
 
     >>> op1 = PauliWord({0:"X", 1:"X"})
     >>> op2 = PauliWord({0:"Y"}) + PauliWord({1:"Y"})
@@ -576,6 +576,7 @@ class PauliSentence(dict):
     + 2j * X(0) @ Z(1)
 
     Or, alternatively, use :func:`~commutator`.
+
     >>> qml.commutator(op1, op2, pauli=True)
 
     Note that we need to specify ``pauli=True`` as :func:`~.commutator` returns PennyLane operators by default.

diff --git a/tests/ops/functions/test_dot.py b/tests/ops/functions/test_dot.py
@@ -18,7 +18,33 @@
 
 import pennylane as qml
 from pennylane.ops import Hamiltonian, Prod, SProd, Sum
-from pennylane.pauli.pauli_arithmetic import PauliSentence
+from pennylane.pauli.pauli_arithmetic import PauliWord, PauliSentence, I, X, Y, Z
+
+pw1 = PauliWord({0: I, 1: X, 2: Y})
+pw2 = PauliWord({0: Z, 2: Z, 4: Z})
+pw3 = PauliWord({"a": X, "b": X, "c": Z})
+pw4 = PauliWord({"a": Y, "b": Z, "c": X})
+pw5 = PauliWord({4: X, 5: X})
+
+ps1 = PauliSentence({pw1: 1.0, pw2: 2.0})
+ps2 = PauliSentence({pw3: 1.0, pw4: 2.0})
+
+op1 = qml.prod(qml.PauliX(1), qml.PauliY(2))
+op2 = qml.prod(qml.PauliZ(0), qml.PauliZ(2), qml.PauliZ(4))
+op3 = qml.prod(qml.PauliX("a"), qml.PauliX("b"), qml.PauliZ("c"))
+op4 = qml.prod(qml.PauliY("a"), qml.PauliZ("b"), qml.PauliX("c"))
+op5 = qml.prod(qml.PauliX(4), qml.PauliX(5))
+
+pw_id = PauliWord({})
+ps_id = PauliSentence({pw_id: 1.0})
+op_id = qml.Identity(0)
+
+X0, Y0, Z0 = PauliWord({0: "X"}), PauliWord({0: "Y"}), PauliWord({0: "Z"})
+XX, YY, ZZ = PauliWord({0: "X", 1: "X"}), PauliWord({0: "Y", 1: "Y"}), PauliWord({0: "Z", 1: "Z"})
+
+H1 = X0 + Y0 + Z0
+H2 = XX + YY + ZZ
+H3 = 1.0 * X0 + 2.0 * Y0 + 3.0 * Z0
 
 
 class TestDotSum:
@@ -160,9 +186,92 @@ def test_dot_jax(self, dtype):
         )
         assert qml.equal(op_sum, op_sum_2)
 
-
-coeffs = [0.12345, 1.2345, 12.345, 123.45, 1234.5, 12345]
-ops = [
+    data_just_words_pauli_false = (
+        ([1.0, 2.0, 3.0], [pw1, pw2, pw_id], [op1, op2, op_id]),
+        ([1.0, 2.0, 3.0], [pw1, pw2, pw_id], [op1, op2, op_id]),
+        ([1.0, 2.0, 3.0, 4.0], [pw1, pw2, pw3, pw_id], [op1, op2, op3, op_id]),
+        ([1.0, 2.0, 3j, 4j], [pw1, pw2, pw3, pw_id], [op1, op2, op3, op_id]),
+        ([1, 1, 1], [pw1, pw1, pw1], [op1, op1, op1]),
+        ([1, 1, 1], [pw_id, pw_id, pw_id], [op_id, op_id, op_id]),
+    )
+
+    @pytest.mark.parametrize("coeff, words, ops", data_just_words_pauli_false)
+    def test_dot_with_just_words_pauli_false(self, coeff, words, ops):
+        """Test operators that are just pauli words"""
+        dot_res = qml.dot(coeff, words, pauli=False)
+        true_res = qml.dot(coeff, ops)
+        assert dot_res == true_res
+
+    data_words_and_sentences_pauli_false = (
+        (
+            [1.0, 2.0, 3.0],
+            [pw1, pw2, ps1],
+        ),
+        (
+            [1.0, 2.0, 3.0, 4.0],
+            [X0, Y0, XX, YY],
+        ),
+        (
+            [1.0, 2.0, 3.0],
+            [H1, H2, H3],
+        ),
+        (
+            [1.0, 2.0, 3.0],
+            [pw3, pw4, ps2],
+        ),  # comparisons for Sum objects with string valued wires
+    )
+
+    @pytest.mark.parametrize("coeff, ops", data_words_and_sentences_pauli_false)
+    def test_dot_with_words_and_sentences_pauli_false(self, coeff, ops):
+        """Test operators that are a mix of pauli words and pauli sentences"""
+        dot_res = qml.dot(coeff, ops, pauli=False)
+        true_res = qml.dot(coeff, [op.operation() for op in ops], pauli=False)
+        assert dot_res == true_res
+
+    data_op_words_and_sentences_pauli_false = (
+        (
+            [1.0, 2.0, 3.0, 1j],
+            [pw1, pw2, ps1, op1],
+            qml.sum(qml.s_prod(3.0 * 1 + 1 + 1j, op1), qml.s_prod(3 * 2.0 + 2, op2)),
+        ),
+        (
+            [2.0, 3.0, 1j],
+            [pw2, ps1, op1],
+            qml.sum(qml.s_prod(3.0 * 1 + 1j, op1), qml.s_prod(3 * 2.0 + 2, op2)),
+        ),
+        (
+            [2.0, 3.0, 1j],
+            [pw2, ps1, op5],
+            qml.sum(qml.s_prod(3.0 * 1, op1), qml.s_prod(3 * 2.0 + 2, op2), qml.s_prod(1j, op5)),
+        ),
+        (
+            [1.0, 2.0, 3.0, 1j],
+            [pw3, pw4, ps2, op3],
+            qml.sum(qml.s_prod(3.0 * 1 + 1 + 1j, op3), qml.s_prod(3 * 2.0 + 2, op4)),
+        ),  # string valued wires
+    )
+
+    @pytest.mark.parametrize("coeff, ops, res", data_op_words_and_sentences_pauli_false)
+    def test_dot_with_ops_words_and_sentences(self, coeff, ops, res):
+        """Test operators that are a mix of PL operators, pauli words and pauli sentences with pauli=False (i.e. returning operators)"""
+        dot_res = qml.dot(coeff, ops, pauli=False).simplify()
+        assert dot_res == res
+
+    def test_identities_with_pauli_words_pauli_false(self):
+        """Test that identities in form of empty PauliWords are treated correctly"""
+        res = qml.dot([2.0, 2.0], [pw_id, pw1], pauli=False)
+        true_res = qml.s_prod(2, qml.sum(op_id, op1))
+        assert res == true_res
+
+    def test_identities_with_pauli_sentences_pauli_false(self):
+        """Test that identities in form of PauliSentences with empty PauliWords are treated correctly"""
+        res = qml.dot([2.0, 2.0], [ps_id, pw1], pauli=False)
+        true_res = qml.s_prod(2, qml.sum(op_id, op1))
+        assert res == true_res
+
+
+coeffs0 = [0.12345, 1.2345, 12.345, 123.45, 1234.5, 12345]
+ops0 = [
     qml.PauliX(0),
     qml.PauliY(1),
     qml.PauliZ(2),
@@ -177,15 +286,15 @@ class TestDotPauliSentence:
 
     def test_dot_returns_pauli_sentence(self):
         """Test that the dot function returns a PauliSentence class."""
-        ps = qml.dot(coeffs, ops, pauli=True)
+        ps = qml.dot(coeffs0, ops0, pauli=True)
         assert isinstance(ps, PauliSentence)
 
     def test_coeffs_and_ops(self):
         """Test that the coefficients and operators of the returned PauliSentence are correct."""
-        ps = qml.dot(coeffs, ops, pauli=True)
+        ps = qml.dot(coeffs0, ops0, pauli=True)
         h = ps.hamiltonian()
-        assert qml.math.allequal(h.coeffs, coeffs)
-        assert all(qml.equal(op1, op2) for op1, op2 in zip(h.ops, ops))
+        assert qml.math.allequal(h.coeffs, coeffs0)
+        assert all(qml.equal(op1, op2) for op1, op2 in zip(h.ops, ops0))
 
     def test_dot_simplifies_linear_combination(self):
         """Test that the dot function groups equal pauli words."""
@@ -200,9 +309,9 @@ def test_dot_simplifies_linear_combination(self):
     def test_dot_returns_hamiltonian_simplified(self):
         """Test that hamiltonian computed from the PauliSentence created by the dot function is equal
         to the simplified hamiltonian."""
-        ps = qml.dot(coeffs, ops, pauli=True)
+        ps = qml.dot(coeffs0, ops0, pauli=True)
         h_ps = ps.hamiltonian()
-        h = Hamiltonian(coeffs, ops)
+        h = Hamiltonian(coeffs0, ops0)
         h.simplify()
         assert qml.equal(h_ps, h)
 
@@ -275,3 +384,69 @@ def test_dot_jax(self):
             }
         )
         assert ps == ps_2
+
+    data_just_words = (
+        (
+            [1.0, 2.0, 3.0, 4.0],
+            [pw1, pw2, pw3, pw_id],
+            PauliSentence({pw1: 1.0, pw2: 2.0, pw3: 3.0, pw_id: 4.0}),
+        ),
+        (
+            [1.5, 2.5, 3.5, 4.5j],
+            [pw1, pw2, pw3, pw_id],
+            PauliSentence({pw1: 1.5, pw2: 2.5, pw3: 3.5, pw_id: 4.5j}),
+        ),
+        ([1.5, 2.5, 3.5], [pw3, pw2, pw1], PauliSentence({pw3: 1.5, pw2: 2.5, pw1: 3.5})),
+        ([1, 1, 1], [PauliWord({0: "X"})] * 3, PauliSentence({PauliWord({0: "X"}): 3.0})),
+        ([1, 1, 1], [PauliWord({})] * 3, PauliSentence({PauliWord({}): 3.0})),
+    )
+
+    @pytest.mark.parametrize("coeff, ops, res", data_just_words)
+    def test_dot_with_just_words(self, coeff, ops, res):
+        """Test operators that are just pauli words"""
+        dot_res = qml.dot(coeff, ops, pauli=True)
+        assert dot_res == res
+
+    data_words_and_sentences = (
+        ([1.0, 2.0, 3.0], [pw1, pw2, ps1], PauliSentence({pw1: 3.0 * 1 + 1, pw2: 3 * 2.0 + 2})),
+        ([1.0, 2.0, 3.0], [pw3, pw4, ps2], PauliSentence({pw3: 3.0 * 1 + 1, pw4: 3 * 2.0 + 2})),
+    )
+
+    @pytest.mark.parametrize("coeff, ops, res", data_words_and_sentences)
+    def test_dot_with_words_and_sentences(self, coeff, ops, res):
+        """Test operators that are a mix of pauli words and pauli sentences"""
+        dot_res = qml.dot(coeff, ops, pauli=True)
+        assert dot_res == res
+
+    data_op_words_and_sentences = (
+        (
+            [1.0, 2.0, 3.0, 1j],
+            [pw1, pw2, ps1, op1],
+            PauliSentence({pw1: 3.0 * 1 + 1 + 1j, pw2: 3 * 2.0 + 2}),
+        ),
+        (
+            [1.0, 2.0, 3.0, 1j],
+            [pw3, pw4, ps2, op3],
+            PauliSentence({pw3: 3.0 * 1 + 1 + 1j, pw4: 3 * 2.0 + 2}),
+        ),
+        ([2.0, 3.0, 1j], [pw2, ps1, op1], PauliSentence({pw1: 3.0 * 1 + 1j, pw2: 3 * 2.0 + 2})),
+        ([2.0, 3.0, 1j], [pw2, ps1, op5], PauliSentence({pw1: 3.0 * 1, pw2: 3 * 2.0 + 2, pw5: 1j})),
+    )
+
+    @pytest.mark.parametrize("coeff, ops, res", data_op_words_and_sentences)
+    def test_dot_with_ops_words_and_sentences(self, coeff, ops, res):
+        """Test operators that are a mix of PL operators, pauli words and pauli sentences"""
+        dot_res = qml.dot(coeff, ops, pauli=True)
+        assert dot_res == res
+
+    def test_identities_with_pauli_words_pauli_true(self):
+        """Test that identities in form of empty PauliWords are treated correctly"""
+        res = qml.dot([2.0, 2.0], [pw_id, pw1], pauli=True)
+        true_res = PauliSentence({pw_id: 2, pw1: 2})
+        assert res == true_res
+
+    def test_identities_with_pauli_sentences_pauli_true(self):
+        """Test that identities in form of PauliSentences with empty PauliWords are treated correctly"""
+        res = qml.dot([2.0, 2.0], [ps_id, pw1], pauli=True)
+        true_res = PauliSentence({pw_id: 2.0, pw1: 2.0})
+        assert res == true_res
diff --git a/tests/pauli/test_pauli_arithmetic.py b/tests/pauli/test_pauli_arithmetic.py
@@ -43,8 +43,6 @@
 
 words = [pw1, pw2, pw3, pw4]
 
-words = [pw1, pw2, pw3, pw4]
-
 ps1 = PauliSentence({pw1: 1.23, pw2: 4j, pw3: -0.5})
 ps2 = PauliSentence({pw1: -1.23, pw2: -4j, pw3: 0.5})
 ps1_hamiltonian = PauliSentence({pw1: 1.23, pw2: 4, pw3: -0.5})