From 909b71f1ac360539fddc58eb5f8c40b531ce672d Mon Sep 17 00:00:00 2001
From: arettig <arettig@google.com>
Date: Tue, 9 Jun 2026 18:39:56 +0000
Subject: [PATCH 1/5] Correct 1d and 2d planewave coulomb operator for jellium

The coulomb operator for jellium was hardcorded to the 3d case
for planewaves: \sim 1/G^2, which gives incorrect results for
3d systems. The correct coulomb operator is now selected based
on the dimensionality of the system.
---
 .../trotter/low_depth_trotter_error_test.py   | 10 ++--
 src/openfermion/hamiltonians/jellium.py       | 59 ++++++++++++++++---
 src/openfermion/hamiltonians/jellium_test.py  | 17 +++++-
 .../hamiltonians/plane_wave_hamiltonian.py    | 13 ++--
 4 files changed, 78 insertions(+), 21 deletions(-)

diff --git a/src/openfermion/circuits/trotter/low_depth_trotter_error_test.py b/src/openfermion/circuits/trotter/low_depth_trotter_error_test.py
index 40736baf2..46501c1b2 100644
--- a/src/openfermion/circuits/trotter/low_depth_trotter_error_test.py
+++ b/src/openfermion/circuits/trotter/low_depth_trotter_error_test.py
@@ -100,7 +100,7 @@ def test_error_bound_using_info_1d(self):
         terms, indices, is_hopping = result
         self.assertAlmostEqual(
             low_depth_second_order_trotter_error_bound(terms, indices, is_hopping),
-            7.4239378440953283,
+            0.00021764888459867673,
         )
 
     def test_error_bound_using_info_1d_with_input_ordering(self):
@@ -119,7 +119,7 @@ def test_error_bound_using_info_1d_with_input_ordering(self):
         terms, indices, is_hopping = result
         self.assertAlmostEqual(
             low_depth_second_order_trotter_error_bound(terms, indices, is_hopping),
-            7.4239378440953283,
+            0.00021764888459867673,
         )
 
     def test_error_bound_using_info_2d_verbose(self):
@@ -138,7 +138,7 @@ def test_error_bound_using_info_2d_verbose(self):
             low_depth_second_order_trotter_error_bound(
                 terms, indices, is_hopping, jellium_only=True, verbose=True
             ),
-            0.052213321121580794,
+            0.0003030867254527683,
         )
 
 
@@ -352,13 +352,13 @@ def test_all_terms_in_standardized_dual_basis_jellium_hamiltonian(self):
             expected_terms.append(FO(((i, 1), ((i + 3) % 4, 0)), -0.012337005501361697))
             expected_terms.append(
                 normal_ordered(
-                    FO(((i, 1), ((i + 1) % 4, 1), (i, 0), ((i + 1) % 4, 0)), 3.1830988618379052)
+                    FO(((i, 1), ((i + 1) % 4, 1), (i, 0), ((i + 1) % 4, 0)), 0.06651568175782213)
                 )
             )
             if i // 2:
                 expected_terms.append(
                     normal_ordered(
-                        FO(((i, 1), ((i + 2) % 4, 1), (i, 0), ((i + 2) % 4, 0)), 22.281692032865351)
+                        FO(((i, 1), ((i + 2) % 4, 1), (i, 0), ((i + 2) % 4, 0)), 0.1320111630094219)
                     )
                 )
 
diff --git a/src/openfermion/hamiltonians/jellium.py b/src/openfermion/hamiltonians/jellium.py
index 98b162c50..1c521c5fb 100644
--- a/src/openfermion/hamiltonians/jellium.py
+++ b/src/openfermion/hamiltonians/jellium.py
@@ -15,6 +15,7 @@
 from typing import Optional
 
 import numpy
+import scipy.special
 
 from openfermion.ops.operators import FermionOperator, QubitOperator
 from openfermion.utils.grid import Grid
@@ -61,6 +62,37 @@ def wigner_seitz_length_scale(
     return length_scale
 
 
+def coulomb_potential_momentum(
+    momenta_squared: float, dimension: int, volume: float, a_1d: float = 1.0
+) -> float:
+    """Return the momentum space Coulomb potential for a given dimension.
+
+    For 1d systems, a soft coulomb potential is used with a regularization
+    parameter of a_1d: v(r) = 1/sqrt(r^2 + a_1d^2)
+
+    Args:
+        momenta_squared (float): The squared momentum vector.
+        dimension (int): The dimension of the system (1, 2, or 3).
+        volume (float): The volume (or area, or length) of the simulation cell.
+        a_1d (float): The regularization parameter for 1D systems (default 1.0).
+
+    Returns:
+        float: The potential coefficient.
+    """
+    if momenta_squared == 0:
+        return 0.0
+    q = numpy.sqrt(momenta_squared)
+    if dimension == 3:
+        return 2.0 * numpy.pi / (volume * momenta_squared)
+    elif dimension == 2:
+        return numpy.pi / (volume * q)
+    elif dimension == 1:
+        # use a soft coulomb potential with reg. param a_1d
+        return scipy.special.k0(q * a_1d) / volume
+    else:
+        raise ValueError(f'Unsupported dimension {dimension}.')
+
+
 def plane_wave_kinetic(
     grid: Grid, spinless: bool = False, e_cutoff: Optional[float] = None
 ) -> FermionOperator:
@@ -119,7 +151,6 @@ def plane_wave_potential(
         operator (FermionOperator)
     """
     # Initialize.
-    prefactor = 2.0 * numpy.pi / grid.volume_scale()
     operator = FermionOperator((), 0.0)
     spins = [None] if spinless else [0, 1]
     if non_periodic and period_cutoff is None:
@@ -169,7 +200,9 @@ def plane_wave_potential(
             continue
 
         # Compute coefficient.
-        coefficient = prefactor / momenta_squared
+        coefficient = coulomb_potential_momentum(
+            momenta_squared, grid.dimensions, grid.volume_scale()
+        )
         if non_periodic:
             coefficient *= 1.0 - numpy.cos(period_cutoff * numpy.sqrt(momenta_squared))
 
@@ -228,7 +261,6 @@ def dual_basis_jellium_model(
     """
     # Initialize.
     n_points = grid.num_points
-    position_prefactor = 2.0 * numpy.pi / grid.volume_scale()
     operator = FermionOperator()
     spins = [None] if spinless else [0, 1]
     if potential and non_periodic and period_cutoff is None:
@@ -268,7 +300,12 @@ def dual_basis_jellium_model(
             if kinetic:
                 kinetic_coefficient += cos_difference * momenta_squared / (2.0 * float(n_points))
             if potential:
-                potential_coefficient += position_prefactor * cos_difference / momenta_squared
+                potential_coefficient += (
+                    coulomb_potential_momentum(
+                        momenta_squared, grid.dimensions, grid.volume_scale()
+                    )
+                    * cos_difference
+                )
         for grid_indices_shift in grid.all_points_indices():
             # Loop over spins and identify interacting orbitals.
             orbital_a = {}
@@ -431,13 +468,15 @@ def jordan_wigner_dual_basis_jellium(
     z_coefficient = 0.0
     for k_indices in grid.all_points_indices():
         momenta = momentum_vectors[k_indices]
-        momenta_squared = momenta.dot(momenta)
+        momenta_squared = momenta_squared_dict[k_indices]
         if momenta_squared == 0:
             continue
 
+        coulomb_potential = coulomb_potential_momentum(momenta_squared, grid.dimensions, volume)
+
         identity_coefficient += momenta_squared / 2.0
-        identity_coefficient -= numpy.pi * float(n_orbitals) / (momenta_squared * volume)
-        z_coefficient += numpy.pi / (momenta_squared * volume)
+        identity_coefficient -= coulomb_potential * float(n_orbitals) / 2.0
+        z_coefficient += coulomb_potential / 2.0
         z_coefficient -= momenta_squared / (4.0 * float(n_orbitals))
     if spinless:
         identity_coefficient /= 2.0
@@ -452,7 +491,6 @@ def jordan_wigner_dual_basis_jellium(
         hamiltonian += qubit_term
 
     # Add ZZ terms and XZX + YZY terms.
-    zz_prefactor = numpy.pi / volume
     xzx_yzy_prefactor = 0.25 / float(n_orbitals)
     for p in range(n_qubits):
         index_p = grid.grid_indices(p, spinless)
@@ -476,7 +514,10 @@ def jordan_wigner_dual_basis_jellium(
 
                 cos_difference = numpy.cos(momenta.dot(difference))
 
-                zpzq_coefficient += zz_prefactor * cos_difference / momenta_squared
+                coulomb_potential = coulomb_potential_momentum(
+                    momenta_squared, grid.dimensions, volume
+                )
+                zpzq_coefficient += coulomb_potential / 2.0 * cos_difference
 
                 if skip_xzx_yzy:
                     continue
diff --git a/src/openfermion/hamiltonians/jellium_test.py b/src/openfermion/hamiltonians/jellium_test.py
index 25f5f7bf8..dcf14a6f6 100644
--- a/src/openfermion/hamiltonians/jellium_test.py
+++ b/src/openfermion/hamiltonians/jellium_test.py
@@ -15,6 +15,7 @@
 import numpy
 
 from openfermion.hamiltonians.jellium import (
+    coulomb_potential_momentum,
     dual_basis_jellium_model,
     dual_basis_kinetic,
     dual_basis_potential,
@@ -256,7 +257,7 @@ def test_model_integration_with_constant(self):
 
     def test_coefficients(self):
         # Test that the coefficients post-JW transform are as claimed in paper.
-        grid = Grid(dimensions=2, length=3, scale=2.0)
+        grid = Grid(dimensions=3, length=2, scale=2.0)
         spinless = 1
         n_orbitals = grid.num_points
         n_qubits = (2 ** (1 - spinless)) * n_orbitals
@@ -431,3 +432,17 @@ def test_plane_wave_period_cutoff(self):
         #     integration test.
         jellium_model(grid, spinless, True, False, None, True)
         jellium_model(grid, spinless, False, False, None, True)
+
+
+class CoulombPotentialMomentumTest(unittest.TestCase):
+    def test_zero_momentum(self):
+        # return 0.0 when momenta_squared == 0
+        val = coulomb_potential_momentum(0.0, dimension=3, volume=10.0)
+        self.assertEqual(val, 0.0)
+
+    def test_unsupported_dimension(self):
+        # crash when number dimensions not in [1,3]
+        with self.assertRaises(ValueError):
+            _ = coulomb_potential_momentum(1.0, dimension=0, volume=10.0)
+        with self.assertRaises(ValueError):
+            _ = coulomb_potential_momentum(1.0, dimension=4, volume=10.0)
diff --git a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
index f37053ec8..408f1d737 100644
--- a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
+++ b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
@@ -16,6 +16,7 @@
 import numpy as np
 
 from openfermion.hamiltonians.jellium import jellium_model, jordan_wigner_dual_basis_jellium
+from openfermion.hamiltonians.jellium import coulomb_potential_momentum
 from openfermion.ops.operators import FermionOperator, QubitOperator
 from openfermion.transforms.repconversions import inverse_fourier_transform
 from openfermion.utils.grid import Grid
@@ -52,7 +53,6 @@ def dual_basis_external_potential(
     Returns:
         FermionOperator: The dual basis operator.
     """
-    prefactor = -4.0 * np.pi / grid.volume_scale()
     if non_periodic and period_cutoff is None:
         period_cutoff = grid.volume_scale() ** (1.0 / grid.dimensions)
     operator = None
@@ -72,8 +72,10 @@ def dual_basis_external_potential(
 
                 cos_index = momenta.dot(coordinate_j - coordinate_p)
                 coefficient = (
-                    prefactor
-                    / momenta_squared
+                    -2.0
+                    * coulomb_potential_momentum(
+                        momenta_squared, grid.dimensions, grid.volume_scale()
+                    )
                     * md.periodic_hash_table[nuclear_term[0]]
                     * np.cos(cos_index)
                 )
@@ -226,7 +228,6 @@ def jordan_wigner_dual_basis_hamiltonian(
         n_qubits = n_orbitals
     else:
         n_qubits = 2 * n_orbitals
-    prefactor = -2 * np.pi / volume
     external_potential = QubitOperator()
 
     for k_indices in grid.all_points_indices():
@@ -244,8 +245,8 @@ def jordan_wigner_dual_basis_hamiltonian(
 
                 cos_index = momenta.dot(coordinate_j - coordinate_p)
                 coefficient = (
-                    prefactor
-                    / momenta_squared
+                    -1.0
+                    * coulomb_potential_momentum(momenta_squared, grid.dimensions, volume)
                     * md.periodic_hash_table[nuclear_term[0]]
                     * np.cos(cos_index)
                 )

From 7ec6ad621a098c56d5e26d4bd3af7af02d7520a4 Mon Sep 17 00:00:00 2001
From: arettig <arettig@google.com>
Date: Tue, 9 Jun 2026 18:44:42 +0000
Subject: [PATCH 2/5] Make truncation of planewave basis set consistent

The planewave basis set when used with an Ecut was truncated for
the kinetic operator in the Hamiltonian, while the nuclear-electronic
term was not truncated, and the electron-electron repulsion term
was truncated incorrectly. Vee and Ven have been changed to match
the kinetic planewave truncation.
---
 src/openfermion/hamiltonians/jellium.py       | 29 ++++++++++++-------
 src/openfermion/hamiltonians/jellium_test.py  |  5 ++--
 .../hamiltonians/plane_wave_hamiltonian.py    | 22 ++++++++++++++
 .../plane_wave_hamiltonian_test.py            | 10 ++++---
 4 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/src/openfermion/hamiltonians/jellium.py b/src/openfermion/hamiltonians/jellium.py
index 1c521c5fb..0f885dbaa 100644
--- a/src/openfermion/hamiltonians/jellium.py
+++ b/src/openfermion/hamiltonians/jellium.py
@@ -183,6 +183,13 @@ def plane_wave_potential(
         for spin in spins:
             orbital_ids[indices_a][spin] = grid.orbital_id(indices_a, spin)
 
+    # Identify active grid indices within the energy cutoff.
+    active_indices = set()
+    for indices in grid.all_points_indices():
+        momenta = grid.momentum_vector(indices)
+        if e_cutoff is None or momenta.dot(momenta) / 2.0 <= e_cutoff:
+            active_indices.add(indices)
+
     # Loop once through all plane waves.
     for omega_indices in grid.all_points_indices():
         shifted_omega_indices = shifted_omega_indices_dict[omega_indices]
@@ -195,10 +202,6 @@ def plane_wave_potential(
         if momenta_squared == 0:
             continue
 
-        # Energy cutoff.
-        if e_cutoff is not None and momenta_squared / 2.0 > e_cutoff:
-            continue
-
         # Compute coefficient.
         coefficient = coulomb_potential_momentum(
             momenta_squared, grid.dimensions, grid.volume_scale()
@@ -206,10 +209,14 @@ def plane_wave_potential(
         if non_periodic:
             coefficient *= 1.0 - numpy.cos(period_cutoff * numpy.sqrt(momenta_squared))
 
-        for grid_indices_a in grid.all_points_indices():
+        for grid_indices_a in active_indices:
             shifted_indices_d = shifted_indices_minus_dict[omega_indices][grid_indices_a]
-            for grid_indices_b in grid.all_points_indices():
+            if shifted_indices_d not in active_indices:
+                continue
+            for grid_indices_b in active_indices:
                 shifted_indices_c = shifted_indices_plus_dict[omega_indices][grid_indices_b]
+                if shifted_indices_c not in active_indices:
+                    continue
 
                 # Loop over spins.
                 for spin_a in spins:
@@ -300,12 +307,12 @@ def dual_basis_jellium_model(
             if kinetic:
                 kinetic_coefficient += cos_difference * momenta_squared / (2.0 * float(n_points))
             if potential:
-                potential_coefficient += (
-                    coulomb_potential_momentum(
-                        momenta_squared, grid.dimensions, grid.volume_scale()
-                    )
-                    * cos_difference
+                coef = coulomb_potential_momentum(
+                    momenta_squared, grid.dimensions, grid.volume_scale()
                 )
+                if non_periodic:
+                    coef *= 1.0 - numpy.cos(period_cutoff * numpy.sqrt(momenta_squared))
+                potential_coefficient += coef * cos_difference
         for grid_indices_shift in grid.all_points_indices():
             # Loop over spins and identify interacting orbitals.
             orbital_a = {}
diff --git a/src/openfermion/hamiltonians/jellium_test.py b/src/openfermion/hamiltonians/jellium_test.py
index dcf14a6f6..257bdc879 100644
--- a/src/openfermion/hamiltonians/jellium_test.py
+++ b/src/openfermion/hamiltonians/jellium_test.py
@@ -396,14 +396,15 @@ def test_plane_wave_energy_cutoff(self):
         grid = Grid(dimensions=1, length=5, scale=1.0)
         spinless = True
         e_cutoff = 20.0
+        n_qubits = grid.num_points if spinless else 2 * grid.num_points
 
         hamiltonian_1 = jellium_model(grid, spinless, True, False)
         jw_1 = jordan_wigner(hamiltonian_1)
-        spectrum_1 = eigenspectrum(jw_1)
+        spectrum_1 = eigenspectrum(jw_1, n_qubits=n_qubits)
 
         hamiltonian_2 = jellium_model(grid, spinless, True, False, e_cutoff)
         jw_2 = jordan_wigner(hamiltonian_2)
-        spectrum_2 = eigenspectrum(jw_2)
+        spectrum_2 = eigenspectrum(jw_2, n_qubits=n_qubits)
 
         max_diff = numpy.amax(numpy.absolute(spectrum_1 - spectrum_2))
         self.assertGreater(max_diff, 0.0)
diff --git a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
index 408f1d737..ee2b80e4a 100644
--- a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
+++ b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
@@ -90,6 +90,25 @@ def dual_basis_external_potential(
     return operator
 
 
+def filter_plane_wave_operator(
+    operator: FermionOperator, grid: Grid, e_cutoff: float, spinless: bool
+) -> FermionOperator:
+    """Filter plane wave operator to only include terms within energy cutoff."""
+    new_operator = FermionOperator()
+    for term, coeff in operator.terms.items():
+        keep = True
+        for mode, _ in term:
+            indices = grid.grid_indices(mode, spinless)
+            momenta = grid.momentum_vector(indices)
+            energy = momenta.dot(momenta) / 2.0
+            if energy > e_cutoff:
+                keep = False
+                break
+        if keep:
+            new_operator += FermionOperator(term, coeff)
+    return new_operator
+
+
 def plane_wave_external_potential(
     grid: Grid,
     geometry: List[Tuple[str, Tuple[Union[int, float], Union[int, float], Union[int, float]]]],
@@ -125,6 +144,9 @@ def plane_wave_external_potential(
     )
     operator = inverse_fourier_transform(dual_basis_operator, grid, spinless)
 
+    if e_cutoff is not None:
+        operator = filter_plane_wave_operator(operator, grid, e_cutoff, spinless)
+
     return operator
 
 
diff --git a/src/openfermion/hamiltonians/plane_wave_hamiltonian_test.py b/src/openfermion/hamiltonians/plane_wave_hamiltonian_test.py
index c6267c812..1e43d99dc 100644
--- a/src/openfermion/hamiltonians/plane_wave_hamiltonian_test.py
+++ b/src/openfermion/hamiltonians/plane_wave_hamiltonian_test.py
@@ -115,14 +115,16 @@ def test_plane_wave_energy_cutoff(self):
         geometry = [('H', (0,)), ('H', (0.8,))]
         grid = Grid(dimensions=1, scale=1.1, length=5)
         e_cutoff = 50.0
+        spinless = True
+        n_qubits = grid.num_points if spinless else 2 * grid.num_points
 
-        h_1 = plane_wave_hamiltonian(grid, geometry, True, True, False)
+        h_1 = plane_wave_hamiltonian(grid, geometry, spinless, True, False)
         jw_1 = jordan_wigner(h_1)
-        spectrum_1 = eigenspectrum(jw_1)
+        spectrum_1 = eigenspectrum(jw_1, n_qubits=n_qubits)
 
-        h_2 = plane_wave_hamiltonian(grid, geometry, True, True, False, e_cutoff)
+        h_2 = plane_wave_hamiltonian(grid, geometry, spinless, True, False, e_cutoff)
         jw_2 = jordan_wigner(h_2)
-        spectrum_2 = eigenspectrum(jw_2)
+        spectrum_2 = eigenspectrum(jw_2, n_qubits=n_qubits)
 
         max_diff = np.amax(np.absolute(spectrum_1 - spectrum_2))
         self.assertGreater(max_diff, 0.0)

From a8690c7ab0df3b07deb0d945d6ebf5bba70fcdc2 Mon Sep 17 00:00:00 2001
From: Adam Rettig <arettig@google.com>
Date: Tue, 9 Jun 2026 15:04:56 -0700
Subject: [PATCH 3/5] Speed up operator planewave truncation

The reduces the computational cost of filtering out high energy planewaves from an operator.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../hamiltonians/plane_wave_hamiltonian.py    | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
index ee2b80e4a..0eff20e2b 100644
--- a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
+++ b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
@@ -94,19 +94,19 @@ def filter_plane_wave_operator(
     operator: FermionOperator, grid: Grid, e_cutoff: float, spinless: bool
 ) -> FermionOperator:
     """Filter plane wave operator to only include terms within energy cutoff."""
-    new_operator = FermionOperator()
+    n_qubits = grid.num_points if spinless else 2 * grid.num_points
+    valid_modes = {}
+    for mode in range(n_qubits):
+        indices = grid.grid_indices(mode, spinless)
+        momenta = grid.momentum_vector(indices)
+        energy = momenta.dot(momenta) / 2.0
+        valid_modes[mode] = energy <= e_cutoff
+
+    new_terms = {}
     for term, coeff in operator.terms.items():
-        keep = True
-        for mode, _ in term:
-            indices = grid.grid_indices(mode, spinless)
-            momenta = grid.momentum_vector(indices)
-            energy = momenta.dot(momenta) / 2.0
-            if energy > e_cutoff:
-                keep = False
-                break
-        if keep:
-            new_operator += FermionOperator(term, coeff)
-    return new_operator
+        if all(valid_modes.get(mode, False) for mode, _ in term):
+            new_terms[term] = coeff
+    return FermionOperator(new_terms)
 
 
 def plane_wave_external_potential(

From f750e6b979724476c3017b6d70c0bc5dcec48514 Mon Sep 17 00:00:00 2001
From: arettig <arettig@google.com>
Date: Tue, 9 Jun 2026 23:26:02 +0000
Subject: [PATCH 4/5] Fix erroneous constructor used in
 filter_plane_wave_operator()

A nonexistent constructor was used when truncating the planewave
basis of a FermionOperator.
---
 .../hamiltonians/plane_wave_hamiltonian.py           | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
index 0eff20e2b..652f1b333 100644
--- a/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
+++ b/src/openfermion/hamiltonians/plane_wave_hamiltonian.py
@@ -95,18 +95,22 @@ def filter_plane_wave_operator(
 ) -> FermionOperator:
     """Filter plane wave operator to only include terms within energy cutoff."""
     n_qubits = grid.num_points if spinless else 2 * grid.num_points
-    valid_modes = {}
+    valid_modes = set()
     for mode in range(n_qubits):
         indices = grid.grid_indices(mode, spinless)
         momenta = grid.momentum_vector(indices)
         energy = momenta.dot(momenta) / 2.0
-        valid_modes[mode] = energy <= e_cutoff
+        if energy <= e_cutoff:
+            valid_modes.add(mode)
 
     new_terms = {}
     for term, coeff in operator.terms.items():
-        if all(valid_modes.get(mode, False) for mode, _ in term):
+        if all(mode in valid_modes for mode, _ in term):
             new_terms[term] = coeff
-    return FermionOperator(new_terms)
+
+    filtered_operator = FermionOperator()
+    filtered_operator.terms = new_terms
+    return filtered_operator
 
 
 def plane_wave_external_potential(

From dc9d31726c84984254bf88a9863973abad091de5 Mon Sep 17 00:00:00 2001
From: arettig <arettig@google.com>
Date: Wed, 10 Jun 2026 15:25:17 +0000
Subject: [PATCH 5/5] Add latex formatting to jellium.py docstring

Formatted equation in jellium.py docstring to use latex for
easier viewing. Several nits were fixed as well.
---
 src/openfermion/hamiltonians/jellium.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/openfermion/hamiltonians/jellium.py b/src/openfermion/hamiltonians/jellium.py
index 0f885dbaa..d39d595c3 100644
--- a/src/openfermion/hamiltonians/jellium.py
+++ b/src/openfermion/hamiltonians/jellium.py
@@ -67,14 +67,14 @@ def coulomb_potential_momentum(
 ) -> float:
     """Return the momentum space Coulomb potential for a given dimension.
 
-    For 1d systems, a soft coulomb potential is used with a regularization
-    parameter of a_1d: v(r) = 1/sqrt(r^2 + a_1d^2)
+    For 1-D systems, a soft coulomb potential is used with a regularization
+    parameter of a: $$v(r) = \frac{1}{\sqrt{r^2 + a^2}}$$
 
     Args:
         momenta_squared (float): The squared momentum vector.
         dimension (int): The dimension of the system (1, 2, or 3).
         volume (float): The volume (or area, or length) of the simulation cell.
-        a_1d (float): The regularization parameter for 1D systems (default 1.0).
+        a_1d (float): The regularization parameter a for 1-D systems (default 1.0).
 
     Returns:
         float: The potential coefficient.