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,47 @@ 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 computational basis of qubits.
+
+        Uses direct bit-manipulation for each term and accumulates COO triplets
+        to avoid repeated sparse matrix addition overhead.
+
+        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,7 +967,7 @@ 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:
+    def matrix(self, projector_qids: Iterable[raw_types.Qid] | None = None) -> sparse.csr_matrix:
         """Returns the matrix of self in computational basis of qubits.
 
         Args:

cirq-core/cirq/ops/linear_combinations_test.py

@@ -1078,6 +1078,40 @@ def test_pauli_sum_matrix() -> None:
     assert np.allclose(H3, paulisum.matrix([q[1], q[2], q[0]]))
 
 
+def test_pauli_sum_sparse_matrix() -> None:
+    q = cirq.LineQubit.range(3)
+    paulisum = cirq.X(q[0]) * cirq.X(q[1]) + cirq.Z(q[0])
+    H1 = np.array(
+        [[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0], [0.0, 1.0, -1.0, 0.0], [1.0, 0.0, 0.0, -1.0]]
+    )
+    # Default qubit ordering.
+    assert np.allclose(H1, paulisum.sparse_matrix().toarray())
+    # Explicit 2-qubit ordering (same as default).
+    assert np.allclose(H1, paulisum.sparse_matrix([q[0], q[1]]).toarray())
+    # Reversed qubit ordering changes the matrix.
+    H2 = np.array(
+        [[1.0, 0.0, 0.0, 1.0], [0.0, -1.0, 1.0, 0.0], [0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, -1.0]]
+    )
+    assert np.allclose(H2, paulisum.sparse_matrix([q[1], q[0]]).toarray())
+    # Adding an extra idle qubit expands to 8x8.
+    H3 = np.array(
+        [
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+            [0.0, -1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            [0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0],
+        ]
+    )
+    assert np.allclose(H3, paulisum.sparse_matrix([q[1], q[2], q[0]]).toarray())
+    # Empty PauliSum should return a zero sparse matrix.
+    empty = cirq.PauliSum.from_pauli_strings([])
+    assert np.allclose(empty.sparse_matrix([q[0], q[1]]).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
@@ -469,6 +470,57 @@ def matrix(self, qubits: Iterable[TKey] | None = None) -> np.ndarray:
         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 computational basis of 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.
+        """
+        qubits = self.qubits if qubits is None else tuple(qubits)
+        if protocols.is_parameterized(self):
+            raise NotImplementedError('Cannot express as matrix when parameterized')
+        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 in qubits:
+            pauli = self.get(q)
+            if pauli is None:
+                continue
+            idx = qubit_to_idx[q]
+            bit = 1 << (n - 1 - idx)
+            if pauli is pauli_gates.X:
+                x_mask |= bit
+            elif pauli is pauli_gates.Y:
+                y_mask |= bit
+            elif pauli is pauli_gates.Z:
+                z_mask |= bit
+
+        cols = np.arange(dim, dtype=np.int32)
+        rows = cols ^ x_mask ^ y_mask
+
+        num_y = y_mask.bit_count()
+        y_phase = (1j**num_y) * ((-1.0) ** np.bitwise_count(cols & y_mask))
+        z_phase = (-1.0) ** np.bitwise_count(cols & z_mask)
+        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

@@ -854,6 +854,11 @@ def test_matrix(pauli_string, qubits, expected_matrix) -> 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, qubits, expected_matrix) -> 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))
@@ -2050,6 +2055,8 @@ def test_parameterization() -> None:
         pst.expectation_from_density_matrix(np.array([]), {})
     with pytest.raises(NotImplementedError, match='as matrix when parameterized'):
         pst.matrix()
+    with pytest.raises(NotImplementedError, match='as matrix when parameterized'):
+        pst.sparse_matrix()
     assert pst**1 == pst
     assert pst**-1 == pst.with_coefficient(1.0 / t)
     assert (-pst) ** 1 == -pst