Add sparse_matrix() to PauliString and PauliSum

cirq-core/cirq/ops/linear_combinations.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import sympy
+from scipy import sparse
 from sympy.logic.boolalg import And, Not, Or, Xor
 
 from cirq import linalg, protocols, qis, value
@@ -32,8 +33,6 @@
 from cirq.value.linear_dict import _format_terms
 
 if TYPE_CHECKING:
-    from scipy.sparse import csr_matrix
-
     import cirq
 
 UnitPauliStringT = frozenset[tuple[raw_types.Qid, pauli_gates.Pauli]]
@@ -589,6 +588,50 @@ def matrix(self, qubits: Iterable[raw_types.Qid] | None = None) -> np.ndarray:
             result += coeff * op.matrix(qubits)
         return result
 
+    def sparse_matrix(self, qubits: Iterable[raw_types.Qid] | None = None) -> sparse.csr_matrix:
+        """Returns the sparse matrix of this `PauliSum` in the computational basis of the qubits.
+
+        For each term we build the sparse matrix via direct bit-manipulation
+        (see `PauliString.sparse_matrix`) and collect its non-zero entries as
+        COO (COOrdinate) triplets (data, row, col).  All triplets are
+        concatenated and a single sparse matrix is built at the end, avoiding
+        the overhead of adding sparse matrices term-by-term.
+
+        Args:
+            qubits: Ordered collection of qubits that determine the subspace
+                in which the matrix representation of the Pauli sum is to
+                be computed. If none is provided the default ordering of
+                `self.qubits` is used.  Qubits present in `qubits` but absent from
+                `self.qubits` are acted on by the identity.
+
+        Returns:
+            A `scipy.sparse.csr_matrix` representing the Pauli sum.
+        """
+        qubits = self.qubits if qubits is None else tuple(qubits)
+        num_qubits = len(qubits)
+        dim = 2**num_qubits
+        all_data: list[np.ndarray] = []
+        all_rows: list[np.ndarray] = []
+        all_cols: list[np.ndarray] = []
+
+        for vec, coeff in self._linear_dict.items():
+            op = _pauli_string_from_unit(vec)
+            term = coeff * op.sparse_matrix(qubits)
+            coo = term.tocoo()
+            all_data.append(coo.data)
+            all_rows.append(coo.row)
+            all_cols.append(coo.col)
+
+        if not all_data:
+            return sparse.csr_matrix((dim, dim), dtype=np.complex128)
+
+        data = np.concatenate(all_data)
+        rows = np.concatenate(all_rows)
+        cols = np.concatenate(all_cols)
+        result = sparse.coo_matrix((data, (rows, cols)), shape=(dim, dim)).tocsr()
+        result.eliminate_zeros()
+        return result
+
     def _has_unitary_(self) -> bool:
         return linalg.is_unitary(self.matrix())
 
@@ -927,8 +970,8 @@ def from_projector_strings(cls, terms: ProjectorString | list[ProjectorString])
     def copy(self) -> ProjectorSum:
         return ProjectorSum(self._linear_dict.copy())
 
-    def matrix(self, projector_qids: Iterable[raw_types.Qid] | None = None) -> csr_matrix:
-        """Returns the matrix of self in computational basis of qubits.
+    def matrix(self, projector_qids: Iterable[raw_types.Qid] | None = None) -> sparse.csr_matrix:
+        """Returns the matrix of self in the computational basis of the qubits.
 
         Args:
             projector_qids: Ordered collection of qubits that determine the subspace in which the

cirq-core/cirq/ops/linear_combinations_test.py

@@ -1078,6 +1078,38 @@ def test_pauli_sum_matrix() -> None:
     assert np.allclose(H3, paulisum.matrix([q[1], q[2], q[0]]))
 
 
+@pytest.mark.parametrize(
+    'paulisum, qubits',
+    (
+        # Single term.
+        (cirq.X(q0), None),
+        # Two terms, default ordering.
+        (cirq.X(q0) * cirq.X(q1) + cirq.Z(q0), None),
+        # Three terms.
+        (cirq.X(q0) + cirq.Y(q1) + cirq.Z(q2), None),
+        # Reversed qubit ordering.
+        (cirq.X(q0) * cirq.X(q1) + cirq.Z(q0), [q1, q0]),
+        # Shuffled ordering with an idle qubit.
+        (cirq.X(q0) * cirq.X(q1) + cirq.Z(q0), [q1, q2, q0]),
+        # Complex coefficients.
+        ((1 + 2j) * cirq.X(q0) * cirq.Y(q1) - 0.5 * cirq.Z(q0), None),
+        # Identity factors included.
+        (cirq.X(q0) * cirq.I(q1) + cirq.Z(q1), None),
+    ),
+)
+def test_pauli_sum_sparse_matrix(paulisum: cirq.PauliSum, qubits: list[cirq.Qid] | None) -> None:
+    actual = paulisum.sparse_matrix(qubits).toarray()
+    expected = paulisum.matrix(qubits)
+    assert np.array_equal(actual, expected)
+
+
+def test_pauli_sum_sparse_matrix_empty() -> None:
+    q = cirq.LineQubit.range(2)
+    empty = cirq.PauliSum.from_pauli_strings([])
+    assert np.array_equal(empty.sparse_matrix().toarray(), np.zeros((1, 1)))
+    assert np.array_equal(empty.sparse_matrix(q).toarray(), np.zeros((4, 4)))
+
+
 def test_pauli_sum_repr() -> None:
     q = cirq.LineQubit.range(2)
     pstr1 = cirq.X(q[0]) * cirq.X(q[1])

cirq-core/cirq/ops/pauli_string.py

@@ -33,6 +33,7 @@
 
 import numpy as np
 import sympy
