Improve _jellium code which computes the Hamiltonian (QubitOperator). (#83)

Author: Spaceenter

src/fermilib/utils/_jellium.py

@@ -365,81 +365,68 @@ def jordan_wigner_position_jellium(grid, spinless=False):
         n_qubits = 2 * n_orbitals
     hamiltonian = QubitOperator()
 
-    # Compute the identity coefficient.
+    # Compute the identity coefficient and the coefficient of local Z terms.
     identity_coefficient = 0.
+    z_coefficient = 0.
     for k_indices in grid.all_points_indices():
         momenta = momentum_vector(k_indices, grid)
         if momenta.any():
-            identity_coefficient += momenta.dot(momenta) / 2.
+            momenta_squared = momenta.dot(momenta)
+            identity_coefficient += momenta_squared / 2.
             identity_coefficient -= (numpy.pi * float(n_orbitals) /
-                                     (momenta.dot(momenta) * volume))
+                                     (momenta_squared * volume))
+            z_coefficient += numpy.pi / (momenta_squared * volume)
+            z_coefficient -= momenta_squared / (4. * float(n_orbitals))
     if spinless:
         identity_coefficient /= 2.
 
     # Add identity term.
     identity_term = QubitOperator((), identity_coefficient)
     hamiltonian += identity_term
 
-    # Compute coefficient of local Z terms.
-    z_coefficient = 0.
-    for k_indices in grid.all_points_indices():
-        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))
-
     # Add local Z terms.
     for qubit in range(n_qubits):
         qubit_term = QubitOperator(((qubit, 'Z'),), z_coefficient)
         hamiltonian += qubit_term
 
-    # Add ZZ terms.
-    prefactor = numpy.pi / volume
+    # Add ZZ terms and XZX + YZY terms.
+    zz_prefactor = numpy.pi / volume
+    xzx_yzy_prefactor = .25 / float(n_orbitals)
     for p in range(n_qubits):
         index_p = grid_indices(p, grid, spinless)
         position_p = position_vector(index_p, grid)
         for q in range(p + 1, n_qubits):
             index_q = grid_indices(q, grid, spinless)
             position_q = position_vector(index_q, grid)
 
-            differences = position_p - position_q
+            difference = position_p - position_q
+
+            skip_xzx_yzy = not spinless and (p + q) % 2
 
             # Loop through momenta.
             zpzq_coefficient = 0.
+            term_coefficient = 0.
             for k_indices in grid.all_points_indices():
                 momenta = momentum_vector(k_indices, grid)
                 if momenta.any():
-                    zpzq_coefficient += prefactor * numpy.cos(
-                        momenta.dot(differences)) / momenta.dot(momenta)
+                    momenta_squared = momenta.dot(momenta)
+                    cos_difference = numpy.cos(momenta.dot(difference))
+
+                    zpzq_coefficient += (zz_prefactor * cos_difference /
+                                         momenta_squared)
+
+                    if skip_xzx_yzy:
+                        continue
+                    term_coefficient += (xzx_yzy_prefactor * cos_difference *
+                                         momenta_squared)
 
-            # Add term.
+            # Add ZZ term.
             qubit_term = QubitOperator(((p, 'Z'), (q, 'Z')), zpzq_coefficient)
             hamiltonian += qubit_term
 
-    # Add XZX + YZY terms.
-    prefactor = .25 / float(n_orbitals)
-    for p in range(n_qubits):
-        index_p = grid_indices(p, grid, spinless)
-        position_p = position_vector(index_p, grid)
-        for q in range(p + 1, n_qubits):
-            if not spinless and (p + q) % 2:
+            # Add XZX + YZY term.
+            if skip_xzx_yzy:
                 continue
-
-            index_q = grid_indices(q, grid, spinless)
-            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)
-                if momenta.any():
-                    term_coefficient += (prefactor *
-                                         momenta.dot(momenta) *
-                                         numpy.cos(momenta.dot(differences)))
-
-            # Add term.
             z_string = tuple((i, 'Z') for i in range(p + 1, q))
             xzx_operators = ((p, 'X'),) + z_string + ((q, 'X'),)
             yzy_operators = ((p, 'Y'),) + z_string + ((q, 'Y'),)