From f8ff9aa935d34859714d7ae61bd25a6eab7b6a6d Mon Sep 17 00:00:00 2001
From: marko1olo <barsukdana@gmail.com>
Date: Mon, 8 Jun 2026 16:32:41 +0400
Subject: [PATCH] Handle unsimplified qubit sparse terms

---
 src/openfermion/linalg/sparse_tools.py      | 14 +++++--
 src/openfermion/linalg/sparse_tools_test.py | 41 ++++++++++++++++++++-
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/openfermion/linalg/sparse_tools.py b/src/openfermion/linalg/sparse_tools.py
index a8baf51ea..7b5121cb9 100644
--- a/src/openfermion/linalg/sparse_tools.py
+++ b/src/openfermion/linalg/sparse_tools.py
@@ -155,9 +155,8 @@ def qubit_operator_sparse(qubit_operator, n_qubits=None):
     column_list = [[]]
 
     # Loop through the terms.
-    for qubit_term in qubit_operator.terms:
+    for qubit_term, coefficient in _iter_simplified_qubit_operator_terms(qubit_operator):
         tensor_factor = 0
-        coefficient = qubit_operator.terms[qubit_term]
         sparse_operators = [coefficient]
         for pauli_operator in qubit_term:
             # Grow space for missing identity operators.
@@ -194,6 +193,13 @@ def qubit_operator_sparse(qubit_operator, n_qubits=None):
     return sparse_operator
 
 
+def _iter_simplified_qubit_operator_terms(qubit_operator):
+    """Yield terms after applying QubitOperator's Pauli simplification."""
+    for qubit_term, coefficient in qubit_operator.terms.items():
+        coefficient, qubit_term = qubit_operator._simplify(qubit_term, coefficient)
+        yield qubit_term, coefficient
+
+
 def get_linear_qubit_operator_diagonal(qubit_operator, n_qubits=None):
     """Return a linear operator's diagonal elements.
 
@@ -221,7 +227,7 @@ def get_linear_qubit_operator_diagonal(qubit_operator, n_qubits=None):
     ones_diagonal = numpy.ones(n_hilbert)
     linear_operator_diagonal = zeros_diagonal
     # Loop through the terms.
-    for qubit_term in qubit_operator.terms:
+    for qubit_term, coefficient in _iter_simplified_qubit_operator_terms(qubit_operator):
         is_zero = False
         tensor_factor = 0
         vecs = [ones_diagonal]
@@ -243,7 +249,7 @@ def get_linear_qubit_operator_diagonal(qubit_operator, n_qubits=None):
             vecs = [v for vp in vec_pairs for v in (vp[0], -vp[1])]
             tensor_factor = pauli_operator[0] + 1
         if not is_zero:
-            linear_operator_diagonal += qubit_operator.terms[qubit_term] * numpy.concatenate(vecs)
+            linear_operator_diagonal += coefficient * numpy.concatenate(vecs)
     return linear_operator_diagonal
 
 
diff --git a/src/openfermion/linalg/sparse_tools_test.py b/src/openfermion/linalg/sparse_tools_test.py
index b889f5859..c312a62de 100644
--- a/src/openfermion/linalg/sparse_tools_test.py
+++ b/src/openfermion/linalg/sparse_tools_test.py
@@ -41,7 +41,12 @@
     MajoranaOperator,
 )
 from openfermion.ops.representations import DiagonalCoulombHamiltonian
-from openfermion.transforms.opconversions import jordan_wigner, get_fermion_operator, normal_ordered
+from openfermion.transforms.opconversions import (
+    jordan_wigner,
+    get_fermion_operator,
+    normal_ordered,
+    symmetry_conserving_bravyi_kitaev,
+)
 from openfermion.transforms.repconversions import fourier_transform, get_interaction_operator
 
 from openfermion.testing.testing_utils import random_hermitian_matrix
@@ -137,6 +142,40 @@ def test_qubit_jw_fermion_integration(self):
         fermion_spectrum = sparse_eigenspectrum(fermion_sparse)
         self.assertAlmostEqual(0.0, numpy.amax(numpy.absolute(fermion_spectrum - qubit_spectrum)))
 
+    def test_qubit_operator_sparse_simplifies_repeated_qubit_factors(self):
+        qubit_operator = QubitOperator()
+        qubit_operator.terms[((0, 'X'), (1, 'Y'), (1, 'X'))] = 1.0
+        expected_operator = -1j * QubitOperator('X0 Z1')
+
+        sparse_operator = qubit_operator_sparse(qubit_operator, n_qubits=2)
+        expected_sparse_operator = qubit_operator_sparse(expected_operator, n_qubits=2)
+
+        self.assertTrue(
+            numpy.allclose(sparse_operator.toarray(), expected_sparse_operator.toarray())
+        )
+
+    def test_get_sparse_operator_with_unsimplified_symmetry_conserving_bk(self):
+        fermion_operator = FermionOperator('0^ 1^')
+        qubit_operator = symmetry_conserving_bravyi_kitaev(fermion_operator, 4, 2)
+        expected_operator = QubitOperator()
+        for qubit_term, coefficient in qubit_operator.terms.items():
+            expected_operator += QubitOperator(qubit_term, coefficient)
+
+        sparse_operator = get_sparse_operator(qubit_operator)
+        expected_sparse_operator = get_sparse_operator(expected_operator)
+
+        self.assertTrue(
+            numpy.allclose(sparse_operator.toarray(), expected_sparse_operator.toarray())
+        )
+
+    def test_get_linear_qubit_operator_diagonal_simplifies_repeated_qubit_factors(self):
+        qubit_operator = QubitOperator()
+        qubit_operator.terms[((0, 'Z'), (1, 'Z'), (1, 'Z'))] = 2.0
+
+        diagonal = get_linear_qubit_operator_diagonal(qubit_operator, n_qubits=2)
+
+        self.assertTrue(numpy.allclose(diagonal, 2.0 * numpy.array([1, 1, -1, -1])))
+
 
 class JordanWignerSparseTest(unittest.TestCase):
     def test_jw_sparse_0create(self):