feat: add PauliStringSum container

CleanDev-Fix · CleanDev-Fix · commit fd4353fb1553 · 2026-06-03T15:57:57.000-04:00
diff --git a/src/braket/quantum_information/__init__.py b/src/braket/quantum_information/__init__.py
@@ -12,8 +12,9 @@
 # language governing permissions and limitations under the License.
 
 """Provides utilities for working with quantum information concepts. It
-includes the PauliString class for representing and manipulating tensor products
-of Pauli operators.
+includes classes for representing and manipulating tensor products and weighted
+sums of Pauli operators.
 """
 
 from braket.quantum_information.pauli_string import PauliString  # noqa: F401
+from braket.quantum_information.pauli_string_sum import PauliStringSum  # noqa: F401
diff --git a/src/braket/quantum_information/pauli_string.py b/src/braket/quantum_information/pauli_string.py
@@ -234,6 +234,8 @@ def __mul__(self, other: PauliString) -> PauliString:
         See Also:
             `braket.quantum_information.PauliString.dot()`
         """
+        if not isinstance(other, PauliString):
+            return NotImplemented
         return self.dot(other)
 
     def __imul__(self, other: PauliString) -> PauliString:
diff --git a/src/braket/quantum_information/pauli_string_sum.py b/src/braket/quantum_information/pauli_string_sum.py
@@ -0,0 +1,238 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+
+from __future__ import annotations
+
+import numbers
+from collections.abc import Iterable
+from itertools import combinations
+
+from braket.circuits.observable import Observable, StandardObservable
+from braket.circuits.observables import I, Sum, TensorProduct, X, Y, Z
+from braket.quantum_information.pauli_string import PauliString
+
+_OBSERVABLE_TO_FACTOR = {I: "I", X: "X", Y: "Y", Z: "Z"}
+
+
+class PauliStringSum:
+    """A weighted sum of Pauli strings."""
+
+    def __init__(self, terms: Iterable[tuple[numbers.Number, str | PauliString]] = ()):
+        """Initializes a ``PauliStringSum``.
+
+        Args:
+            terms (Iterable[tuple[numbers.Number, str | PauliString]]): Pairs of coefficient and
+                Pauli string.
+        """
+        self._terms: dict[str, numbers.Number] = {}
+        for coefficient, pauli_string in terms:
+            self._add_term(coefficient, PauliString(pauli_string))
+
+    @classmethod
+    def from_list(cls, terms: Iterable[tuple[numbers.Number, str | PauliString]]) -> PauliStringSum:
+        """Builds a ``PauliStringSum`` from a list of weighted Pauli strings."""
+        return cls(terms)
+
+    @property
+    def terms(self) -> tuple[tuple[numbers.Number, PauliString], ...]:
+        """tuple[tuple[numbers.Number, PauliString], ...]: The weighted Pauli terms."""
+        return tuple(
+            (coefficient, PauliString(pauli)) for pauli, coefficient in self._terms.items()
+        )
+
+    @property
+    def qubit_count(self) -> int:
+        """int: The number of qubits in the largest Pauli string term."""
+        if not self._terms:
+            return 0
+        return max(PauliString(pauli).qubit_count for pauli in self._terms)
+
+    def to_list(self) -> list[tuple[numbers.Number, str]]:
+        """Returns a list representation of the weighted Pauli strings."""
+        return [(coefficient, pauli) for pauli, coefficient in self._terms.items()]
+
+    def to_sum(self) -> Sum:
+        """Converts the weighted Pauli strings into a circuit ``Sum`` observable."""
+        if not self._terms:
+            raise ValueError("Cannot convert an empty PauliStringSum to Sum")
+        observables = []
+        for coefficient, pauli in self.terms:
+            observables.append(coefficient * pauli.to_unsigned_observable(include_trivial=True))
+        return Sum(observables)
+
+    def commutes_with(self, other: PauliStringSum | PauliString | str) -> bool:
+        """Returns whether all terms commute with ``other``."""
+        other_sum = self._coerce(other)
+        return all(
+            self._pauli_strings_commute(left, right)
+            for _, left in self.terms
+            for _, right in other_sum.terms
+        )
+
+    def is_self_commuting(self) -> bool:
+        """Returns whether all terms in this sum commute with each other."""
+        return all(
+            self._pauli_strings_commute(left, right)
+            for (_, left), (_, right) in combinations(self.terms, 2)
+        )
+
+    @classmethod
+    def from_sum(cls, observable_sum: Sum) -> PauliStringSum:
+        """Builds a ``PauliStringSum`` from a circuit ``Sum`` observable."""
+        if not isinstance(observable_sum, Sum):
+            raise TypeError("Expected a Sum observable")
+        return cls(cls._term_from_observable(observable) for observable in observable_sum.summands)
+
+    def __add__(self, other: PauliStringSum | PauliString | str) -> PauliStringSum:
+        other_sum = self._coerce(other)
+        return PauliStringSum((*self.terms, *other_sum.terms))
+
+    def __radd__(self, other: PauliStringSum | PauliString | str) -> PauliStringSum:
+        return self + other
+
+    def __sub__(self, other: PauliStringSum | PauliString | str) -> PauliStringSum:
+        return self + (-1 * self._coerce(other))
+
+    def __rsub__(self, other: PauliStringSum | PauliString | str) -> PauliStringSum:
+        return self._coerce(other) + (-self)
+
+    def __neg__(self) -> PauliStringSum:
+        return -1 * self
+
+    def __mul__(self, other: numbers.Number | PauliString | str) -> PauliStringSum:
+        if isinstance(other, numbers.Number):
+            return PauliStringSum((coefficient * other, pauli) for coefficient, pauli in self.terms)
+        pauli = PauliString(other)
+        qubit_count = max(self.qubit_count, pauli.qubit_count)
+        right = self._pad_pauli_string(pauli, qubit_count)
+        return PauliStringSum(
+            (coefficient, self._pad_pauli_string(term, qubit_count).dot(right))
+            for coefficient, term in self.terms
+        )
+
+    def __rmul__(self, other: numbers.Number | PauliString | str) -> PauliStringSum:
+        if isinstance(other, numbers.Number):
+            return self * other
+        pauli = PauliString(other)
+        qubit_count = max(self.qubit_count, pauli.qubit_count)
+        left = self._pad_pauli_string(pauli, qubit_count)
+        return PauliStringSum(
+            (coefficient, left.dot(self._pad_pauli_string(term, qubit_count)))
+            for coefficient, term in self.terms
+        )
+
+    def __contains__(self, item: str | PauliString) -> bool:
+        _, pauli = self._canonical_term(PauliString(item))
+        return pauli in self._terms
+
+    def __getitem__(self, item: int) -> tuple[numbers.Number, PauliString]:
+        return self.terms[item]
+
+    def __iter__(self):
+        return iter(self.terms)
+
+    def __len__(self) -> int:
+        return len(self._terms)
+
+    def __eq__(self, other: PauliStringSum) -> bool:
+        if not isinstance(other, PauliStringSum):
+            return False
+        return self._terms == other._terms
+
+    def __repr__(self) -> str:
+        return f"PauliStringSum({self.to_list()})"
+
+    def _add_term(self, coefficient: numbers.Number, pauli_string: PauliString) -> None:
+        if not isinstance(coefficient, numbers.Number):
+            raise TypeError("PauliStringSum coefficients must be numbers")
+        coefficient, pauli = self._canonical_term(pauli_string, coefficient)
+        if coefficient == 0:
+            return
+        new_coefficient = self._terms.get(pauli, 0) + coefficient
+        if new_coefficient == 0:
+            self._terms.pop(pauli, None)
+        else:
+            self._terms[pauli] = new_coefficient
+
+    @staticmethod
+    def _canonical_term(
+        pauli_string: PauliString, coefficient: numbers.Number = 1
+    ) -> tuple[numbers.Number, str]:
+        factors = ["I"] * pauli_string.qubit_count
+        for qubit in range(pauli_string.qubit_count):
+            factors[qubit] = "IXYZ"[pauli_string[qubit]]
+        return coefficient * pauli_string.phase, f"+{''.join(factors)}"
+
+    @staticmethod
+    def _coerce(other: PauliStringSum | PauliString | str) -> PauliStringSum:
+        if isinstance(other, PauliStringSum):
+            return other
+        return PauliStringSum([(1, other)])
+
+    @staticmethod
+    def _pauli_strings_commute(left: PauliString, right: PauliString) -> bool:
+        anticommuting_factors = 0
+        qubit_count = max(left.qubit_count, right.qubit_count)
+        for qubit in range(qubit_count):
+            left_factor = left[qubit] if qubit < left.qubit_count else 0
+            right_factor = right[qubit] if qubit < right.qubit_count else 0
+            if left_factor and right_factor and left_factor != right_factor:
+                anticommuting_factors += 1
+        return anticommuting_factors % 2 == 0
+
+    @staticmethod
+    def _pad_pauli_string(pauli_string: PauliString, qubit_count: int) -> PauliString:
+        if pauli_string.qubit_count == qubit_count:
+            return pauli_string
+        factors = ["I"] * qubit_count
+        for qubit in range(pauli_string.qubit_count):
+            factors[qubit] = "IXYZ"[pauli_string[qubit]]
+        sign = "-" if pauli_string.phase < 0 else "+"
+        return PauliString(f"{sign}{''.join(factors)}")
+
+    @staticmethod
+    def _term_from_observable(observable: Observable) -> tuple[numbers.Number, str]:
+        coefficient = observable.coefficient
+        unscaled = observable._unscaled()
+        if isinstance(unscaled, StandardObservable):
+            factors = PauliStringSum._factors_from_standard(unscaled)
+        elif isinstance(unscaled, TensorProduct):
+            factors = PauliStringSum._factors_from_tensor_product(unscaled)
+        else:
+            raise TypeError(f"Unsupported observable type {type(observable).__name__}")
+        return coefficient, f"+{''.join(factors)}"
+
+    @staticmethod
+    def _factors_from_standard(observable: StandardObservable) -> list[str]:
+        factor = PauliStringSum._factor_from_standard(observable)
+        if observable.targets:
+            factors = ["I"] * (max(observable.targets) + 1)
+            factors[int(observable.targets[0])] = factor
+            return factors
+        return [factor]
+
+    @staticmethod
+    def _factors_from_tensor_product(observable: TensorProduct) -> list[str]:
+        if observable.targets:
+            factors = ["I"] * (max(observable.targets) + 1)
+            for factor, target in zip(observable.factors, observable.targets, strict=True):
+                factors[int(target)] = PauliStringSum._factor_from_standard(factor)
+            return factors
+        return [PauliStringSum._factor_from_standard(factor) for factor in observable.factors]
+
+    @staticmethod
+    def _factor_from_standard(observable: StandardObservable) -> str:
+        for observable_type, factor in _OBSERVABLE_TO_FACTOR.items():
+            if isinstance(observable, observable_type):
+                return factor
+        raise TypeError(f"Unsupported observable factor {type(observable).__name__}")
diff --git a/test/unit_tests/braket/quantum_information/test_pauli_string_sum.py b/test/unit_tests/braket/quantum_information/test_pauli_string_sum.py
@@ -0,0 +1,120 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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
+
+from braket.circuits.observables import H, X, Y, Z
+from braket.quantum_information import PauliString, PauliStringSum
+
+
+def test_initializes_from_weighted_list_and_combines_terms():
+    pauli_sum = PauliStringSum.from_list([(1.5, "XZ"), (2.0, "-XZ"), (3.0, "YY")])
+
+    assert pauli_sum.to_list() == [(-0.5, "+XZ"), (3.0, "+YY")]
+
+
+def test_addition_subtraction_and_scalar_multiplication():
+    first = PauliStringSum([(1.0, "XI")])
+    second = PauliStringSum([(2.0, "IZ")])
+
+    assert (first + second - PauliString("XI")).to_list() == [(2.0, "+IZ")]
+    assert (0.5 * second).to_list() == [(1.0, "+IZ")]
+    assert (-second).to_list() == [(-2.0, "+IZ")]
+
+
+def test_multiplication_by_pauli_string():
+    pauli_sum = PauliStringSum([(2.0, "XY"), (3.0, "ZZ")])
+
+    assert (pauli_sum * PauliString("YZ")).to_list() == [(-2.0, "+ZX"), (-3.0, "+XI")]
+
+
+def test_multiplication_by_pauli_string_pads_mixed_width_terms():
+    pauli_sum = PauliStringSum([(1.0, "X"), (2.0, "IZ")])
+
+    assert (pauli_sum * PauliString("ZZ")).to_list() == [(-1.0, "+YZ"), (2.0, "+ZI")]
+
+
+def test_left_multiplication_by_pauli_string_preserves_order():
+    pauli_sum = PauliStringSum([(2.0, "X")])
+
+    assert (pauli_sum * PauliString("Z")).to_list() == [(-2.0, "+Y")]
+    assert (PauliString("Z") * pauli_sum).to_list() == [(2.0, "+Y")]
+
+
+def test_indexing_and_membership():
+    pauli_sum = PauliStringSum([(2.0, "XI"), (3.0, "IZ")])
+
+    assert "XI" in pauli_sum
+    assert PauliString("-IZ") in pauli_sum
+    assert pauli_sum[0] == (2.0, PauliString("XI"))
+    assert list(pauli_sum) == [(2.0, PauliString("XI")), (3.0, PauliString("IZ"))]
+
+
+def test_to_sum_and_from_sum_round_trip():
+    pauli_sum = PauliStringSum([(2.0, "XY"), (-3.0, "ZI")])
+
+    observable_sum = pauli_sum.to_sum()
+
+    assert PauliStringSum.from_sum(observable_sum) == pauli_sum
+
+
+def test_equality_is_independent_of_term_insertion_order():
+    first = PauliStringSum([(1.0, "XI"), (2.0, "IZ")])
+    second = PauliStringSum([(2.0, "IZ"), (1.0, "XI")])
+
+    assert first == second
+
+
+def test_reverse_subtraction():
+    pauli_sum = PauliStringSum([(2.0, "XI")])
+
+    assert (PauliString("IZ") - pauli_sum).to_list() == [(1, "+IZ"), (-2.0, "+XI")]
+
+
+def test_from_sum_with_targeted_observables():
+    observable_sum = 2.0 * (X(0) @ Y(2)) - 3.0 * Z(1)
+
+    assert PauliStringSum.from_sum(observable_sum).to_list() == [
+        (2.0, "+XIY"),
+        (-3.0, "+IZ"),
+    ]
+
+
+def test_commutation_checks():
+    commuting_sum = PauliStringSum([(1.0, "XX"), (2.0, "YY")])
+    non_commuting_sum = PauliStringSum([(1.0, "XI"), (2.0, "ZI")])
+
+    assert commuting_sum.is_self_commuting()
+    assert commuting_sum.commutes_with("ZZ")
+    assert not non_commuting_sum.is_self_commuting()
+    assert not non_commuting_sum.commutes_with("YI")
+
+
+def test_empty_sum_cannot_convert_to_observable_sum():
+    with pytest.raises(ValueError, match="empty PauliStringSum"):
+        PauliStringSum().to_sum()
+
+
+def test_non_numeric_coefficients_are_rejected():
+    with pytest.raises(TypeError, match="coefficients must be numbers"):
+        PauliStringSum([("weight", "XI")])
+
+
+def test_from_sum_rejects_non_sum_observable():
+    with pytest.raises(TypeError, match="Expected a Sum observable"):
+        PauliStringSum.from_sum(X())
+
+
+def test_from_sum_rejects_non_pauli_observable_terms():
+    with pytest.raises(TypeError, match="Unsupported observable factor H"):
+        PauliStringSum.from_sum(H() + X())
