Switch position_kinetic_operator to Grid (#44)

Author: Strilanc

src/fermilib/utils/_jellium.py

@@ -258,53 +258,47 @@ def momentum_potential_operator(n_dimensions, grid_length,
     return operator
 
 
-def position_kinetic_operator(n_dimensions, grid_length,
-                              length_scale, spinless=False):
+def position_kinetic_operator(grid, spinless=False):
     """Return the kinetic 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: Bool, whether to use the spinless model or not.
 
     Returns:
         operator: An instance of the FermionOperator class.
     """
     # Initialize.
-    n_points = grid_length ** n_dimensions
+    n_points = grid.num_points()
     operator = FermionOperator()
     if spinless:
         spins = [None]
     else:
         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 += (
                         numpy.cos(momenta.dot(differences)) *
                         momenta.dot(momenta) / (2. * float(n_points)))
 
             # 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.length, grid_indices_a, spin)
+                orbital_b = orbital_id(grid.length, grid_indices_b, spin)
 
                 # Add interaction term.
                 operators = ((orbital_a, 1), (orbital_b, 0))
@@ -397,10 +391,7 @@ def jellium_model(grid, spinless=False, momentum_space=True):
                                                    grid.scale,
                                                    spinless)
     else:
-        hamiltonian = position_kinetic_operator(grid.dimensions,
-                                                grid.length,
-                                                grid.scale,
-                                                spinless)
+        hamiltonian = position_kinetic_operator(grid, spinless)
         hamiltonian += position_potential_operator(grid.dimensions,
                                                    grid.length,
                                                    grid.scale,

src/fermilib/utils/_jellium_test.py

@@ -12,7 +12,6 @@
 
 from __future__ import absolute_import
 
-import itertools
 import unittest
 
 import numpy
@@ -129,15 +128,10 @@ def test_momentum_vector(self):
     def test_kinetic_integration(self):
 
         # Compute kinetic energy operator in both momentum and position space.
-        n_dimensions = 2
-        grid_length = 2
-        length_scale = 3.
+        grid = Grid(dimensions=2, length=2, scale=3.)
         spinless = False
-        momentum_kinetic = momentum_kinetic_operator(
-            Grid(n_dimensions, grid_length, length_scale),
-            spinless)
-        position_kinetic = position_kinetic_operator(
-            n_dimensions, grid_length, length_scale, spinless)
+        momentum_kinetic = momentum_kinetic_operator(grid, spinless)
+        position_kinetic = position_kinetic_operator(grid, spinless)
 
         # Diagonalize and confirm the same energy.
         jw_momentum = jordan_wigner(momentum_kinetic)
@@ -196,22 +190,19 @@ def test_model_integration(self):
     def test_coefficients(self):
 
         # Test that the coefficients post-JW transform are as claimed in paper.
-        n_dimensions = 2
-        grid_length = 3
-        length_scale = 2.
+        grid = Grid(dimensions=2, length=3, scale=2.)
         spinless = 1
-        n_orbitals = grid_length ** n_dimensions
+        n_orbitals = grid.num_points()
         n_qubits = (2 ** (1 - spinless)) * n_orbitals
-        volume = length_scale ** n_dimensions
+        volume = grid.volume_scale()
 
         # Kinetic operator.
-        kinetic = position_kinetic_operator(
-            n_dimensions, grid_length, length_scale, spinless)
+        kinetic = position_kinetic_operator(grid, spinless)
         qubit_kinetic = jordan_wigner(kinetic)
 
         # Potential operator.
         potential = position_potential_operator(
-            n_dimensions, grid_length, length_scale, spinless)
+            grid.dimensions, grid.length, grid.scale, spinless)
         qubit_potential = jordan_wigner(potential)
 
         # Check identity.
@@ -221,10 +212,8 @@ def test_coefficients(self):
 
         paper_kinetic_coefficient = 0.
         paper_potential_coefficient = 0.
-        for indices in itertools.product(range(grid_length),
-                                         repeat=n_dimensions):
-            momenta = momentum_vector(
-                indices, grid_length, length_scale)
+        for indices in grid.all_points_indices():
+            momenta = momentum_vector(indices, grid.length, grid.scale)
             paper_kinetic_coefficient += float(
                 n_qubits) * momenta.dot(momenta) / float(4. * n_orbitals)
 
@@ -246,10 +235,8 @@ def test_coefficients(self):
 
             paper_kinetic_coefficient = 0.
             paper_potential_coefficient = 0.
-            for indices in itertools.product(range(grid_length),
-                                             repeat=n_dimensions):
-                momenta = momentum_vector(
-                    indices, grid_length, length_scale)
+            for indices in grid.all_points_indices():
+                momenta = momentum_vector(indices, grid.length, grid.scale)
                 paper_kinetic_coefficient -= momenta.dot(
                     momenta) / float(4. * n_orbitals)
 
@@ -269,27 +256,25 @@ def test_coefficients(self):
         else:
             spins = [0, 1]
 
-        for indices_a in itertools.product(range(grid_length),
-                                           repeat=n_dimensions):
-            for indices_b in itertools.product(range(grid_length),
-                                               repeat=n_dimensions):
+        for indices_a in grid.all_points_indices():
+            for indices_b in grid.all_points_indices():
 
                 paper_kinetic_coefficient = 0.
                 paper_potential_coefficient = 0.
 
                 position_a = position_vector(
-                    indices_a, grid_length, length_scale)
+                    indices_a, grid.length, grid.scale)
                 position_b = position_vector(
-                    indices_b, grid_length, length_scale)
+                    indices_b, grid.length, grid.scale)
                 differences = position_b - position_a
 
                 for spin_a in spins:
                     for spin_b in spins:
 
                         p = orbital_id(
-                            grid_length, indices_a, spin_a)
+                            grid.length, indices_a, spin_a)
                         q = orbital_id(
-                            grid_length, indices_b, spin_b)
+                            grid.length, indices_b, spin_b)
 
                         if p == q:
                             continue
@@ -300,11 +285,9 @@ def test_coefficients(self):
                         else:
                             potential_coefficient = 0.
 
-                        for indices_c in \
-                                itertools.product(range(grid_length),
-                                                  repeat=n_dimensions):
+                        for indices_c in grid.all_points_indices():
                             momenta = momentum_vector(
-                                indices_c, grid_length, length_scale)
+                                indices_c, grid.length, grid.scale)
 
                             if momenta.any():
                                 potential_contribution = numpy.pi * numpy.cos(