Switch position_potential_operator to Grid (#48)

* Switch position_kinetic_operator to Grid

* Switch momentum_potential_operator to Grid

* Switch jordan_wigner_position_jellium to Grid

* Switch position_potential_operator to Grid

Author: Strilanc

src/fermilib/utils/_jellium.py

@@ -13,13 +13,11 @@
 """This module constructs Hamiltonians for the uniform electron gas."""
 from __future__ import absolute_import
 
-import itertools
 import numpy
+from projectq.ops import QubitOperator
 
 from fermilib.ops import FermionOperator
 
-from projectq.ops import QubitOperator
-
 
 # Exceptions.
 class OrbitalSpecificationError(Exception):
@@ -302,22 +300,18 @@ def position_kinetic_operator(grid, spinless=False):
     return operator
 
 
-def position_potential_operator(n_dimensions, grid_length,
-                                length_scale, spinless=False):
+def position_potential_operator(grid, spinless=False):
     """Return the potential operator in position space second quantization.
 
     Args:
-        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 (Grid): The discretization to use.
         spinless: Boole, whether to use the spinless model or not.
 
     Returns:
         operator: An instance of the FermionOperator class.
     """
     # Initialize.
-    n_points = grid_length ** n_dimensions
-    volume = length_scale ** float(n_dimensions)
+    volume = grid.volume_scale()
     prefactor = 2. * numpy.pi / volume
     operator = FermionOperator()
     if spinless:
@@ -326,32 +320,29 @@ def position_potential_operator(n_dimensions, grid_length,
         spins = [0, 1]
 
     # Loop once through all lattice sites.
-    for grid_indices_a in itertools.product(range(grid_length),
-                                            repeat=n_dimensions):
+    for grid_indices_a in grid.all_points_indices():
         coordinates_a = position_vector(
-            grid_indices_a, grid_length, length_scale)
-        for grid_indices_b in itertools.product(range(grid_length),
-                                                repeat=n_dimensions):
+            grid_indices_a, grid.length, grid.scale)
+        for grid_indices_b in grid.all_points_indices():
             coordinates_b = position_vector(
-                grid_indices_b, grid_length, length_scale)
+                grid_indices_b, grid.length, grid.scale)
             differences = coordinates_b - coordinates_a
 
             # Compute coefficient.
             coefficient = 0.
-            for momenta_indices in itertools.product(range(grid_length),
-                                                     repeat=n_dimensions):
+            for momenta_indices in grid.all_points_indices():
                 momenta = momentum_vector(
-                    momenta_indices, grid_length, length_scale)
+                    momenta_indices, grid.length, grid.scale)
                 if momenta.any():
                     coefficient += (
                         prefactor * numpy.cos(momenta.dot(differences)) /
                         momenta.dot(momenta))
 
             # Loop over spins and identify interacting orbitals.
             for spin_a in spins:
-                orbital_a = orbital_id(grid_length, grid_indices_a, spin_a)
+                orbital_a = orbital_id(grid.length, grid_indices_a, spin_a)
                 for spin_b in spins:
-                    orbital_b = orbital_id(grid_length, grid_indices_b, spin_b)
+                    orbital_b = orbital_id(grid.length, grid_indices_b, spin_b)
 
                     # Add interaction term.
                     if orbital_a != orbital_b:
@@ -383,10 +374,7 @@ def jellium_model(grid, spinless=False, momentum_space=True):
         hamiltonian += momentum_potential_operator(grid, spinless)
     else:
         hamiltonian = position_kinetic_operator(grid, spinless)
-        hamiltonian += position_potential_operator(grid.dimensions,
-                                                   grid.length,
-                                                   grid.scale,
-                                                   spinless)
+        hamiltonian += position_potential_operator(grid, spinless)
     return hamiltonian
 
 

src/fermilib/utils/_jellium_test.py

@@ -151,8 +151,7 @@ def test_potential_integration(self):
         grid = Grid(dimensions=2, length=3, scale = 2.)
         spinless = 1
         momentum_potential = momentum_potential_operator(grid, spinless)
-        position_potential = position_potential_operator(
-            grid.dimensions, grid.length, grid.scale, spinless)
+        position_potential = position_potential_operator(grid, spinless)
 
         # Diagonalize and confirm the same energy.
         jw_momentum = jordan_wigner(momentum_potential)
@@ -198,8 +197,7 @@ def test_coefficients(self):
         qubit_kinetic = jordan_wigner(kinetic)
 
         # Potential operator.
-        potential = position_potential_operator(
-            grid.dimensions, grid.length, grid.scale, spinless)
+        potential = position_potential_operator(grid, spinless)
         qubit_potential = jordan_wigner(potential)
 
         # Check identity.