From 9df8156e3ddd04e4cc2bb7f4a8534ef9bdab3bdd Mon Sep 17 00:00:00 2001 From: Lucas Wagner Date: Wed, 3 Dec 2025 09:24:09 -0600 Subject: [PATCH 1/4] sample_many now actually writes to the HDF file every block. --- pyqmc/method/hdftools.py | 7 ++++ pyqmc/method/sample_many.py | 84 +++++++++++++++++++++++++------------ pyqmc/observables/obdm.py | 1 - 3 files changed, 64 insertions(+), 28 deletions(-) diff --git a/pyqmc/method/hdftools.py b/pyqmc/method/hdftools.py index 9222ae0b..88d0941e 100644 --- a/pyqmc/method/hdftools.py +++ b/pyqmc/method/hdftools.py @@ -35,6 +35,13 @@ def setup_hdf(f, data, attr): def append_hdf(f, data): + """ + Add data from one block to the HDF file. + + f should be an h5py file object + data should be a dictionary of numpy arrays that represent one block. + + """ for k, it in data.items(): if k not in f.keys(): itnp = np.array(it) diff --git a/pyqmc/method/sample_many.py b/pyqmc/method/sample_many.py index e6f7d05c..561a20f3 100644 --- a/pyqmc/method/sample_many.py +++ b/pyqmc/method/sample_many.py @@ -15,20 +15,16 @@ import numpy as np import pyqmc.method.mc as mc import scipy.stats -import pyqmc.method.linemin as linemin -import pyqmc.gpu as gpu -import os +#import pyqmc.method.linemin as linemin +#import pyqmc.gpu as gpu +#import os import h5py +import pyqmc.method.hdftools as hdftools -def run_vmc_many(wfs, configs, energy, nreps=1, **kwargs): - # data gets saved at the end of each rep - for i in range(nreps): - sample_overlap(wfs, configs, energy, **kwargs) def hdf_save(hdf_file, weighted, unweighted, attr, configs): - import pyqmc.method.hdftools as hdftools if hdf_file is not None: fulldata = dict(weighted=weighted, unweighted=unweighted) @@ -39,21 +35,10 @@ def hdf_save(hdf_file, weighted, unweighted, attr, configs): for label, data in fulldata.items(): hdf.create_group(label) for label, data in fulldata.items(): - extend_hdf(hdf[label], data) + hdftools.append_hdf(hdf[label], data) configs.to_hdf(hdf) -def extend_hdf(f, data): - for k, it in data.items(): - if k not in f.keys(): - f.create_dataset( - k, (0, *it.shape[1:]), maxshape=(None, *it.shape[1:]), dtype=it.dtype - ) - n = it.shape[0] - f[k].resize((f[k].shape[0] + n), axis=0) - f[k][-n:] = it - - def compute_weights(wfs): """ computes psi_i* psi_j / rho for all i,j and for each configuration. @@ -82,18 +67,23 @@ def invert_list_of_dicts(A, asarray=True): return {k: [a[k] for a in A] for k in A[0].keys()} -def sample_overlap_worker(wfs, configs, tstep, nsteps, nblocks, energy): - r"""Run nstep Metropolis steps to sample a distribution proportional to - :math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i] +def sample_overlap_run(wfs, configs, tstep, nsteps, nblocks, energy, + hdf_file=None, client=None, npartitions=None): + """ + Use a single core to sample over blocks """ nconf, nelec, _ = configs.configs.shape weighted = [] unweighted = [] for block in range(nblocks): print("-", end="", flush=True) - w, u, _ = sample_overlap_block(wfs, configs, tstep, nsteps, energy) + if client is None: + w, u, configs = sample_overlap_worker(wfs, configs, tstep, nsteps, energy) + else: + w, u, configs = sample_overlap_client(wfs, configs, tstep, nsteps, energy, client, npartitions) weighted.append(w) unweighted.append(u) + hdf_save(hdf_file, w, u, dict(tstep=tstep), configs) # here we modify the data so that weighted and unweighted are dictionaries of arrays # Access as weighted[quantity][block, ...] @@ -101,8 +91,46 @@ def sample_overlap_worker(wfs, configs, tstep, nsteps, nblocks, energy): unweighted = invert_list_of_dicts(unweighted) return weighted, unweighted, configs +def sample_overlap_client(wfs, configs, tstep, nsteps, energy, client, npartitions): + """ + Sample nblocks, saving every block. + wfs: list of wave functions + configs: pyqmc.config.Config + tstep: float + nsteps: int + energy: Accumulator object + client: futures client + npartitions: number of jobs to submit to the client. + """ + config = configs.split(npartitions) + runs = [ + client.submit(sample_overlap_worker, wfs, conf, tstep, nsteps, energy) + for conf in config + ] + allresults = list(zip(*[r.result() for r in runs])) #weighted, unweighted, configs + configs.join(allresults[2]) + confweight = np.array([len(c.configs) for c in config], dtype=float) + confweight /= np.mean(confweight) * npartitions + weighted_block = {} + for k in allresults[0][0].keys(): + weighted_block[k] = np.sum( + [res[k] * w for res, w in zip(allresults[0], confweight)], axis=0 + ) + unweighted_block = {} + for k in allresults[1][0].keys(): + unweighted_block[k] = np.sum( + [res[k] * w for res, w in zip(allresults[1], confweight)], axis=0 + ) + + + return weighted_block, unweighted_block, configs + + -def sample_overlap_block(wfs, configs, tstep, nsteps, energy): +def sample_overlap_worker(wfs, configs, tstep, nsteps, energy): + r"""Run nstep Metropolis steps to sample a distribution proportional to + :math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i] + """ for wf in wfs: wf.recompute(configs) weighted_block = {} @@ -176,9 +204,11 @@ def sample_overlap( npartitions=None, ): """ """ + + return sample_overlap_run(wfs, configs, tstep, nsteps, nblocks, energy, hdf_file, client, npartitions) + if client is None: # otherwise running in parallel - w, u, _ = sample_overlap_worker(wfs, configs, tstep, nsteps, nblocks, energy) - hdf_save(hdf_file, w, u, dict(tstep=tstep), configs) + w, u, _ = sample_overlap_single(wfs, configs, tstep, nsteps, nblocks, energy, hdf_file) return w, u, configs if npartitions is None: diff --git a/pyqmc/observables/obdm.py b/pyqmc/observables/obdm.py index ab520dcf..883dfa22 100644 --- a/pyqmc/observables/obdm.py +++ b/pyqmc/observables/obdm.py @@ -203,7 +203,6 @@ def shapes(self): "norm": (norb,), } - def sample_onebody(configs, orbitals, spin, nsamples=1, tstep=0.5): r""" For a set of orbitals defined by orb_coeff, return samples from :math:`f(r) = \sum_i \phi_i(r)^2`. From a223e4d57b778aa5f484d19846c19e47fdc09ddc Mon Sep 17 00:00:00 2001 From: Lucas Wagner Date: Fri, 5 Dec 2025 11:54:33 -0600 Subject: [PATCH 2/4] restart --- pyqmc/method/sample_many.py | 33 ++++++--------------------------- 1 file changed, 6 insertions(+), 27 deletions(-) diff --git a/pyqmc/method/sample_many.py b/pyqmc/method/sample_many.py index 561a20f3..f81aecc2 100644 --- a/pyqmc/method/sample_many.py +++ b/pyqmc/method/sample_many.py @@ -17,7 +17,7 @@ import scipy.stats #import pyqmc.method.linemin as linemin #import pyqmc.gpu as gpu -#import os +import os import h5py import pyqmc.method.hdftools as hdftools @@ -204,35 +204,14 @@ def sample_overlap( npartitions=None, ): """ """ + if os.path.isfile(hdf_file): + with h5py.File(hdf_file, "r") as f: + with h5py.File(continue_from, "r") as hdf: + if "configs" in hdf.keys(): + configs.load_hdf(hdf) return sample_overlap_run(wfs, configs, tstep, nsteps, nblocks, energy, hdf_file, client, npartitions) - if client is None: # otherwise running in parallel - w, u, _ = sample_overlap_single(wfs, configs, tstep, nsteps, nblocks, energy, hdf_file) - return w, u, configs - - if npartitions is None: - npartitions = sum(client.nthreads().values()) - - coord = configs.split(npartitions) - runs = [ - client.submit(sample_overlap_worker, wfs, conf, tstep, nsteps, nblocks, energy) - for conf in coord - ] - allresults = list(zip(*[r.result() for r in runs])) - configs.join(allresults[2]) - confweight = np.array([len(c.configs) for c in coord], dtype=float) - weighted = {} - for k, it in invert_list_of_dicts(allresults[0]).items(): - weighted[k] = np.average(it, weights=confweight, axis=0) - unweighted = {} - for k, it in invert_list_of_dicts(allresults[1]).items(): - unweighted[k] = np.average(it, weights=confweight, axis=0) - - hdf_save(hdf_file, weighted, unweighted, dict(tstep=tstep), configs) - return weighted, unweighted, configs - - def normalize(weighted, unweighted): """ (more or less) correctly average the output from sample_overlap From 0eae646109318d1bbac06d245bd338f0c73dfea0 Mon Sep 17 00:00:00 2001 From: Lucas Wagner Date: Fri, 5 Dec 2025 11:55:14 -0600 Subject: [PATCH 3/4] Fixed restart. --- pyqmc/method/sample_many.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyqmc/method/sample_many.py b/pyqmc/method/sample_many.py index f81aecc2..15c25f5a 100644 --- a/pyqmc/method/sample_many.py +++ b/pyqmc/method/sample_many.py @@ -206,7 +206,8 @@ def sample_overlap( """ """ if os.path.isfile(hdf_file): with h5py.File(hdf_file, "r") as f: - with h5py.File(continue_from, "r") as hdf: + with h5py.File(hdf_file, "r") as hdf: + print("Restarting") if "configs" in hdf.keys(): configs.load_hdf(hdf) From f6e53e1790539bb11c55eb59b7e996f3c0ed466f Mon Sep 17 00:00:00 2001 From: Lucas Wagner Date: Fri, 5 Dec 2025 14:21:59 -0600 Subject: [PATCH 4/4] Fixed another bug; check for None. --- pyqmc/method/sample_many.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyqmc/method/sample_many.py b/pyqmc/method/sample_many.py index 15c25f5a..13c6d428 100644 --- a/pyqmc/method/sample_many.py +++ b/pyqmc/method/sample_many.py @@ -204,7 +204,7 @@ def sample_overlap( npartitions=None, ): """ """ - if os.path.isfile(hdf_file): + if hdf_file is not None and os.path.isfile(hdf_file): with h5py.File(hdf_file, "r") as f: with h5py.File(hdf_file, "r") as hdf: print("Restarting")