Switch fourier_transform to Grid (#52)

Author: Strilanc

examples/fermilib_demo.ipynb

@@ -625,7 +625,7 @@
     "print(momentum_qubit_operator)\n",
     "\n",
     "# Fourier transform the Hamiltonian to the position basis.\n",
-    "position_hamiltonian = fourier_transform(momentum_hamiltonian, grid.dimensions, grid.length, grid.scale, spinless)\n",
+    "position_hamiltonian = fourier_transform(momentum_hamiltonian, grid, spinless)\n",
     "position_qubit_operator = jordan_wigner(position_hamiltonian)\n",
     "position_qubit_operator.compress()\n",
     "print('')\n",

src/fermilib/utils/_plane_wave_hamiltonian.py

@@ -149,9 +149,8 @@ def plane_wave_hamiltonian(grid, geometry,
     return jellium_op + external_potential
 
 
-def fourier_transform(hamiltonian, n_dimensions, grid_length, length_scale,
-                      spinless):
-    """Apply Fourier tranform to change hamiltonian in plane wave basis.
+def fourier_transform(hamiltonian, grid, spinless):
+    """Apply Fourier transform to change hamiltonian in plane wave basis.
 
     .. math::
 
@@ -160,27 +159,22 @@ def fourier_transform(hamiltonian, n_dimensions, grid_length, length_scale,
 
     Args:
         hamiltonian: The hamiltonian in plane wave basis.
-        n_dimensions: An int giving the number of dimensions for the model.
-        grid_length: Int, the number of points in one dimension of the grid.
-        length_scale: Float, the real space length of a box dimension.
+        grid: The discretization to use.
         spinless: Bool, whether to use the spinless model or not.
 
     Returns:
         hamiltonian_t: An instance of the FermionOperator class.
     """
     return _fourier_transform_helper(hamiltonian=hamiltonian,
-                                     n_dimensions=n_dimensions,
-                                     grid_length=grid_length,
-                                     length_scale=length_scale,
+                                     grid=grid,
                                      spinless=spinless,
                                      factor=+1,
                                      vec_func_1=momentum_vector,
                                      vec_func_2=position_vector)
 
 
-def inverse_fourier_transform(hamiltonian, n_dimensions, grid_length,
-                              length_scale, spinless):
-    """Apply Fourier tranform to change hamiltonian in plane wave dual basis.
+def inverse_fourier_transform(hamiltonian, grid, spinless):
+    """Apply Fourier transform to change hamiltonian in plane wave dual basis.
 
     .. math::
 
@@ -189,31 +183,23 @@ def inverse_fourier_transform(hamiltonian, n_dimensions, grid_length,
 
     Args:
         hamiltonian: The hamiltonian in plane wave dual basis.
-        n_dimensions: An int giving the number of dimensions for the model.
-        grid_length: Int, the number of points in one dimension of the grid.
-        length_scale: Float, the real space length of a box dimension.
+        grid: The discretization to use.
         spinless: Bool, whether to use the spinless model or not.
 
     Returns:
         hamiltonian_t: An instance of the FermionOperator class.
     """
     return _fourier_transform_helper(hamiltonian=hamiltonian,
-                                     n_dimensions=n_dimensions,
-                                     grid_length=grid_length,
-                                     length_scale=length_scale,
+                                     grid=grid,
                                      spinless=spinless,
                                      factor=-1,
                                      vec_func_1=position_vector,
                                      vec_func_2=momentum_vector)
 
 
-def _fourier_transform_helper(hamiltonian, n_dimensions, grid_length,
-                              length_scale, spinless, factor,
-                              vec_func_1, vec_func_2):
+def _fourier_transform_helper(hamiltonian, grid, spinless,
+                              factor, vec_func_1, vec_func_2):
     hamiltonian_t = None
-    grid = Grid(dimensions=n_dimensions,
-                length=grid_length,
-                scale=length_scale)
 
     for term in hamiltonian.terms:
         transformed_term = None
@@ -227,7 +213,7 @@ def _fourier_transform_helper(hamiltonian, n_dimensions, grid_length,
                     spin = None
                 else:
                     spin = ladder_operator[0] % 2
-                orbital = orbital_id(grid_length, indices_2, spin)
+                orbital = orbital_id(grid.length, indices_2, spin)
                 exp_index = factor * 1.0j * numpy.dot(vec_1, vec_2)
                 if ladder_operator[1] == 1:
                     exp_index *= -1.0
@@ -239,7 +225,7 @@ def _fourier_transform_helper(hamiltonian, n_dimensions, grid_length,
                 else:
                     new_basis += element
 
-            new_basis *= numpy.sqrt(1.0/float(grid_length**n_dimensions))
+            new_basis *= numpy.sqrt(1.0/float(grid.num_points()))
 
             if transformed_term is None:
                 transformed_term = new_basis

src/fermilib/utils/_plane_wave_hamiltonian_test.py

@@ -40,9 +40,7 @@ def test_fourier_transform(self):
                 grid, geometry, spinless, True)
             h_dual_basis = plane_wave_hamiltonian(
                 grid, geometry, spinless, False)
-            h_plane_wave_t = fourier_transform(
-                h_plane_wave, grid.dimensions, grid.length,
-                grid.scale, spinless)
+            h_plane_wave_t = fourier_transform(h_plane_wave, grid, spinless)
             self.assertTrue(normal_ordered(h_plane_wave_t).isclose(
                 normal_ordered(h_dual_basis)))
 
@@ -56,8 +54,7 @@ def test_inverse_fourier_transform_1d(self):
             h_dual_basis = plane_wave_hamiltonian(
                 grid, geometry, spinless, False)
             h_dual_basis_t = inverse_fourier_transform(
-                h_dual_basis, grid.dimensions, grid.length,
-                grid.scale, spinless)
+                h_dual_basis, grid, spinless)
             self.assertTrue(normal_ordered(h_dual_basis_t).isclose(
                 normal_ordered(h_plane_wave)))
 
@@ -68,8 +65,7 @@ def test_inverse_fourier_transform_2d(self):
         h_plane_wave = plane_wave_hamiltonian(grid, geometry, spinless, True)
         h_dual_basis = plane_wave_hamiltonian(grid, geometry, spinless, False)
         h_dual_basis_t = inverse_fourier_transform(
-            h_dual_basis, grid.dimensions, grid.length, grid.scale,
-            spinless)
+            h_dual_basis, grid, spinless)
         self.assertTrue(normal_ordered(h_dual_basis_t).isclose(
             normal_ordered(h_plane_wave)))