Switch position_vector/momentum_vector to Grid (#50)

* 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

* Switch position_vector to Grid

* Switch momentum_vector to Grid

Author: Strilanc

src/fermilib/utils/_jellium.py

@@ -92,42 +92,37 @@ def grid_indices(qubit_id, n_dimensions, grid_length, spinless):
     return grid_indices
 
 
-def position_vector(position_indices, grid_length, length_scale):
+def position_vector(position_indices, grid):
     """Given grid point coordinate, return position vector with dimensions.
 
     Args:
         position_indices: List or tuple of integers giving grid point
             coordinate. Allowed values are ints in [0, grid_length).
-        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.
 
     Returns:
-        position_vector: A numpy array giving the position vector with
-        dimensions.
+        position_vector (numpy.ndarray[float])
     """
     # Raise exceptions.
     if isinstance(position_indices, int):
         position_indices = [position_indices]
-    if (not isinstance(grid_length, int) or
-        max(position_indices) >= grid_length or
-            min(position_indices) < 0.):
+    if not all(0 <= e < grid.length for e in position_indices):
         raise OrbitalSpecificationError(
             'Position indices must be integers in [0, grid_length).')
 
     # Compute position vector.
-    adjusted_vector = numpy.array(position_indices, float) - grid_length // 2
-    position_vector = length_scale * adjusted_vector / float(grid_length)
+    adjusted_vector = numpy.array(position_indices, float) - grid.length // 2
+    position_vector = grid.scale * adjusted_vector / float(grid.length)
     return position_vector
 
 
-def momentum_vector(momentum_indices, grid_length, length_scale):
+def momentum_vector(momentum_indices, grid):
     """Given grid point coordinate, return momentum vector with dimensions.
 
     Args:
         momentum_indices: List or tuple of integers giving momentum indices.
             Allowed values are ints in [0, grid_length).
-        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.
 
         Returns:
             momentum_vector: A numpy array giving the momentum vector with
@@ -136,15 +131,13 @@ def momentum_vector(momentum_indices, grid_length, length_scale):
     # Raise exceptions.
     if isinstance(momentum_indices, int):
         momentum_indices = [momentum_indices]
-    if (not isinstance(grid_length, int) or
-        max(momentum_indices) >= grid_length or
-            min(momentum_indices) < 0.):
+    if not all(0 <= e < grid.length for e in momentum_indices):
         raise OrbitalSpecificationError(
             'Momentum indices must be integers in [0, grid_length).')
 
     # Compute momentum vector.
-    adjusted_vector = numpy.array(momentum_indices, float) - grid_length // 2
-    momentum_vector = 2. * numpy.pi * adjusted_vector / length_scale
+    adjusted_vector = numpy.array(momentum_indices, float) - grid.length // 2
+    momentum_vector = 2. * numpy.pi * adjusted_vector / grid.scale
     return momentum_vector
 
 
@@ -167,7 +160,7 @@ def momentum_kinetic_operator(grid, spinless=False):
 
     # Loop once through all plane waves.
     for momenta_indices in grid.all_points_indices():
-        momenta = momentum_vector(momenta_indices, grid.length, grid.scale)
+        momenta = momentum_vector(momenta_indices, grid)
         coefficient = momenta.dot(momenta) / 2.
 
         # Loop over spins.
@@ -206,8 +199,7 @@ def momentum_potential_operator(grid, spinless=False):
                                  index in omega_indices]
 
         # Get the momenta vectors.
-        omega_momenta = momentum_vector(
-            omega_indices, grid.length, grid.scale)
+        omega_momenta = momentum_vector(omega_indices, grid)
 
         # Skip if omega momentum is zero.
         if not omega_momenta.any():
@@ -270,18 +262,15 @@ def position_kinetic_operator(grid, spinless=False):
 
     # Loop once through all lattice sites.
     for grid_indices_a in grid.all_points_indices():
-        coordinates_a = position_vector(
-            grid_indices_a, grid.length, grid.scale)
+        coordinates_a = position_vector(grid_indices_a, grid)
         for grid_indices_b in grid.all_points_indices():
-            coordinates_b = position_vector(
-                grid_indices_b, grid.length, grid.scale)
+            coordinates_b = position_vector(grid_indices_b, grid)
             differences = coordinates_b - coordinates_a
 
             # Compute coefficient.
             coefficient = 0.
             for momenta_indices in grid.all_points_indices():
-                momenta = momentum_vector(
-                    momenta_indices, grid.length, grid.scale)
+                momenta = momentum_vector(momenta_indices, grid)
                 if momenta.any():
                     coefficient += (
                         numpy.cos(momenta.dot(differences)) *
@@ -321,18 +310,15 @@ def position_potential_operator(grid, spinless=False):
 
     # Loop once through all lattice sites.
     for grid_indices_a in grid.all_points_indices():
-        coordinates_a = position_vector(
-            grid_indices_a, grid.length, grid.scale)
+        coordinates_a = position_vector(grid_indices_a, grid)
         for grid_indices_b in grid.all_points_indices():
-            coordinates_b = position_vector(
-                grid_indices_b, grid.length, grid.scale)
+            coordinates_b = position_vector(grid_indices_b, grid)
             differences = coordinates_b - coordinates_a
 
             # Compute coefficient.
             coefficient = 0.
             for momenta_indices in grid.all_points_indices():
-                momenta = momentum_vector(
-                    momenta_indices, grid.length, grid.scale)
+                momenta = momentum_vector(momenta_indices, grid)
                 if momenta.any():
                     coefficient += (
                         prefactor * numpy.cos(momenta.dot(differences)) /
@@ -400,7 +386,7 @@ def jordan_wigner_position_jellium(grid, spinless=False):
     # Compute the identity coefficient.
     identity_coefficient = 0.
     for k_indices in grid.all_points_indices():
-        momenta = momentum_vector(k_indices, grid.length, grid.scale)
+        momenta = momentum_vector(k_indices, grid)
         if momenta.any():
             identity_coefficient += momenta.dot(momenta) / 2.
             identity_coefficient -= (numpy.pi * float(n_orbitals) /
@@ -415,7 +401,7 @@ def jordan_wigner_position_jellium(grid, spinless=False):
     # Compute coefficient of local Z terms.
     z_coefficient = 0.
     for k_indices in grid.all_points_indices():
-        momenta = momentum_vector(k_indices, grid.length, grid.scale)
+        momenta = momentum_vector(k_indices, grid)
         if momenta.any():
             z_coefficient += numpy.pi / (momenta.dot(momenta) * volume)
             z_coefficient -= momenta.dot(momenta) / (4. * float(n_orbitals))
@@ -429,17 +415,17 @@ def jordan_wigner_position_jellium(grid, spinless=False):
     prefactor = numpy.pi / volume
     for p in range(n_qubits):
         index_p = grid_indices(p, grid.dimensions, grid.length, spinless)
-        position_p = position_vector(index_p, grid.length, grid.scale)
+        position_p = position_vector(index_p, grid)
         for q in range(p + 1, n_qubits):
             index_q = grid_indices(q, grid.dimensions, grid.length, spinless)
-            position_q = position_vector(index_q, grid.length, grid.scale)
+            position_q = position_vector(index_q, grid)
 
             differences = position_p - position_q
 
             # Loop through momenta.
             zpzq_coefficient = 0.
             for k_indices in grid.all_points_indices():
-                momenta = momentum_vector(k_indices, grid.length, grid.scale)
+                momenta = momentum_vector(k_indices, grid)
                 if momenta.any():
                     zpzq_coefficient += prefactor * numpy.cos(
                         momenta.dot(differences)) / momenta.dot(momenta)
@@ -452,20 +438,20 @@ def jordan_wigner_position_jellium(grid, spinless=False):
     prefactor = .25 / float(n_orbitals)
     for p in range(n_qubits):
         index_p = grid_indices(p, grid.dimensions, grid.length, spinless)
-        position_p = position_vector(index_p, grid.length, grid.scale)
+        position_p = position_vector(index_p, grid)
         for q in range(p + 1, n_qubits):
             if not spinless and (p + q) % 2:
                 continue
 
             index_q = grid_indices(q, grid.dimensions, grid.length, spinless)
-            position_q = position_vector(index_q, grid.length, grid.scale)
+            position_q = position_vector(index_q, grid)
 
             differences = position_p - position_q
 
             # Loop through momenta.
             term_coefficient = 0.
             for k_indices in grid.all_points_indices():
-                momenta = momentum_vector(k_indices, grid.length, grid.scale)
+                momenta = momentum_vector(k_indices, grid)
                 if momenta.any():
                     term_coefficient += prefactor * momenta.dot(momenta) * \
                         numpy.cos(momenta.dot(differences))

src/fermilib/utils/_jellium_test.py

@@ -64,62 +64,52 @@ def test_orbital_id(self):
     def test_position_vector(self):
 
         # Test in 1D.
-        grid_length = 4
-        length_scale = 4.
-        test_output = [position_vector(i, grid_length, length_scale)
-                       for i in range(grid_length)]
+        grid = Grid(dimensions=1, length=4, scale=4.)
+        test_output = [position_vector(i, grid)
+                       for i in range(grid.length)]
         correct_output = [-2, -1, 0, 1]
         self.assertEqual(correct_output, test_output)
 
-        grid_length = 11
-        length_scale = 2. * numpy.pi
-        for i in range(grid_length):
+        grid = Grid(dimensions=1, length=11, scale=2. * numpy.pi)
+        for i in range(grid.length):
             self.assertAlmostEqual(
-                -position_vector(i, grid_length, length_scale),
-                position_vector(
-                    grid_length - i - 1, grid_length, length_scale))
+                -position_vector(i, grid),
+                position_vector(grid.length - i - 1, grid))
 
         # Test in 2D.
-        grid_length = 3
-        length_scale = 3.
+        grid = Grid(dimensions=2, length=3, scale=3.)
         test_input = []
         test_output = []
         for i in range(3):
             for j in range(3):
                 test_input += [(i, j)]
-                test_output += [position_vector(
-                    (i, j), grid_length, length_scale)]
+                test_output += [position_vector((i, j), grid)]
         correct_output = numpy.array([[-1., -1.], [-1., 0.], [-1., 1.],
                                       [0., -1.], [0., 0.], [0., 1.],
                                       [1., -1.], [1., 0.], [1., 1.]])
         self.assertAlmostEqual(0., numpy.amax(test_output - correct_output))
 
     def test_momentum_vector(self):
-        grid_length = 3
-        length_scale = 2. * numpy.pi
-        test_output = [momentum_vector(i, grid_length, length_scale)
-                       for i in range(grid_length)]
+        grid = Grid(dimensions=1, length=3, scale=2. * numpy.pi)
+        test_output = [momentum_vector(i, grid)
+                       for i in range(grid.length)]
         correct_output = [-1., 0, 1.]
         self.assertEqual(correct_output, test_output)
 
-        grid_length = 11
-        length_scale = 2. * numpy.pi
-        for i in range(grid_length):
+        grid = Grid(dimensions=1, length=11, scale=2. * numpy.pi)
+        for i in range(grid.length):
             self.assertAlmostEqual(
-                -momentum_vector(i, grid_length, length_scale),
-                momentum_vector(
-                    grid_length - i - 1, grid_length, length_scale))
+                -momentum_vector(i, grid),
+                momentum_vector(grid.length - i - 1, grid))
 
         # Test in 2D.
-        grid_length = 3
-        length_scale = 2. * numpy.pi
+        grid = Grid(dimensions=2, length=3, scale=2. * numpy.pi)
         test_input = []
         test_output = []
         for i in range(3):
             for j in range(3):
                 test_input += [(i, j)]
-                test_output += [momentum_vector(
-                    (i, j), grid_length, length_scale)]
+                test_output += [momentum_vector((i, j), grid)]
         correct_output = numpy.array([[-1, -1], [-1, 0], [-1, 1],
                                       [0, -1], [0, 0], [0, 1],
                                       [1, -1], [1, 0], [1, 1]])
@@ -208,7 +198,7 @@ def test_coefficients(self):
         paper_kinetic_coefficient = 0.
         paper_potential_coefficient = 0.
         for indices in grid.all_points_indices():
-            momenta = momentum_vector(indices, grid.length, grid.scale)
+            momenta = momentum_vector(indices, grid)
             paper_kinetic_coefficient += float(
                 n_qubits) * momenta.dot(momenta) / float(4. * n_orbitals)
 
@@ -231,7 +221,7 @@ def test_coefficients(self):
             paper_kinetic_coefficient = 0.
             paper_potential_coefficient = 0.
             for indices in grid.all_points_indices():
-                momenta = momentum_vector(indices, grid.length, grid.scale)
+                momenta = momentum_vector(indices, grid)
                 paper_kinetic_coefficient -= momenta.dot(
                     momenta) / float(4. * n_orbitals)
 
@@ -257,10 +247,8 @@ def test_coefficients(self):
                 paper_kinetic_coefficient = 0.
                 paper_potential_coefficient = 0.
 
-                position_a = position_vector(
-                    indices_a, grid.length, grid.scale)
-                position_b = position_vector(
-                    indices_b, grid.length, grid.scale)
+                position_a = position_vector(indices_a, grid)
+                position_b = position_vector(indices_b, grid)
                 differences = position_b - position_a
 
                 for spin_a in spins:
@@ -281,8 +269,7 @@ def test_coefficients(self):
                             potential_coefficient = 0.
 
                         for indices_c in grid.all_points_indices():
-                            momenta = momentum_vector(
-                                indices_c, grid.length, grid.scale)
+                            momenta = momentum_vector(indices_c, grid)
 
                             if momenta.any():
                                 potential_contribution = numpy.pi * numpy.cos(

src/fermilib/utils/_plane_wave_hamiltonian.py

@@ -23,9 +23,6 @@
 from fermilib.utils._jellium import (orbital_id, grid_indices, position_vector,
                                      momentum_vector, jellium_model)
 from fermilib.utils._molecular_data import periodic_hash_table
-from fermilib.utils._grid import Grid
-
-from projectq.ops import QubitOperator
 
 
 def dual_basis_u_operator(grid, geometry, spinless):
@@ -49,12 +46,11 @@ def dual_basis_u_operator(grid, geometry, spinless):
         spins = [0, 1]
 
     for pos_indices in grid.all_points_indices():
-        coordinate_p = position_vector(pos_indices, grid.length, grid.scale)
+        coordinate_p = position_vector(pos_indices, grid)
         for nuclear_term in geometry:
             coordinate_j = numpy.array(nuclear_term[1], float)
             for momenta_indices in grid.all_points_indices():
-                momenta = momentum_vector(momenta_indices, grid.length,
-                                          grid.scale)
+                momenta = momentum_vector(momenta_indices, grid)
                 momenta_squared = momenta.dot(momenta)
                 if momenta_squared < EQ_TOLERANCE:
                     continue
@@ -100,8 +96,7 @@ def plane_wave_u_operator(grid, geometry, spinless):
             grid_indices_p_q = [
                 (indices_p[i] - indices_q[i] + shift) % grid.length
                 for i in range(grid.dimensions)]
-            momenta_p_q = momentum_vector(grid_indices_p_q, grid.length,
-                                          grid.scale)
+            momenta_p_q = momentum_vector(grid_indices_p_q, grid)
             momenta_p_q_squared = momenta_p_q.dot(momenta_p_q)
             if momenta_p_q_squared < EQ_TOLERANCE:
                 continue
@@ -218,17 +213,20 @@ def _fourier_transform_helper(hamiltonian, n_dimensions, grid_length,
                               length_scale, 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
         for ladder_operator in term:
             indices_1 = grid_indices(ladder_operator[0], n_dimensions,
                                      grid_length, spinless)
-            vec_1 = vec_func_1(indices_1, grid_length, length_scale)
+            vec_1 = vec_func_1(indices_1, grid)
             new_basis = None
             for indices_2 in itertools.product(range(grid_length),
                                                repeat=n_dimensions):
-                vec_2 = vec_func_2(indices_2, grid_length, length_scale)
+                vec_2 = vec_func_2(indices_2, grid)
                 if spinless:
                     spin = None
                 else: