Fix planewave based jellium model

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)
                     )
                 )
 

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 @@
     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 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 a for 1-D 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 @@
         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:
@@ -152,6 +183,13 @@
         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]
@@ -164,19 +202,21 @@
         if momenta_squared == 0:
             continue
 
-        # Energy cutoff.
-        if e_cutoff is not None and momenta_squared / 2.0 > e_cutoff:
-            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))
 
-        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:
@@ -228,7 +268,6 @@
     """
     # 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 +307,12 @@
             if kinetic:
                 kinetic_coefficient += cos_difference * momenta_squared / (2.0 * float(n_points))
             if potential:
-                potential_coefficient += position_prefactor * cos_difference / momenta_squared
+                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 = {}
@@ -431,13 +475,15 @@
     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 +498,6 @@
         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 +521,10 @@
 
                 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

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
@@ -395,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)
@@ -431,3 +433,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)

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)
                 )
@@ -88,6 +90,29 @@ 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."""
+    n_qubits = grid.num_points if spinless else 2 * grid.num_points
+    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
+        if energy <= e_cutoff:
+            valid_modes.add(mode)
+
+    new_terms = {}
+    for term, coeff in operator.terms.items():
+        if all(mode in valid_modes for mode, _ in term):
+            new_terms[term] = coeff
+
+    filtered_operator = FermionOperator()
+    filtered_operator.terms = new_terms
+    return filtered_operator
+
+
 def plane_wave_external_potential(
     grid: Grid,
     geometry: List[Tuple[str, Tuple[Union[int, float], Union[int, float], Union[int, float]]]],
@@ -123,6 +148,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
 
 
@@ -226,7 +254,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 +271,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)
                 )

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)