+from scipy import sparse
 
 from cirq import _compat, linalg, protocols, qis, value
 from cirq._compat import deprecated
@@ -451,7 +452,7 @@ def __str__(self) -> str:
         return prefix + '*'.join(factors)
 
     def matrix(self, qubits: Iterable[TKey] | None = None) -> np.ndarray:
-        """Returns the matrix of self in computational basis of qubits.
+        """Returns the matrix of self in the computational basis of the qubits.
 
         Args:
             qubits: Ordered collection of qubits that determine the subspace
@@ -460,15 +461,72 @@ def matrix(self, qubits: Iterable[TKey] | None = None) -> np.ndarray:
                 the identity. Defaults to `self.qubits`.
 
         Raises:
-            NotImplementedError: If this PauliString is parameterized.
+            NotImplementedError: If this `PauliString` is parameterized.
         """
         qubits = self.qubits if qubits is None else qubits
         factors = [self.get(q, default=identity.I) for q in qubits]
         if protocols.is_parameterized(self):
-            raise NotImplementedError('Cannot express as matrix when parameterized')
+            raise NotImplementedError('Cannot express a parameterized PauliString as a matrix.')
         assert isinstance(self.coefficient, complex)
         return linalg.kron(self.coefficient, *[protocols.unitary(f) for f in factors])
 
+    def sparse_matrix(self, qubits: Iterable[TKey] | None = None) -> sparse.csr_matrix:
+        """Returns the sparse matrix of self in the computational basis of the qubits.
+
+        Uses a direct bit-manipulation algorithm that avoids Kronecker products
+        by computing row/col indices and phases for each basis state directly.
+
+        Args:
+            qubits: Ordered collection of qubits that determine the subspace
+                in which the matrix representation of the Pauli string is to
+                be computed. Qubits absent from `self.qubits` are acted on by
+                the identity. Defaults to `self.qubits`.
+
+        Returns:
+            A `scipy.sparse.csr_matrix` representing the Pauli string.
+
+        Raises:
+            NotImplementedError: If this `PauliString` is parameterized.
+            AssertionError: If an unexpected Pauli gate instance is encountered.
+        """
+        qubits = self.qubits if qubits is None else tuple(qubits)
+        if protocols.is_parameterized(self):
+            raise NotImplementedError('Cannot express a parameterized PauliString as a matrix.')
+        assert isinstance(self.coefficient, complex)
+
+        n = len(qubits)
+        dim = 1 << n
+        qubit_to_idx = {q: i for i, q in enumerate(qubits)}
+
+        x_mask = y_mask = z_mask = 0
+        for q, pauli in self.items():
+            if q not in qubit_to_idx:
+                continue
+            idx = qubit_to_idx[q]
+            bit = 1 << (n - 1 - idx)
+            match pauli:
+                case pauli_gates.X:
+                    x_mask |= bit
+                case pauli_gates.Y:
+                    y_mask |= bit
+                case pauli_gates.Z:
+                    z_mask |= bit
+                case _:  # pragma: no cover
+                    raise AssertionError(
+                        "Unhandled instance of Pauli gate. "
+                        "Expected one of (cirq.X, cirq.Y, cirq.Z)."
+                    )
+
+        cols = np.arange(dim, dtype=np.int32)
+        rows = cols ^ x_mask ^ y_mask
+
+        num_y = y_mask.bit_count()
+        y_phase = (1j**num_y) * np.where(np.bitwise_count(cols & y_mask) & 1, -1.0, 1.0)
+        z_phase = np.where(np.bitwise_count(cols & z_mask) & 1, -1.0, 1.0)
+        data = self.coefficient * y_phase * z_phase
+
+        return sparse.coo_matrix((data, (rows, cols)), shape=(dim, dim)).tocsr()
+
     def _has_unitary_(self) -> bool:
         if self._is_parameterized_():
             return False

cirq-core/cirq/ops/pauli_string_test.py

@@ -850,10 +850,19 @@ def _pauli_string_matrix_cases():
 
 
 @pytest.mark.parametrize('pauli_string, qubits, expected_matrix', _pauli_string_matrix_cases())
-def test_matrix(pauli_string, qubits, expected_matrix) -> None:
+def test_matrix(
+    pauli_string: cirq.PauliString, qubits: tuple[cirq.Qid, ...] | None, expected_matrix: np.ndarray
+) -> None:
     assert np.allclose(pauli_string.matrix(qubits), expected_matrix)
 
 
+@pytest.mark.parametrize('pauli_string, qubits, expected_matrix', _pauli_string_matrix_cases())
+def test_sparse_matrix(
+    pauli_string: cirq.PauliString, qubits: tuple[cirq.Qid, ...] | None, expected_matrix: np.ndarray
+) -> None:
+    assert np.allclose(pauli_string.sparse_matrix(qubits).toarray(), expected_matrix)
+
+
 def test_unitary_matrix() -> None:
     a, b = cirq.LineQubit.range(2)
     assert not cirq.has_unitary(2 * cirq.X(a) * cirq.Z(b))
@@ -2048,8 +2057,14 @@ def test_parameterization() -> None:
         pst.expectation_from_state_vector(np.array([]), {})
     with pytest.raises(NotImplementedError, match='parameterized'):
         pst.expectation_from_density_matrix(np.array([]), {})
-    with pytest.raises(NotImplementedError, match='as matrix when parameterized'):
+    with pytest.raises(
+        NotImplementedError, match='Cannot express a parameterized PauliString as a matrix'
+    ):
         pst.matrix()
+    with pytest.raises(
+        NotImplementedError, match='Cannot express a parameterized PauliString as a matrix'
+    ):
+        pst.sparse_matrix()
     assert pst**1 == pst
     assert pst**-1 == pst.with_coefficient(1.0 / t)
     assert (-pst) ** 1 == -pst