Switch orbital_id to Grid (#57)

Author: Strilanc

src/fermilib/utils/_jellium.py

@@ -24,11 +24,11 @@ class OrbitalSpecificationError(Exception):
     pass
 
 
-def orbital_id(grid_length, grid_coordinates, spin=None):
+def orbital_id(grid, grid_coordinates, spin=None):
     """Return the tensor factor of a orbital with given coordinates and spin.
 
     Args:
-        grid_length: Int, the number of points in one dimension of the grid.
+        grid (Grid): The discretization to use.
         grid_coordinates: List or tuple of ints giving coordinates of grid
             element. Acceptable to provide an int (instead of tuple or list)
             for 1D case.
@@ -47,8 +47,8 @@ def orbital_id(grid_length, grid_coordinates, spin=None):
     for dimension, grid_coordinate in enumerate(grid_coordinates):
 
         # Make sure coordinate is an integer in the correct bounds.
-        if isinstance(grid_coordinate, int) and grid_coordinate < grid_length:
-            tensor_factor += grid_coordinate * (grid_length ** dimension)
+        if isinstance(grid_coordinate, int) and grid_coordinate < grid.length:
+            tensor_factor += grid_coordinate * (grid.length ** dimension)
 
         else:
             # Raise for invalid model.
@@ -164,7 +164,7 @@ def momentum_kinetic_operator(grid, spinless=False):
 
         # Loop over spins.
         for spin in spins:
-            orbital = orbital_id(grid.length, momenta_indices, spin)
+            orbital = orbital_id(grid, momenta_indices, spin)
 
             # Add interaction term.
             operators = ((orbital, 1), (orbital, 0))
@@ -220,15 +220,11 @@ def momentum_potential_operator(grid, spinless=False):
 
                 # Loop over spins.
                 for spin_a in spins:
-                    orbital_a = orbital_id(
-                        grid.length, grid_indices_a, spin_a)
-                    orbital_d = orbital_id(
-                        grid.length, shifted_indices_d, spin_a)
+                    orbital_a = orbital_id(grid, grid_indices_a, spin_a)
+                    orbital_d = orbital_id(grid, shifted_indices_d, spin_a)
                     for spin_b in spins:
-                        orbital_b = orbital_id(
-                            grid.length, grid_indices_b, spin_b)
-                        orbital_c = orbital_id(
-                            grid.length, shifted_indices_c, spin_b)
+                        orbital_b = orbital_id(grid, grid_indices_b, spin_b)
+                        orbital_c = orbital_id(grid, shifted_indices_c, spin_b)
 
                         # Add interaction term.
                         if (orbital_a != orbital_b) and \
@@ -277,8 +273,8 @@ def position_kinetic_operator(grid, spinless=False):
 
             # Loop over spins and identify interacting orbitals.
             for spin in spins:
-                orbital_a = orbital_id(grid.length, grid_indices_a, spin)
-                orbital_b = orbital_id(grid.length, grid_indices_b, spin)
+                orbital_a = orbital_id(grid, grid_indices_a, spin)
+                orbital_b = orbital_id(grid, grid_indices_b, spin)
 
                 # Add interaction term.
                 operators = ((orbital_a, 1), (orbital_b, 0))
@@ -325,9 +321,9 @@ def position_potential_operator(grid, spinless=False):
 
             # 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, 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, grid_indices_b, spin_b)
 
                     # Add interaction term.
                     if orbital_a != orbital_b:

src/fermilib/utils/_jellium_test.py

@@ -37,28 +37,25 @@ class JelliumTest(unittest.TestCase):
     def test_orbital_id(self):
 
         # Test in 1D with spin.
-        grid_length = 5
+        grid = Grid(dimensions=1, length=5, scale=1.0)
         input_coords = [0, 1, 2, 3, 4]
         tensor_factors_up = [1, 3, 5, 7, 9]
         tensor_factors_down = [0, 2, 4, 6, 8]
 
-        test_output_up = [orbital_id(
-            grid_length, i, 1) for i in input_coords]
-        test_output_down = [orbital_id(
-            grid_length, i, 0) for i in input_coords]
+        test_output_up = [orbital_id(grid, i, 1) for i in input_coords]
+        test_output_down = [orbital_id(grid, i, 0) for i in input_coords]
 
         self.assertEqual(test_output_up, tensor_factors_up)
         self.assertEqual(test_output_down, tensor_factors_down)
 
         with self.assertRaises(OrbitalSpecificationError):
-            orbital_id(5, 6, 1)
+            orbital_id(grid, 6, 1)
 
         # Test in 2D without spin.
-        grid_length = 3
+        grid = Grid(dimensions=2, length=3, scale=1.0)
         input_coords = [(0, 0), (0, 1), (1, 2)]
         tensor_factors = [0, 3, 7]
-        test_output = [orbital_id(
-            grid_length, i) for i in input_coords]
+        test_output = [orbital_id(grid, i) for i in input_coords]
         self.assertEqual(test_output, tensor_factors)
 
     def test_position_vector(self):
@@ -253,10 +250,8 @@ def test_coefficients(self):
                 for spin_a in spins:
                     for spin_b in spins:
 
-                        p = orbital_id(
-                            grid.length, indices_a, spin_a)
-                        q = orbital_id(
-                            grid.length, indices_b, spin_b)
+                        p = orbital_id(grid, indices_a, spin_a)
+                        q = orbital_id(grid, indices_b, spin_b)
 
                         if p == q:
                             continue

src/fermilib/utils/_plane_wave_hamiltonian.py

@@ -58,8 +58,7 @@ def dual_basis_u_operator(grid, geometry, spinless):
                                numpy.exp(exp_index))
 
                 for spin_p in spins:
-                    orbital_p = orbital_id(
-                        grid.length, pos_indices, spin_p)
+                    orbital_p = orbital_id(grid, pos_indices, spin_p)
                     operators = ((orbital_p, 1), (orbital_p, 0))
                     if operator is None:
                         operator = FermionOperator(operators, coefficient)
@@ -108,8 +107,8 @@ def plane_wave_u_operator(grid, geometry, spinless):
                                numpy.exp(exp_index))
 
                 for spin in spins:
-                    orbital_p = orbital_id(grid.length, indices_p, spin)
-                    orbital_q = orbital_id(grid.length, indices_q, spin)
+                    orbital_p = orbital_id(grid, indices_p, spin)
+                    orbital_q = orbital_id(grid, indices_q, spin)
                     operators = ((orbital_p, 1), (orbital_q, 0))
                     if operator is None:
                         operator = FermionOperator(operators, coefficient)
@@ -215,7 +214,7 @@ def _fourier_transform_helper(hamiltonian, grid, spinless,
                     spin = None
                 else:
                     spin = ladder_operator[0] % 2
-                orbital = orbital_id(grid.length, indices_2, spin)
+                orbital = orbital_id(grid, indices_2, spin)
                 exp_index = factor * 1.0j * numpy.dot(vec_1, vec_2)
                 if ladder_operator[1] == 1:
                     exp_index *= -1.0