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)