From 0f297fb37457c44f7430cc04f96e6e2d936b9338 Mon Sep 17 00:00:00 2001
From: Fionn Malone <fionn.malone@gmail.com>
Date: Wed, 26 Jun 2024 08:03:22 -0700
Subject: [PATCH] Fix test failure.

---
 recirq/qcqmc/optimize_wf.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py
index 704a0001..49e1c4ae 100644
--- a/recirq/qcqmc/optimize_wf.py
+++ b/recirq/qcqmc/optimize_wf.py
@@ -27,8 +27,16 @@
 import scipy.optimize
 import scipy.sparse
 
-from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data,
-                          fermion_mode, hamiltonian, layer_spec, trial_wf)
+from recirq.qcqmc import (
+    afqmc_circuits,
+    afqmc_generators,
+    converters,
+    data,
+    fermion_mode,
+    hamiltonian,
+    layer_spec,
+    trial_wf,
+)
 
 
 def get_and_check_energy(
@@ -499,6 +507,7 @@ def evaluate_energy_and_gradient(
 
     return cost_val, np.concatenate((two_body_grad, one_body_grad))
 
+
 def compute_finite_difference_grad(
     n_orb: int,
     n_elec: int,
@@ -520,7 +529,7 @@ def compute_finite_difference_grad(
         initial_wf: The initial wavefunction (typically Hartree--Fock)
         restricted: Whether we're using a restricted ansatz or not.
     """
-    generators = get_pp_plus_gate_generators(
+    generators = afqmc_generators.get_pp_plus_gate_generators(
         n_elec=n_elec, heuristic_layers=tuple(), do_pp=True
     )
     one_body_gradient = np.zeros_like(one_body_params)
@@ -574,7 +583,6 @@ def compute_finite_difference_grad(
     return one_body_gradient, two_body_gradient
 
 
-
 def objective(
     params: np.ndarray,
     n_one_body_params: int,