feat: Add PauliSum container (#1270)

CleanDev-Fix · web-flow · commit 1299de6678ac · 2026-06-20T15:45:28.000-07: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_sum import PauliStringSum, PauliSum  # 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:
@@ -255,6 +257,55 @@ def __imul__(self, other: PauliString) -> PauliString:
         """
         return self.dot(other, inplace=True)
 
+    def __add__(self, other: PauliString | str):
+        """Operator overload for addition with another Pauli string.
+
+        Args:
+            other (PauliString | str): The Pauli string to add.
+
+        Returns:
+            PauliSum: The sum containing both Pauli strings.
+        """
+        if not isinstance(other, (PauliString, str)):
+            return NotImplemented
+        from braket.quantum_information.pauli_sum import PauliSum  # noqa: PLC0415
+
+        return PauliSum([(1, self), (1, other)])
+
+    def __radd__(self, other: PauliString | str):
+        """Operator overload for reverse addition with another Pauli string.
+
+        Args:
+            other (PauliString | str): The Pauli string to add.
+
+        Returns:
+            PauliSum: The sum containing both Pauli strings.
+        """
+        if not isinstance(other, (PauliString, str)):
+            return NotImplemented
+        from braket.quantum_information.pauli_sum import PauliSum  # noqa: PLC0415
+
+        return PauliSum([(1, other), (1, self)])
+
+    def commutes_with(self, other: PauliString | str) -> bool:
+        """Returns whether this Pauli string commutes with another Pauli string.
+
+        Args:
+            other (PauliString | str): The Pauli string to check against.
+
+        Returns:
+            bool: Whether the Pauli strings commute.
+        """
+        other_pauli = PauliString(other)
+        anticommuting_factors = 0
+        qubit_count = max(self._qubit_count, other_pauli._qubit_count)
+        for qubit in range(qubit_count):
+            left_factor = self[qubit] if qubit < self._qubit_count else 0
+            right_factor = other_pauli[qubit] if qubit < other_pauli._qubit_count else 0
+            if left_factor and right_factor and left_factor != right_factor:
+                anticommuting_factors += 1
+        return anticommuting_factors % 2 == 0
+
     def power(self, n: int, inplace: bool = False) -> PauliString:
         """Composes Pauli string with itself n times.
 
diff --git a/src/braket/quantum_information/pauli_sum.py b/src/braket/quantum_information/pauli_sum.py
@@ -0,0 +1,252 @@
+# 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, starmap
+
+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 PauliSum:
+    """A weighted sum of Pauli strings."""
+
+    def __init__(self, terms: Iterable[tuple[numbers.Number, str | PauliString]] = ()):
+        """Initializes a ``PauliSum``.
+
+        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))
+        self._all_terms_commute = self._compute_all_terms_commute()
+
+    @property
+    def all_terms_commute(self) -> bool:
+        """bool: Whether all terms in the sum commute with each other."""
+        return self._all_terms_commute
+
+    @classmethod
+    def from_list(cls, terms: Iterable[tuple[numbers.Number, str | PauliString]]) -> PauliSum:
+        """Builds a ``PauliSum`` 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 PauliSum 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: PauliSum | PauliString | str) -> bool:
+        """Returns whether all terms commute with ``other``."""
+        other_sum = self._coerce(other)
+        if not self.all_terms_commute or not other_sum.all_terms_commute:
+            return False
+        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 self._all_terms_commute
+
+    @classmethod
+    def from_sum(cls, observable_sum: Sum) -> PauliSum:
+        """Builds a ``PauliSum`` 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: PauliSum | PauliString | str) -> PauliSum:
+        other_sum = self._coerce(other)
+        return PauliSum((*self.terms, *other_sum.terms))
+
+    def __radd__(self, other: PauliSum | PauliString | str) -> PauliSum:
+        return self + other
+
+    def __sub__(self, other: PauliSum | PauliString | str) -> PauliSum:
+        return self + (-1 * self._coerce(other))
+
+    def __rsub__(self, other: PauliSum | PauliString | str) -> PauliSum:
+        return self._coerce(other) + (-self)
+
+    def __neg__(self) -> PauliSum:
+        return -1 * self
+
+    def __mul__(self, other: numbers.Number | PauliString | str) -> PauliSum:
+        if isinstance(other, numbers.Number):
+            return PauliSum((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 PauliSum(
+            (coefficient, self._pad_pauli_string(term, qubit_count).dot(right))
+            for coefficient, term in self.terms
+        )
+
+    def __rmul__(self, other: numbers.Number | PauliString | str) -> PauliSum:
+        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 PauliSum(
+            (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: PauliSum) -> bool:
+        if not isinstance(other, PauliSum):
+            return False
+        return self._terms == other._terms
+
+    def __repr__(self) -> str:
+        return f"PauliSum({self.to_list()})"
+
+    def _add_term(self, coefficient: numbers.Number, pauli_string: PauliString) -> None:
+        if not isinstance(coefficient, numbers.Number):
+            raise TypeError("PauliSum 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
+
+    def _compute_all_terms_commute(self) -> bool:
+        if len(self._terms) <= 1:
+            return True
+        paulis = [PauliString(pauli) for pauli in self._terms]
+        return all(starmap(self._pauli_strings_commute, combinations(paulis, 2)))
+
+    @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: PauliSum | PauliString | str) -> PauliSum:
+        if isinstance(other, PauliSum):
+            return other
+        return PauliSum([(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 = PauliSum._factors_from_standard(unscaled)
+        elif isinstance(unscaled, TensorProduct):
+            factors = PauliSum._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 = PauliSum._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)] = PauliSum._factor_from_standard(factor)
+            return factors
+        return [PauliSum._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__}")
+
+
+PauliStringSum = PauliSum
diff --git a/test/unit_tests/braket/quantum_information/test_pauli_string.py b/test/unit_tests/braket/quantum_information/test_pauli_string.py
@@ -21,7 +21,7 @@
 from braket.circuits import gates
 from braket.circuits.circuit import Circuit
 from braket.circuits.observables import I, X, Y, Z
-from braket.quantum_information import PauliString
+from braket.quantum_information import PauliString, PauliStringSum
 
 ORDER = ["I", "X", "Y", "Z"]
 PAULI_INDEX_MATRICES = {
@@ -172,6 +172,35 @@ def test_dot(circ_arg_1, circ_arg_2, circ_res):
     assert circ1 == PauliString(circ_res)
 
 
+def test_pauli_string_addition_builds_pauli_string_sum():
+    pauli_sum = PauliString("XI") + PauliString("-IZ")
+
+    assert pauli_sum == PauliStringSum([(1, "XI"), (1, "-IZ")])
+
+
+def test_pauli_string_addition_accepts_strings_and_reverse_addition():
+    assert PauliString("XI") + "IZ" == PauliStringSum([(1, "XI"), (1, "IZ")])
+    assert "IZ" + PauliString("XI") == PauliStringSum([(1, "IZ"), (1, "XI")])
+
+
+def test_pauli_string_operator_overloads_reject_unsupported_types():
+    pauli_string = PauliString("XI")
+
+    with pytest.raises(TypeError):
+        pauli_string * "IZ"
+    with pytest.raises(TypeError):
+        pauli_string + 1
+    with pytest.raises(TypeError):
+        1 + pauli_string
+
+
+def test_pauli_string_commutation_check():
+    assert PauliString("XX").commutes_with("YY")
+    assert PauliString("X").commutes_with("IZ")
+    assert PauliString("IZ").commutes_with("X")
+    assert not PauliString("XI").commutes_with("ZI")
+
+
 @pytest.mark.xfail(raises=ValueError)
 @pytest.mark.parametrize(
     "circ1, circ2, operation",
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
