Threaded ensemble bugfix

pyqmc/method/ensemble_optimization_threaded.py

@@ -176,6 +176,7 @@ def evaluate_gradients_threaded(
             )
             transform_list = updater[wfi]
             for sub_iteration in range(len(transform_list)):
+                transform = transform_list[sub_iteration]
                 overlap_workers_thread = threader.submit(
                     pyqmc.method.sample_many.sample_overlap,
                     wfs[0 : wfi + 1],
@@ -319,6 +320,7 @@ def optimize_ensemble(
             npartitions=npartitions,
             vmc_kwargs=vmc_kwargs,
             overlap_kwargs=overlap_kwargs,
+            overlap_thread_weight=overlap_thread_weight,
         )
         for wfi, wf in enumerate(wfs):
             transform_list = updater[wfi]

tests/unit/test_evaluate_gradients_threaded.py

@@ -0,0 +1,48 @@
+import pyqmc.api as pyq
+import copy
+import pyqmc.observables.accumulators
+from concurrent.futures import ProcessPoolExecutor
+
+
+def test_transform_consistent_with_wf(H2_casci):
+    from pyqmc.method.ensemble_optimization_wfbywf import StochasticReconfigurationWfbyWf
+    from pyqmc.method.ensemble_optimization_threaded import evaluate_gradients_threaded
+    mol, mf, mc = H2_casci
+    mcs = [copy.copy(mc) for i in range(2)]
+    for i in range(2):
+        mcs[i].ci = mc.ci[i]
+
+    energy = pyq.EnergyAccumulator(mol)
+    sr_accumulator = []
+    tol = 1e-20
+    wfs = []
+    for i in range(2):
+
+        wf, to_opt = pyq.generate_slater(mol, mf, mc=mcs[i], optimize_determinants=True, tol = tol)
+        wfs.append(wf)
+        sr_accumulator.append(
+            [
+                StochasticReconfigurationWfbyWf(
+                    energy,
+                    pyqmc.observables.accumulators.LinearTransform(
+                        wf.parameters, to_opt
+                    ),
+                )
+            ]
+        )
+
+    configs = pyq.initial_guess(mol, 100)
+    configs_ensemble = [
+    [[copy.deepcopy(configs) for _ in range(2)] for _ in range(len(sr_accumulator[wfi]))]
+    for wfi in range(2)
+]   
+    with ProcessPoolExecutor() as executor:
+        _, data_unweighted, configs = pyqmc.method.sample_many.sample_overlap(
+                wfs,
+                configs_ensemble[0][0][0],
+                None,
+                client=executor,
+                npartitions=1
+        )
+        evaluate_gradients_threaded(wfs, configs_ensemble, sr_accumulator, client=executor)
+