Quality of life improvements

examples/trial_wf/slater_geminal_etb_optimize.ipynb

@@ -388,15 +388,16 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2025-12-08T15:03:17.285861Z",
-     "start_time": "2025-12-08T15:03:17.077966Z"
+     "end_time": "2025-12-08T15:10:10.463409Z",
+     "start_time": "2025-12-08T15:10:10.198661Z"
     }
    },
    "cell_type": "code",
    "source": [
     "g = sns.relplot(data =df, x=\"n_s\", y=\"energy\", hue=\"n_p\", col=\"n_d\")\n",
     "for ax in g.axes.flat:\n",
-    "    ax.axhline(ccvtz_energy, linestyle=\"--\", color=\"k\", label=\"ccvtz energy\")\n"
+    "    ax.axhline(ccvtz_energy, linestyle=\"--\", color=\"k\", label=\"ccvtz energy\")\n",
+    "plt.savefig(\"slater_geminal_etb_optimize.png\")\n"
    ],
    "id": "5d8cd4fbb3aae616",
    "outputs": [
@@ -414,7 +415,7 @@
      }
     }
    ],
-   "execution_count": 8
+   "execution_count": 9
   },
   {
    "metadata": {

examples/workflow/he_recipe.py

@@ -11,11 +11,11 @@
 #
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
+import pyscf
+import pyqmc
+import pyqmc.recipes
 
 if __name__ == "__main__":
-    import pyscf
-    import pyqmc
-    import pyqmc.recipes
 
     mol = pyscf.gto.M(
         atom="He 0. 0. 0.", basis="ccECP_cc-pVDZ", ecp="ccecp", unit="bohr"
@@ -28,7 +28,7 @@
     slater_kws = {"optimize_orbitals": True}
 
     pyqmc.recipes.OPTIMIZE(
-        "he_dft.hdf5", "he_sj.hdf5", jastrow_kws=jastrow_kws, slater_kws=slater_kws
+        "he_dft.hdf5", "he_sj.hdf5", jastrow_kws=jastrow_kws, slater_kws=slater_kws, verbose=True
     )
 
     pyqmc.recipes.VMC(
@@ -48,5 +48,6 @@
         accumulators={"rdm1": True},
         jastrow_kws=jastrow_kws,
         slater_kws=slater_kws,
+        verbose = True,
         **{"nblocks": 4000, "tstep": 0.02},
     )

pyqmc/method/dmc.py

@@ -304,9 +304,10 @@ def dmc_propagate_parallel(wf, configs, weights, client, npartitions, *args, **k
     return block_avg, configs, weights
 
 
-def branch(configs, weights, var_trigger_branch=0.05):
+def branch_trigger_variance(configs, weights, var_trigger_branch=0.05):
     """
     Perform branching on a set of walkers using the 'stochastic comb'
+    This is triggered
 
     Walkers are resampled with probability proportional to the weights, and the new weights are all set to be equal to the average weight.
 
@@ -317,14 +318,38 @@ def branch(configs, weights, var_trigger_branch=0.05):
 
     """
     if np.var(weights) < var_trigger_branch:
-        return (
-            configs,
-            weights,
-            {
-                "max branches": 0,
-                "Number of walkers killed": 0,
-            },
-        )
+        return branch(configs, weights)
+
+
+def branch_trigger_maxweight(configs, weights, trigger_size=2.0):
+    """
+    Perform branching on a set of walkers using the 'stochastic comb'
+    This is triggered
+
+    Walkers are resampled with probability proportional to the weights, and the new weights are all set to be equal to the average weight.
+
+    :parameter configs: (nconfig,nelec,3) walker coordinates
+    :parameter weights: (nconfig,) walker weights
+    :parameter var_trigger_branch: variance at which to perform branching.
+    :returns: resampled walker configurations and weights all equal to average weight
+
+    """
+    if np.max(weights) < trigger_size:
+        return branch(configs, weights)
+
+
+def branch(configs, weights):
+    """
+    Perform branching on a set of walkers using the 'stochastic comb'
+
+    Walkers are resampled with probability proportional to the weights, and the new weights are all set to be equal to the average weight.
+
+    :parameter configs: (nconfig,nelec,3) walker coordinates
+    :parameter weights: (nconfig,) walker weights
+    :parameter var_trigger_branch: variance at which to perform branching.
+    :returns: resampled walker configurations and weights all equal to average weight
+
+    """
 
     nconfig = configs.configs.shape[0]
     if np.any(weights > 2.0):
@@ -405,6 +430,7 @@ def rundmc(
     branchtime=None,
     stepoffset=None,
     nsteps=None,
+    branch_trigger = 2.0
 ):
     """
     Run DMC
@@ -428,6 +454,7 @@ def rundmc(
     :parameter int vmc_warmup: If starting a run, how many VMC warmup blocks to run
     :parameter int branchcut_start: Used in computing weights. Recommended for "experts only".
     :parameter float feedback: Feedback strength for controlling normalization. Recommended for "experts only".
+    :parameter float branch_trigger: At what point do we trigger a branching event.
     :returns: (df,coords,weights)
       df: A list of dictionaries nblocks long that contains all results from the accumulators.
 
@@ -554,7 +581,7 @@ def rundmc(
         df_["weight_std"] = np.std(weights)
         df_["nsteps_per_block"] = nsteps_per_block
 
-        configs, weights, branch_info = branch(configs, weights)
+        configs, weights, branch_info = branch_trigger_maxweight(configs, weights, branch_trigger)
         df_.update(branch_info)
         df.append(df_)
         dmc_file(hdf_file, df_, {}, configs, weights)

pyqmc/method/linemin.py

@@ -253,7 +253,7 @@ def line_minimization(
 
             # Correlated sampling line minimization.
             steps = np.linspace(-steprange / (npts - 2), steprange, npts)
-            dps, update_report = pgrad.delta_p(steps, data, verbose=verbose)
+            dps, update_report = pgrad.delta_p(steps, data, verbose=False)
             step_data.update(update_report)
             params = [x0 + dp for dp in dps]
 
@@ -283,11 +283,11 @@ def line_minimization(
             step_data["correlated_energy_std"] = en_std
 
             if verbose:
-                print("energies from correlated sampling", en)
+                #print("energies from correlated sampling", en)
                 print("Chose to move", est_min, "from", steps[0], "to", steps[-1])
-                print("weight variance", np.var(w, axis=1))
-                print("avg weights", np.mean(correlated_data["weight"], axis=1))
-                print("energy standard deviation", en_std)
+                #print("weight variance", np.var(w, axis=1))
+                #print("avg weights", np.mean(correlated_data["weight"], axis=1))
+                #print("energy standard deviation", en_std)
 
             set_wf_params(wf, x0, pgrad)
             opt_hdf(hdf_file, step_data, attr, coords, wf.parameters)

pyqmc/observables/stochastic_reconfiguration.py

@@ -59,6 +59,7 @@ def __init__(
         nodal_cutoff=1e-3,
         eps=1e-1,
         inverse_strategy="pseudo_inverse",
+        verbose = False
     ):
         """
         eps here is the regularization for SR.
@@ -68,6 +69,7 @@ def __init__(
         self.nodal_cutoff = nodal_cutoff
         self.eps = eps
         self.inverse_strategy = inverse_strategy
+        self.verbose = verbose
 
     def __call__(self, configs, wf):
         pgrad = wf.pgradient()
@@ -164,7 +166,7 @@ def delta_p(self, steps, data: dict, verbose=False):
             "SRdot": np.dot(pgrad, v) / (np.linalg.norm(v) * np.linalg.norm(pgrad)),
         }
 
-        if verbose:
+        if verbose or self.verbose:
             eigvals = np.linalg.eigvals(Sij)
             print("eigenvalues of Sij", eigvals)
             print("Gradient norm: ", np.linalg.norm(pgrad))