modify low_rank for one-body-matrix truncation reference

src/openfermion/circuits/low_rank.py

@@ -159,22 +159,22 @@ def low_rank_two_body_decomposition(two_body_coefficients,
             one_body_correction, truncation_value)
 
 
-def prepare_one_body_squared_evolution(one_body_matrix, spin_basis=True):
+def prepare_one_body_squared_evolution(one_body_matrix, spin_basis=True, truncation_ref=None):
     r"""Get Givens angles and DiagonalHamiltonian to simulate squared one-body.
 
     The goal here will be to prepare to simulate evolution under
     $(\sum_{pq} h_{pq} a^\dagger_p a_q)^2$ by decomposing as
     $R e^{-i \sum_{pq} V_{pq} n_p n_q} R^\dagger$ where
     $R$ is a basis transformation matrix.
 
-    TODO: Add option for truncation based on one-body eigenvalues.
-
     Args:
         one_body_matrix (ndarray of floats): an N by N array storing the
             coefficients of a one-body operator to be squared. For instance,
             in the above the elements of this matrix are $h_{pq}$.
         spin_basis (bool): Whether the matrix is passed in the
             spin orbital basis.
+        truncation_ref (float): Eigenvalues under this threshold reference will
+            be truncated for density matrix calculation.
 
     Returns:
         density_density_matrix(ndarray of floats) an N by N array storing
@@ -206,6 +206,18 @@ def prepare_one_body_squared_evolution(one_body_matrix, spin_basis=True):
                                                  numpy.eye(2))
         eigenvalues = numpy.kron(eigenvalues, numpy.ones(2))
 
+    # If the truncation reference is valid, 
+    # build truncated eigenvalues and eigenvectors
+    # based on the reference for density matrix
+    if truncation_ref is not None and isinstance(truncation_ref, float):
+        valid_indices = numpy.where(eigenvalues >= truncation_ref)[0]
+        retained_indices = numpy.where(eigenvalues < truncation_ref)[0]
+        eigenvalues = eigenvalues[valid_indices]
+        eigenvectors = eigenvectors[:, valid_indices]
+        print(f"Retained {len(valid_indices)} eigenvalues and truncated {len(retained_indices)} eigenvalues from the one-body-matrix.")
+    else:
+        raise ValueError("truncation_ref argument must be a valid float.")
+
     # Obtain the diagonal two-body matrix.
     density_density_matrix = numpy.outer(eigenvalues, eigenvalues)
 

src/openfermion/circuits/low_rank_test.py

@@ -328,3 +328,11 @@ def test_one_body_squared_nonhermitian_raises_error(self):
         with self.assertRaises(ValueError):
             prepare_one_body_squared_evolution(one_body_matrix,
                                                spin_basis=False)
+
+    def test_one_body_squared_truncated(self):
+        one_body_matrix = numpy.array([[2.0, 1.5, 0.2],[1.5, 3.0, 1.0],[0.2, 1.0, 4.0]])
+        truncation_reference = 1.0
+        with self.assertRaises(ValueError):
+            prepare_one_body_squared_evolution(one_body_matrix, 
+                                               spin_basis=False, 
+                                               truncation_ref=truncation_reference)