diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ab4578de..fb8c232e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,11 +13,11 @@ jobs: - name: Setup MPI uses: mpi4py/setup-mpi@v1 with: - mpi: openmpi + mpi: openmpi - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -r requirements.txt + pip install -r requirements.txt pip install -r dev/dev.txt - name: pylint run: | @@ -32,7 +32,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: [ '3.8', '3.9', '3.10' ] + python-version: [ '3.10', '3.11', '3.12' ] mpi: - openmpi runs-on: ubuntu-latest @@ -53,7 +53,7 @@ jobs: pip install -r dev/dev.txt pip install pyblock # https://github.com/JoonhoLee-Group/ipie/issues/278 - pip install "pyscf<=2.3.0" + pip install pyscf - name: Install package run: | # HACK FOR LEGACY CODE! @@ -108,11 +108,11 @@ jobs: - name: Setup MPI uses: mpi4py/setup-mpi@v1 with: - mpi: openmpi + mpi: openmpi - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -r requirements.txt + pip install -r requirements.txt pip install -r dev/dev.txt - name: Install package run: | @@ -131,14 +131,14 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -r requirements.txt + pip install -r requirements.txt pip install pyblock # https://github.com/JoonhoLee-Group/ipie/issues/278 - pip install "pyscf<=2.3.0" fqe + pip install pyscf fqe - name: Install package run: | python -m pip install -e . - name: Test Examples timeout-minutes: 10 run: | - python dev/run_tests.py --examples \ No newline at end of file + python dev/run_tests.py --examples diff --git a/.gitignore b/.gitignore index 19ea23c8..611a5f27 100644 --- a/.gitignore +++ b/.gitignore @@ -108,4 +108,6 @@ ipie/qmc/tests/reference_data/**/*.h5 *.txt *wheels* *FCIDUMP* -*out* \ No newline at end of file +*out* + +*.code-workspace diff --git a/.pylintrc b/.pylintrc index 9ba5f042..59d6545e 100644 --- a/.pylintrc +++ b/.pylintrc @@ -23,8 +23,8 @@ output-format=colorized disable=all enable= no-member, - # renable once classes are tidied -; attribute-defined-outside-init, + # renable once classes are tidied + # attribute-defined-outside-init, consider-using-f-string, useless-object-inheritance, unused-variable, @@ -68,4 +68,4 @@ enable= no-else-return, # https://github.com/pylint-dev/pylint/issues/2178 # ignore urls -ignore-long-lines=^\s*(# )??$|^\s*(\w*\s*=\s*)?(\"|\').*(\"|\'),?\s*$ \ No newline at end of file +ignore-long-lines=^\s*(# )??$|^\s*(\w*\s*=\s*)?(\"|\').*(\"|\'),?\s*$ diff --git a/dev/run_tests.py b/dev/run_tests.py old mode 100644 new mode 100755 index 710fb22a..03dd61ff --- a/dev/run_tests.py +++ b/dev/run_tests.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 import argparse import glob import os diff --git a/examples/03-custom_observable/run_afqmc.py b/examples/03-custom_observable/run_afqmc.py index a4100b36..e97e455a 100644 --- a/examples/03-custom_observable/run_afqmc.py +++ b/examples/03-custom_observable/run_afqmc.py @@ -15,6 +15,8 @@ # Author: Fionn Malone # +from typing import Dict + import numpy as np from pyscf import gto, scf @@ -86,11 +88,11 @@ def __init__(self, ham): # Must specify that we're dealing with array valued estimator self.scalar_estimator = False - def compute_estimator(self, system, walker_batch, hamiltonian, trial_wavefunction): - trial_wavefunction.calc_greens_function(walker_batch, build_full=True) - numer = np.einsum("w,wii->i", walker_batch.weight, walker_batch.Ga + walker_batch.Gb) + def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None): + trial.calc_greens_function(walkers, build_full=True) + numer = np.einsum("w,wii->i", walkers.weight, walkers.Ga + walkers.Gb) self["DiagGNumer"] = numer - self["DiagGDenom"] = sum(walker_batch.weight) + self["DiagGDenom"] = sum(walkers.weight) afqmc = build_afqmc_driver(comm, nelec=mol.nelec) @@ -127,25 +129,28 @@ def __init__(self, ham): # Must specify that we're dealing with array valued estimator self.scalar_estimator = False - def compute_estimator(self, system, walker_batch, hamiltonian, trial_wavefunction): - trial_wavefunction.calc_greens_function(walker_batch, build_full=True) + def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None): + trial.calc_greens_function(walkers, build_full=True) numer = np.array( [ - np.einsum("w,wij->ij", walker_batch.weight, walker_batch.Ga), - np.einsum("w,wij->ij", walker_batch.weight, walker_batch.Gb), + np.einsum("w,wij->ij", walkers.weight, walkers.Ga), + np.einsum("w,wij->ij", walkers.weight, walkers.Gb), ] ) # For multidimensional arrays we must flatten the data self["GNumer"] = numer.ravel() - self["GDenom"] = sum(walker_batch.weight) + self["GDenom"] = sum(walkers.weight) afqmc = build_afqmc_driver(comm, nelec=mol.nelec) # Let us override the number of blocks to keep it short afqmc.params.num_blocks = 20 # We can now add this to the estimator handler object in the afqmc driver -add_est = {"diagG": Diagonal1RDM(ham=afqmc.hamiltonian), "1RDM": Mixed1RDM(ham=afqmc.hamiltonian)} +add_est: Dict[str, EstimatorBase] = { + "diagG": Diagonal1RDM(ham=afqmc.hamiltonian), + "1RDM": Mixed1RDM(ham=afqmc.hamiltonian), +} afqmc.run(additional_estimators=add_est) # We can extract the qmc data as as a pandas data frame like so from ipie.analysis.extraction import extract_observable diff --git a/examples/07-custom_trial/run_afqmc.py b/examples/07-custom_trial/run_afqmc.py index 58b34375..48379be0 100644 --- a/examples/07-custom_trial/run_afqmc.py +++ b/examples/07-custom_trial/run_afqmc.py @@ -91,7 +91,7 @@ def __init__( trial=trial, ) - def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): + def compute_estimator(self, system, walkers, hamiltonian, trial): trial.calc_greens_function(walkers) # Need to be able to dispatch here energy = local_energy_batch(system, hamiltonian, walkers, trial) diff --git a/examples/08-custom_walker/run_afqmc.py b/examples/08-custom_walker/run_afqmc.py index 4125c5a0..4dd545da 100644 --- a/examples/08-custom_walker/run_afqmc.py +++ b/examples/08-custom_walker/run_afqmc.py @@ -94,7 +94,7 @@ def __init__( trial=trial, ) - def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): + def compute_estimator(self, system, walkers, hamiltonian, trial): trial.calc_greens_function(walkers) # Need to be able to dispatch here energy = local_energy_batch(system, hamiltonian, walkers, trial) diff --git a/examples/16-ft_afqmc/run_afqmc.py b/examples/16-ft_afqmc/run_afqmc.py index 272da7e4..77783a78 100644 --- a/examples/16-ft_afqmc/run_afqmc.py +++ b/examples/16-ft_afqmc/run_afqmc.py @@ -1,11 +1,12 @@ import json -import numpy +import numpy from ueg import UEG -from ipie.config import MPI + from ipie.addons.thermal.qmc.calc import build_thermal_afqmc_driver -from ipie.analysis.extraction import extract_observable from ipie.analysis.autocorr import reblock_by_autocorr +from ipie.analysis.extraction import extract_observable +from ipie.config import MPI comm = MPI.COMM_WORLD @@ -65,7 +66,7 @@ # 3. Run thermal AFQMC calculation. afqmc.run(verbose=verbose) afqmc.finalise() -afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) +afqmc.estimators.compute_estimators(hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers) if comm.rank == 0: energy_data = extract_observable(afqmc.estimators.filename, "energy") diff --git a/ipie/addons/free_projection/estimators/energy.py b/ipie/addons/free_projection/estimators/energy.py index bd434374..2425dfa6 100644 --- a/ipie/addons/free_projection/estimators/energy.py +++ b/ipie/addons/free_projection/estimators/energy.py @@ -30,7 +30,7 @@ def __init__( ): super().__init__(system, ham, trial, filename) - def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): + def compute_estimator(self, system, walkers, hamiltonian, trial): trial.calc_greens_function(walkers) # Need to be able to dispatch here energy = local_energy(system, hamiltonian, walkers, trial) diff --git a/ipie/addons/free_projection/estimators/tests/test_estimators.py b/ipie/addons/free_projection/estimators/tests/test_estimators.py index ec130a0e..7bf692ab 100644 --- a/ipie/addons/free_projection/estimators/tests/test_estimators.py +++ b/ipie/addons/free_projection/estimators/tests/test_estimators.py @@ -68,8 +68,12 @@ def test_estimator_handler_fp(): handler["energy1"] = estim handler.json_string = "" handler.initialize(comm) - handler.compute_estimators(comm, system, ham, trial, walker_batch) - handler.compute_estimators(comm, system, ham, trial, walker_batch) + handler.compute_estimators( + system=system, hamiltonian=ham, trial=trial, walker_batch=walker_batch + ) + handler.compute_estimators( + system=system, hamiltonian=ham, trial=trial, walker_batch=walker_batch + ) if __name__ == "__main__": diff --git a/ipie/addons/free_projection/qmc/fp_afqmc.py b/ipie/addons/free_projection/qmc/fp_afqmc.py index 6b794052..f0d971ca 100644 --- a/ipie/addons/free_projection/qmc/fp_afqmc.py +++ b/ipie/addons/free_projection/qmc/fp_afqmc.py @@ -62,7 +62,7 @@ def build( trial_wavefunction, walkers=None, num_walkers: int = 100, - seed: int = None, + seed: Optional[int] = None, num_steps_per_block: int = 25, num_blocks: int = 100, timestep: float = 0.005, @@ -282,7 +282,7 @@ def setup_estimators( def run( self, - psi=None, + walkers=None, estimator_filename="estimate.h5", verbose=True, additional_estimators: Optional[Dict[str, EstimatorBase]] = None, @@ -300,8 +300,8 @@ def run( """ self.setup_timers() tzero_setup = time.time() - if psi is not None: - self.walkers = psi + if walkers is not None: + self.walkers = walkers self.setup_timers() eshift = 0.0 self.walkers.orthogonalise() @@ -360,7 +360,10 @@ def run( start = time.time() if step % self.params.num_steps_per_block == 0: self.estimators[block_number].compute_estimators( - comm, self.system, self.hamiltonian, self.trial, self.walkers + system=self.system, + hamiltonian=self.hamiltonian, + trial=self.trial, + walker_batch=self.walkers, ) self.estimators[block_number].print_block( comm, diff --git a/ipie/addons/thermal/analysis/extraction.py b/ipie/addons/thermal/analysis/extraction.py index 913cafb5..5f725f21 100644 --- a/ipie/addons/thermal/analysis/extraction.py +++ b/ipie/addons/thermal/analysis/extraction.py @@ -19,6 +19,7 @@ from ipie.utils.misc import get_from_dict + def set_info(frame, md): ncols = len(frame.columns) system = md.get("system") @@ -61,4 +62,3 @@ def set_info(frame, md): frame["cholesky_treshold"] = chol return list(frame.columns[ncols:]) - diff --git a/ipie/addons/thermal/analysis/thermal_analysis.py b/ipie/addons/thermal/analysis/thermal_analysis.py index b9cba603..7fefbeba 100644 --- a/ipie/addons/thermal/analysis/thermal_analysis.py +++ b/ipie/addons/thermal/analysis/thermal_analysis.py @@ -18,21 +18,16 @@ #!/usr/bin/env python -import sys import argparse - import glob +import sys + import numpy -import scipy.optimize import pandas as pd - -from ipie.analysis.extraction import ( - extract_observable, - get_metadata, - get_sys_param - ) +import scipy.optimize from ipie.addons.thermal.analysis.extraction import set_info +from ipie.analysis.extraction import extract_observable, get_metadata def parse_args(args): @@ -49,18 +44,10 @@ def parse_args(args): Command line arguments. """ - parser = argparse.ArgumentParser(description = __doc__) - parser.add_argument('-c', '--chem-pot', dest='fit_chem_pot', - action='store_true', default=False, - help='Estimate optimal chemical potential') - parser.add_argument('-n', '--nav', dest='nav', type=float, - help='Target electron density.') - parser.add_argument('-o', '--order', dest='order', type=int, - default=3, help='Order polynomial to fit.') - parser.add_argument('-p', '--plot', dest='plot', action='store_true', - help='Plot density vs. mu.') - parser.add_argument('-f', nargs='+', dest='filenames', - help='Space-separated list of files to analyse.') + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "-f", nargs="+", dest="filenames", help="Space-separated list of files to analyse." + ) options = parser.parse_args(args) @@ -76,9 +63,9 @@ def analyse(files, block_idx=1): files = sorted(files) for f in files: - data_energy = extract_observable(f, name='energy', block_idx=block_idx) - data_nav = extract_observable(f, name='nav', block_idx=block_idx) - data = pd.concat([data_energy, data_nav['Nav']], axis=1) + data_energy = extract_observable(f, name="energy", block_idx=block_idx) + data_nav = extract_observable(f, name="nav", block_idx=block_idx) + data = pd.concat([data_energy, data_nav["Nav"]], axis=1) md = get_metadata(f) keys = set_info(data, md) sims.append(data[1:]) @@ -108,20 +95,6 @@ def nav_mu(mu, coeffs): return numpy.polyval(coeffs, mu) -def find_chem_pot(data, target, vol, order=3, plot=False): - print(f"# System volume: {vol}.") - print(f"# Target number of electrons: {vol * target}.") - nav = data.Nav.values / vol - nav_error = data.Nav_error.values / vol - # Half filling special case where error bar is zero. - zeros = numpy.where(nav_error == 0)[0] - nav_error[zeros] = 1e-8 - mus = data.mu.values - delta = nav - target - s = 0 - e = len(delta) - rmin = None - def main(args): """Run reblocking and data analysis on PAUXY output. @@ -136,7 +109,7 @@ def main(args): """ options = parse_args(args) - if '*' in options.filenames[0]: + if "*" in options.filenames[0]: files = glob.glob(options.filenames[0]) else: @@ -144,20 +117,8 @@ def main(args): data = analyse(files) - if options.fit_chem_pot: - name = get_sys_param(files[0], 'name') - vol = 1. - mu = find_chem_pot(data, options.nav, vol, - order=options.order, plot=options.plot) - - if mu is not None: - print("# Optimal chemical potential found to be: {}.".format(mu)) - - else: - print("# Failed to find chemical potential.") - print(data.to_string(index=False)) -if __name__ == '__main__': +if __name__ == "__main__": main(sys.argv[1:]) diff --git a/ipie/addons/thermal/estimators/energy.py b/ipie/addons/thermal/estimators/energy.py index 7ebe8637..5cf34b03 100644 --- a/ipie/addons/thermal/estimators/energy.py +++ b/ipie/addons/thermal/estimators/energy.py @@ -18,20 +18,19 @@ from typing import Union -from ipie.utils.backend import arraylib as xp -from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol -from ipie.estimators.energy import EnergyEstimator - -from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers -from ipie.addons.thermal.estimators.thermal import one_rdm_from_G from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.estimators.energy import EnergyEstimator +from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol +from ipie.utils.backend import arraylib as xp def local_energy( - hamiltonian: Union[GenericRealChol, GenericComplexChol], - walkers: UHFThermalWalkers): + hamiltonian: Union[GenericRealChol, GenericComplexChol], walkers: UHFThermalWalkers +): energies = xp.zeros((walkers.nwalkers, 3), dtype=xp.complex128) - + for iw in range(walkers.nwalkers): # Want the full Green's function when calculating observables. walkers.calc_greens_function(iw, slice_ix=walkers.stack[iw].nslice) @@ -46,9 +45,13 @@ class ThermalEnergyEstimator(EnergyEstimator): def __init__(self, system=None, hamiltonian=None, trial=None, filename=None): super().__init__(system=system, ham=hamiltonian, trial=trial, filename=filename) - def compute_estimator(self, walkers, hamiltonian, trial, istep=1): + def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None): # Need to be able to dispatch here. # Re-calculated Green's function in `local_energy`. + if hamiltonian is None: + raise ValueError("Hamiltonian must not be none.") + if walkers is None: + raise ValueError("walkers must not be none.") energy = local_energy(hamiltonian, walkers) self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) self._data["EDenom"] = xp.sum(walkers.weight) diff --git a/ipie/addons/thermal/estimators/generic.py b/ipie/addons/thermal/estimators/generic.py index 0ad971fa..cddc7a15 100644 --- a/ipie/addons/thermal/estimators/generic.py +++ b/ipie/addons/thermal/estimators/generic.py @@ -50,7 +50,7 @@ def local_energy_generic_cholesky(hamiltonian: GenericRealChol, P): Xb = hamiltonian.chol.T.dot(Pb.real.ravel()) + 1.0j * hamiltonian.chol.T.dot(Pb.imag.ravel()) X = Xa + Xb ecoul = 0.5 * numpy.dot(X, X) - + # Ex. PaT = Pa.T.copy() PbT = Pb.T.copy() @@ -94,7 +94,7 @@ def local_energy_generic_cholesky(hamiltonian: GenericComplexChol, P): nbasis = hamiltonian.nbasis nchol = hamiltonian.nchol Pa, Pb = P[0], P[1] - + # Ecoul. XAa = hamiltonian.A.T.dot(Pa.ravel()) XAb = hamiltonian.A.T.dot(Pb.ravel()) diff --git a/ipie/addons/thermal/estimators/greens_function.py b/ipie/addons/thermal/estimators/greens_function.py index c39abd19..562aa170 100644 --- a/ipie/addons/thermal/estimators/greens_function.py +++ b/ipie/addons/thermal/estimators/greens_function.py @@ -19,6 +19,7 @@ import numpy import scipy.linalg + def greens_function(A): r"""Construct Greens function from density matrix. @@ -50,7 +51,7 @@ def greens_function(A): def greens_function_qr_strat(walkers, iw, slice_ix=None, inplace=True): - """Compute the Green's function for walker with index `iw` at time + """Compute the Green's function for walker with index `iw` at time `slice_ix`. Uses the Stratification method (DOI 10.1109/IPDPS.2012.37) """ stack_iw = walkers.stack[iw] @@ -119,18 +120,22 @@ def greens_function_qr_strat(walkers, iw, slice_ix=None, inplace=True): if inplace: if spin == 0: walkers.Ga[iw] = numpy.dot( - numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T) + ) else: walkers.Gb[iw] = numpy.dot( - numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T) + ) else: if spin == 0: Ga_iw = numpy.dot( - numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T) + ) else: Gb_iw = numpy.dot( - numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T) + ) return Ga_iw, Gb_iw diff --git a/ipie/addons/thermal/estimators/handler.py b/ipie/addons/thermal/estimators/handler.py index 8ae4b8b7..87bc5c9e 100644 --- a/ipie/addons/thermal/estimators/handler.py +++ b/ipie/addons/thermal/estimators/handler.py @@ -21,11 +21,10 @@ import os from typing import Tuple, Union -from ipie.config import config, MPI -from ipie.estimators.handler import EstimatorHandler - from ipie.addons.thermal.estimators.energy import ThermalEnergyEstimator from ipie.addons.thermal.estimators.particle_number import ThermalNumberEstimator +from ipie.config import config, MPI +from ipie.estimators.handler import EstimatorHandler # Some supported (non-custom) estimators _predefined_estimators = { @@ -106,16 +105,14 @@ def __init__( # TODO: Replace this, should be built outside for obs in observables: try: - est = _predefined_estimators[obs]( - hamiltonian=hamiltonian, - trial=trial) + est = _predefined_estimators[obs](hamiltonian=hamiltonian, trial=trial) self[obs] = est except KeyError: raise RuntimeError(f"unknown observable: {obs}") if verbose: print("# Finished settting up estimator object.") - def compute_estimators(self, hamiltonian, trial, walker_batch): + def compute_estimators(self, system=None, hamiltonian=None, trial=None, walker_batch=None): """Update estimators with bached walkers. Parameters @@ -132,7 +129,7 @@ def compute_estimators(self, hamiltonian, trial, walker_batch): # TODO: generalize for different block groups (loop over groups) offset = self.num_walker_props for k, e in self.items(): - e.compute_estimator(walker_batch, hamiltonian, trial) + e.compute_estimator(walkers=walker_batch, hamiltonian=hamiltonian, trial=trial) start = offset + self.get_offset(k) end = start + int(self[k].size) self.local_estimates[start:end] += e.data @@ -154,7 +151,7 @@ def print_time_slice(self, comm, time_slice, walker_state): offset = walker_state.size if comm.rank == 0: - k = 'energy' + k = "energy" e = self[k] start = offset + self.get_offset(k) end = start + int(self[k].size) @@ -162,15 +159,14 @@ def print_time_slice(self, comm, time_slice, walker_state): e.post_reduce_hook(estim_data) etotal = estim_data[e.get_index("ETotal")] - k = 'nav' + k = "nav" e = self[k] start = offset + self.get_offset(k) end = start + int(self[k].size) estim_data = self.global_estimates[start:end] e.post_reduce_hook(estim_data) nav = estim_data[e.get_index("Nav")] - + print(f"cut : {time_slice} {nav.real} {etotal.real}") self.zero() - diff --git a/ipie/addons/thermal/estimators/local_energy.py b/ipie/addons/thermal/estimators/local_energy.py index 34a06775..5b57ae78 100644 --- a/ipie/addons/thermal/estimators/local_energy.py +++ b/ipie/addons/thermal/estimators/local_energy.py @@ -25,10 +25,12 @@ from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky from ipie.addons.thermal.estimators.thermal import one_rdm_from_G + def local_energy_from_density_matrix( - hamiltonian: Union[GenericRealChol, GenericComplexChol], - trial: Union[OneBody, MeanField], - P: numpy.ndarray): + hamiltonian: Union[GenericRealChol, GenericComplexChol], + trial: Union[OneBody, MeanField], + P: numpy.ndarray, +): """Compute local energy from a given density matrix P. Parameters @@ -48,5 +50,6 @@ def local_energy_from_density_matrix( assert len(P) == 2 return local_energy_generic_cholesky(hamiltonian, P) + def local_energy(hamiltonian, walker, trial): return local_energy_from_density_matrix(hamiltonian, trial, one_rdm_from_G(walker.G)) diff --git a/ipie/addons/thermal/estimators/particle_number.py b/ipie/addons/thermal/estimators/particle_number.py index 63fa1f60..e5a6a888 100644 --- a/ipie/addons/thermal/estimators/particle_number.py +++ b/ipie/addons/thermal/estimators/particle_number.py @@ -17,9 +17,10 @@ # import numpy -from ipie.utils.backend import arraylib as xp -from ipie.estimators.estimator_base import EstimatorBase + from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.estimators.estimator_base import EstimatorBase +from ipie.utils.backend import arraylib as xp def particle_number(dmat: numpy.ndarray): @@ -64,12 +65,13 @@ def __init__(self, hamiltonian=None, trial=None, filename=None): # Must specify that we're dealing with array valued estimator. self.scalar_estimator = True - def compute_estimator(self, walkers, hamiltonian, trial): + def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None): + if walkers is None: + raise ValueError("Walkers cannot be none in estimator.") for iw in range(walkers.nwalkers): # Want the full Green's function when calculating observables. walkers.calc_greens_function(iw, slice_ix=walkers.stack[iw].nslice) - nav_iw = particle_number(one_rdm_from_G( - xp.array([walkers.Ga[iw], walkers.Gb[iw]]))) + nav_iw = particle_number(one_rdm_from_G(xp.array([walkers.Ga[iw], walkers.Gb[iw]]))) self._data["NavNumer"] += walkers.weight[iw] * nav_iw.real self._data["NavDenom"] = sum(walkers.weight) diff --git a/ipie/addons/thermal/estimators/tests/__init__.py b/ipie/addons/thermal/estimators/tests/__init__.py index 381108b3..e2aed039 100644 --- a/ipie/addons/thermal/estimators/tests/__init__.py +++ b/ipie/addons/thermal/estimators/tests/__init__.py @@ -11,4 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - diff --git a/ipie/addons/thermal/estimators/tests/test_estimators.py b/ipie/addons/thermal/estimators/tests/test_estimators.py index 15c9ebb7..e25dc6f7 100644 --- a/ipie/addons/thermal/estimators/tests/test_estimators.py +++ b/ipie/addons/thermal/estimators/tests/test_estimators.py @@ -16,19 +16,19 @@ # Joonho Lee # -import pytest import tempfile -import numpy from typing import Tuple, Union -from ipie.config import MPI -from ipie.hamiltonians.generic import Generic as HamGeneric -from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +import numpy +import pytest from ipie.addons.thermal.estimators.energy import ThermalEnergyEstimator -from ipie.addons.thermal.estimators.particle_number import ThermalNumberEstimator from ipie.addons.thermal.estimators.handler import ThermalEstimatorHandler +from ipie.addons.thermal.estimators.particle_number import ThermalNumberEstimator from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers +from ipie.config import MPI +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol # System params. nup = 5 @@ -38,7 +38,7 @@ nbasis = 10 # Thermal AFQMC params. -mu = -10. +mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 10 @@ -51,23 +51,34 @@ seed = 7 numpy.random.seed(seed) + @pytest.mark.unit def test_energy_estimator(): # Test. objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] assert isinstance(hamiltonian, GenericRealChol) chol = hamiltonian.chol # GenericRealChol. re_estim = ThermalEnergyEstimator(hamiltonian=hamiltonian, trial=trial) - re_estim.compute_estimator(walkers, hamiltonian, trial) + re_estim.compute_estimator(walkers=walkers, hamiltonian=hamiltonian, trial=trial) assert len(re_estim.names) == 5 assert re_estim["ENumer"].real == pytest.approx(24.66552451455761) assert re_estim["ETotal"] == pytest.approx(0.0) @@ -80,17 +91,20 @@ def test_energy_estimator(): header = re_estim.header_to_text data_to_text = re_estim.data_to_text(tmp) assert len(data_to_text.split()) == 5 - + # GenericComplexChol. cx_chol = numpy.array(chol, dtype=numpy.complex128) cx_hamiltonian = HamGeneric( - numpy.array(hamiltonian.H1, dtype=numpy.complex128), cx_chol, - hamiltonian.ecore, verbose=False) + numpy.array(hamiltonian.H1, dtype=numpy.complex128), + cx_chol, + hamiltonian.ecore, + verbose=False, + ) assert isinstance(cx_hamiltonian, GenericComplexChol) cx_estim = ThermalEnergyEstimator(hamiltonian=cx_hamiltonian, trial=trial) - cx_estim.compute_estimator(walkers, cx_hamiltonian, trial) + cx_estim.compute_estimator(walkers=walkers, hamiltonian=cx_hamiltonian, trial=trial) assert len(cx_estim.names) == 5 assert cx_estim["ENumer"].real == pytest.approx(24.66552451455761) assert cx_estim["ETotal"] == pytest.approx(0.0) @@ -103,23 +117,33 @@ def test_energy_estimator(): header = cx_estim.header_to_text data_to_text = cx_estim.data_to_text(tmp) assert len(data_to_text.split()) == 5 - + numpy.testing.assert_allclose(re_estim.data, cx_estim.data) @pytest.mark.unit def test_number_estimator(): # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=True, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=True, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] estim = ThermalNumberEstimator(hamiltonian=hamiltonian, trial=trial) - estim.compute_estimator(walkers, hamiltonian, trial) + estim.compute_estimator(walkers=walkers, hamiltonian=hamiltonian, trial=trial) assert len(estim.names) == 3 assert estim["NavNumer"].real == pytest.approx(ne * nwalkers) assert estim["Nav"] == pytest.approx(0.0) @@ -132,41 +156,44 @@ def test_number_estimator(): header = estim.header_to_text data_to_text = estim.data_to_text(tmp) assert len(data_to_text.split()) == 3 - + @pytest.mark.unit def test_estimator_handler(): with tempfile.NamedTemporaryFile() as tmp1, tempfile.NamedTemporaryFile() as tmp2: # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=True, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - - estim = ThermalEnergyEstimator(hamiltonian=hamiltonian, trial=trial, - filename=tmp1.name) + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=True, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + + estim = ThermalEnergyEstimator(hamiltonian=hamiltonian, trial=trial, filename=tmp1.name) estim.print_to_stdout = False comm = MPI.COMM_WORLD handler = ThermalEstimatorHandler( - comm, - hamiltonian, - trial, - observables=("energy",), - filename=tmp2.name) + comm, hamiltonian, trial, observables=("energy",), filename=tmp2.name + ) handler["energy1"] = estim handler.json_string = "" handler.initialize(comm) - handler.compute_estimators(hamiltonian, trial, walkers) + handler.compute_estimators(hamiltonian=hamiltonian, trial=trial, walker_batch=walkers) if __name__ == "__main__": test_energy_estimator() test_number_estimator() test_estimator_handler() - - - diff --git a/ipie/addons/thermal/estimators/tests/test_generic.py b/ipie/addons/thermal/estimators/tests/test_generic.py index 4205ab58..50eb221c 100644 --- a/ipie/addons/thermal/estimators/tests/test_generic.py +++ b/ipie/addons/thermal/estimators/tests/test_generic.py @@ -22,6 +22,7 @@ try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers + _no_cython = False except ModuleNotFoundError: @@ -33,10 +34,13 @@ from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G -from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky +from ipie.legacy.estimators.generic import ( + local_energy_generic_cholesky as legacy_local_energy_generic_cholesky, +) comm = MPI.COMM_WORLD + @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @pytest.mark.unit def test_local_energy_cholesky(mf_trial=False): @@ -47,7 +51,7 @@ def test_local_energy_cholesky(mf_trial=False): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 @@ -62,30 +66,49 @@ def test_local_energy_cholesky(mf_trial=False): # Test. objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - P = one_rdm_from_G(trial.G) + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + P = one_rdm_from_G(trial.G) eloc = local_energy_generic_cholesky(hamiltonian, P) # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] legacy_P = legacy_one_rdm_from_G(legacy_trial.G) - legacy_eloc = legacy_local_energy_generic_cholesky( - legacy_system, legacy_hamiltonian, legacy_P) - + legacy_eloc = legacy_local_energy_generic_cholesky(legacy_system, legacy_hamiltonian, legacy_P) + numpy.testing.assert_allclose(trial.G, legacy_trial.G, atol=1e-10) numpy.testing.assert_allclose(P, legacy_P, atol=1e-10) numpy.testing.assert_allclose(eloc, legacy_eloc, atol=1e-10) -if __name__ == '__main__': +if __name__ == "__main__": test_local_energy_cholesky(mf_trial=True) diff --git a/ipie/addons/thermal/estimators/tests/test_generic_complex.py b/ipie/addons/thermal/estimators/tests/test_generic_complex.py index 2902f61b..7de7e36a 100644 --- a/ipie/addons/thermal/estimators/tests/test_generic_complex.py +++ b/ipie/addons/thermal/estimators/tests/test_generic_complex.py @@ -33,7 +33,7 @@ nbasis = 10 # Thermal AFQMC params. -mu = -10. +mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 @@ -46,29 +46,43 @@ seed = 7 numpy.random.seed(seed) + @pytest.mark.unit def test_local_energy_vs_real(): # Test. objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - walkers = objs['walkers'] - hamiltonian = objs['hamiltonian'] + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + walkers = objs["walkers"] + hamiltonian = objs["hamiltonian"] assert isinstance(hamiltonian, GenericRealChol) - + chol = hamiltonian.chol cx_chol = numpy.array(chol, dtype=numpy.complex128) cx_hamiltonian = HamGeneric( - numpy.array(hamiltonian.H1, dtype=numpy.complex128), cx_chol, - hamiltonian.ecore, verbose=False) + numpy.array(hamiltonian.H1, dtype=numpy.complex128), + cx_chol, + hamiltonian.ecore, + verbose=False, + ) assert isinstance(cx_hamiltonian, GenericComplexChol) - + for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) energy = local_energy_generic_cholesky(hamiltonian, P) cx_energy = local_energy_generic_cholesky(cx_hamiltonian, P) numpy.testing.assert_allclose(energy, cx_energy, atol=1e-10) @@ -78,45 +92,56 @@ def test_local_energy_vs_real(): def test_local_energy_vs_eri(): # Test. objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, debug=debug, complex_integrals=True, with_eri=True, - seed=seed, verbose=verbose) - trial = objs['trial'] - walkers = objs['walkers'] - hamiltonian = objs['hamiltonian'] + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + debug=debug, + complex_integrals=True, + with_eri=True, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + walkers = objs["walkers"] + hamiltonian = objs["hamiltonian"] assert isinstance(hamiltonian, GenericComplexChol) - eri = objs['eri'].reshape(nbasis, nbasis, nbasis, nbasis) - + eri = objs["eri"].reshape(nbasis, nbasis, nbasis, nbasis) + chol = hamiltonian.chol.copy() nchol = chol.shape[1] chol = chol.reshape(nbasis, nbasis, nchol) # Check if chol and eri are consistent. - eri_chol = numpy.einsum('mnx,slx->mnls', chol, chol.conj()) + eri_chol = numpy.einsum("mnx,slx->mnls", chol, chol.conj()) numpy.testing.assert_allclose(eri, eri_chol, atol=1e-10) for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) Pa, Pb = P Ptot = Pa + Pb etot, e1, e2 = local_energy_generic_cholesky(hamiltonian, P) # Test 1-body term. h1e = hamiltonian.H1[0] - e1ref = numpy.einsum('ij,ij->', h1e, Ptot) + e1ref = numpy.einsum("ij,ij->", h1e, Ptot) numpy.testing.assert_allclose(e1, e1ref, atol=1e-10) # Test 2-body term. - ecoul = 0.5 * numpy.einsum('ijkl,ij,kl->', eri, Ptot, Ptot) - exx = -0.5 * numpy.einsum('ijkl,il,kj->', eri, Pa, Pa) - exx -= 0.5 * numpy.einsum('ijkl,il,kj->', eri, Pb, Pb) + ecoul = 0.5 * numpy.einsum("ijkl,ij,kl->", eri, Ptot, Ptot) + exx = -0.5 * numpy.einsum("ijkl,il,kj->", eri, Pa, Pa) + exx -= 0.5 * numpy.einsum("ijkl,il,kj->", eri, Pb, Pb) e2ref = ecoul + exx numpy.testing.assert_allclose(e2, e2ref, atol=1e-10) - + etotref = e1ref + e2ref numpy.testing.assert_allclose(etot, etotref, atol=1e-10) -if __name__ == '__main__': +if __name__ == "__main__": test_local_energy_vs_real() test_local_energy_vs_eri() diff --git a/ipie/addons/thermal/estimators/thermal.py b/ipie/addons/thermal/estimators/thermal.py index 198f965c..b478c1ed 100644 --- a/ipie/addons/thermal/estimators/thermal.py +++ b/ipie/addons/thermal/estimators/thermal.py @@ -19,6 +19,7 @@ import numpy import scipy.linalg + def one_rdm_from_G(G): r"""Compute one-particle reduced density matrix from Green's function. @@ -38,6 +39,7 @@ def one_rdm_from_G(G): I = numpy.identity(G.shape[-1]) return numpy.array([I - G[0].T, I - G[1].T], dtype=numpy.complex128) + def one_rdm_stable(BT, num_slices): nbasis = BT.shape[-1] G = [] diff --git a/ipie/addons/thermal/propagation/force_bias.py b/ipie/addons/thermal/propagation/force_bias.py index 47fd2138..815f8aa4 100644 --- a/ipie/addons/thermal/propagation/force_bias.py +++ b/ipie/addons/thermal/propagation/force_bias.py @@ -23,6 +23,7 @@ from ipie.addons.thermal.estimators.thermal import one_rdm_from_G from ipie.utils.backend import arraylib as xp + @plum.dispatch def construct_force_bias(hamiltonian: GenericRealChol, walkers): r"""Compute optimal force bias. @@ -69,4 +70,3 @@ def construct_force_bias(hamiltonian: GenericComplexChol, walkers): vbias[iw, nchol:] = hamiltonian.B.T.dot(P[0].ravel()) + hamiltonian.B.T.dot(P[1].ravel()) return vbias - diff --git a/ipie/addons/thermal/propagation/operations.py b/ipie/addons/thermal/propagation/operations.py index 661f96b1..90b311fd 100644 --- a/ipie/addons/thermal/propagation/operations.py +++ b/ipie/addons/thermal/propagation/operations.py @@ -18,6 +18,7 @@ from ipie.utils.backend import arraylib as xp + def apply_exponential(VHS, exp_nmax): """Apply exponential propagator of the HS transformation @@ -42,4 +43,4 @@ def apply_exponential(VHS, exp_nmax): Temp = VHS.dot(Temp) / n phi += Temp - return phi # Shape (nbasis, nbasis). + return phi # Shape (nbasis, nbasis). diff --git a/ipie/addons/thermal/propagation/phaseless_base.py b/ipie/addons/thermal/propagation/phaseless_base.py index c5fac4dc..42578482 100644 --- a/ipie/addons/thermal/propagation/phaseless_base.py +++ b/ipie/addons/thermal/propagation/phaseless_base.py @@ -35,6 +35,7 @@ # TODO: Add lowrank implementation. See: https://github.com/JoonhoLee-Group/ipie/issues/302 # Ref: 10.1103/PhysRevB.80.214116 for bounds. + @plum.dispatch def construct_mean_field_shift(hamiltonian: GenericRealChol, trial): r"""Compute mean field shift. @@ -50,7 +51,7 @@ def construct_mean_field_shift(hamiltonian: GenericRealChol, trial): tmp_real = numpy.dot(hamiltonian.chol.T, P.real) tmp_imag = numpy.dot(hamiltonian.chol.T, P.imag) mf_shift = 1.0j * tmp_real - tmp_imag - return mf_shift # Shape (nchol,). + return mf_shift # Shape (nchol,). @plum.dispatch @@ -69,11 +70,11 @@ def construct_mean_field_shift(hamiltonian: GenericComplexChol, trial): mf_shift = numpy.zeros(hamiltonian.nfields, dtype=hamiltonian.chol.dtype) mf_shift[:nchol] = 1j * numpy.dot(hamiltonian.A.T, P.ravel()) mf_shift[nchol:] = 1j * numpy.dot(hamiltonian.B.T, P.ravel()) - return mf_shift # Shape (nchol,). + return mf_shift # Shape (nchol,). class PhaselessBase(ContinuousBase): - """A base class for generic continuous HS transform FT-AFQMC propagators.""" + """A base class for generic continuous HS transform FT-AFQMC propagators.""" def __init__(self, timestep, mu, lowrank=False, exp_nmax=6, verbose=False): super().__init__(timestep, verbose=verbose) @@ -88,7 +89,6 @@ def __init__(self, timestep, mu, lowrank=False, exp_nmax=6, verbose=False): self.lowrank = lowrank self.exp_nmax = exp_nmax - def build(self, hamiltonian, trial=None, walkers=None, mpi_handler=None, verbose=False): # dt/2 one-body propagator start = time.time() @@ -111,7 +111,6 @@ def build(self, hamiltonian, trial=None, walkers=None, mpi_handler=None, verbose self.mf_core = hamiltonian.ecore + 0.5 * numpy.dot(self.mf_shift, self.mf_shift) self.mf_const_fac = cmath.exp(-self.dt * self.mf_core) - @plum.dispatch def construct_one_body_propagator(self, hamiltonian: GenericRealChol): r"""Construct mean-field shifted one-body propagator. @@ -132,12 +131,11 @@ def construct_one_body_propagator(self, hamiltonian: GenericRealChol): shift = 1j * numpy.einsum("mx,x->m", hamiltonian.chol, self.mf_shift).reshape(nb, nb) muN = self.mu * numpy.identity(nb, dtype=hamiltonian.H1.dtype) H1 = hamiltonian.h1e_mod - numpy.array([shift + muN, shift + muN]) - expH1 = numpy.array([ - scipy.linalg.expm(-0.5 * self.dt * H1[0]), - scipy.linalg.expm(-0.5 * self.dt * H1[1])]) - return expH1 # Shape (nbasis, nbasis). + expH1 = numpy.array( + [scipy.linalg.expm(-0.5 * self.dt * H1[0]), scipy.linalg.expm(-0.5 * self.dt * H1[1])] + ) + return expH1 # Shape (nbasis, nbasis). - @plum.dispatch def construct_one_body_propagator(self, hamiltonian: GenericComplexChol): r"""Construct mean-field shifted one-body propagator. @@ -161,11 +159,10 @@ def construct_one_body_propagator(self, hamiltonian: GenericComplexChol): shift += 1j * numpy.einsum("mx,x->m", hamiltonian.B, self.mf_shift[nchol:]).reshape(nb, nb) muN = self.mu * numpy.identity(nb, dtype=hamiltonian.H1.dtype) H1 = hamiltonian.h1e_mod - numpy.array([shift + muN, shift + muN]) - expH1 = numpy.array([ - scipy.linalg.expm(-0.5 * self.dt * H1[0]), - scipy.linalg.expm(-0.5 * self.dt * H1[1])]) - return expH1 # Shape (nbasis, nbasis). - + expH1 = numpy.array( + [scipy.linalg.expm(-0.5 * self.dt * H1[0]), scipy.linalg.expm(-0.5 * self.dt * H1[1])] + ) + return expH1 # Shape (nbasis, nbasis). def construct_two_body_propagator(self, walkers, hamiltonian, trial, debug=False): r"""Construct two-body propagator. @@ -187,7 +184,7 @@ def construct_two_body_propagator(self, walkers, hamiltonian, trial, debug=False Trial dnsity matrix. """ # Optimal force bias - xbar = xp.zeros((walkers.nwalkers, hamiltonian.nfields)) + xbar = xp.zeros((walkers.nwalkers, hamiltonian.nfields)) start_time = time.time() self.vbias = construct_force_bias(hamiltonian, walkers) xbar = -self.sqrt_dt * (1j * self.vbias - self.mf_shift) @@ -198,18 +195,22 @@ def construct_two_body_propagator(self, walkers, hamiltonian, trial, debug=False # Normally distrubted auxiliary fields. xi = xp.random.normal(0.0, 1.0, hamiltonian.nfields * walkers.nwalkers).reshape( - walkers.nwalkers, hamiltonian.nfields) + walkers.nwalkers, hamiltonian.nfields + ) + + if debug: + self.xi = xi # For debugging. + xshifted = xi - xbar # Shape (nwalkers, nfields). - if debug: self.xi = xi # For debugging. - xshifted = xi - xbar # Shape (nwalkers, nfields). - # Constant factor arising from force bias and mean field shift - cmf = -self.sqrt_dt * xp.einsum("wx,x->w", xshifted, self.mf_shift) # Shape (nwalkers,). + cmf = -self.sqrt_dt * xp.einsum("wx,x->w", xshifted, self.mf_shift) # Shape (nwalkers,). # Constant factor arising from shifting the propability distribution. - cfb = xp.einsum("wx,wx->w", xi, xbar) - 0.5 * xp.einsum("wx,wx->w", xbar, xbar) # Shape (nwalkers,). + cfb = xp.einsum("wx,wx->w", xi, xbar) - 0.5 * xp.einsum( + "wx,wx->w", xbar, xbar + ) # Shape (nwalkers,). - xshifted = xshifted.T.copy() # Shape (nfields, nwalkers). - VHS = self.construct_VHS(hamiltonian, xshifted) # Shape (nwalkers, nbasis, nbasis). + xshifted = xshifted.T.copy() # Shape (nfields, nwalkers). + VHS = self.construct_VHS(hamiltonian, xshifted) # Shape (nwalkers, nbasis, nbasis). return cmf, cfb, xshifted, VHS def propagate_walkers_one_body(self, walkers): @@ -217,11 +218,12 @@ def propagate_walkers_one_body(self, walkers): def propagate_walkers_two_body(self, walkers, hamiltonian, trial): pass - - def propagate_walkers(self, walkers, hamiltonian, trial, eshift=0., debug=False): + + def propagate_walkers(self, walkers, hamiltonian, trial, eshift=0.0, debug=False): start_time = time.time() cmf, cfb, xshifted, VHS = self.construct_two_body_propagator( - walkers, hamiltonian, trial, debug=debug) + walkers, hamiltonian, trial, debug=debug + ) assert walkers.nwalkers == xshifted.shape[-1] self.timer.tvhs += time.time() - start_time assert len(VHS.shape) == 3 @@ -230,7 +232,7 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=0., debug=False) for iw in range(walkers.nwalkers): stack = walkers.stack[iw] phi = xp.identity(VHS[iw].shape[-1], dtype=xp.complex128) - BV = apply_exponential(phi, VHS[iw], self.exp_nmax) # Shape (nbasis, nbasis). + BV = apply_exponential(phi, VHS[iw], self.exp_nmax) # Shape (nbasis, nbasis). B = numpy.array([BV.dot(self.BH1[0]), BV.dot(self.BH1[1])]) B = numpy.array([self.BH1[0].dot(B[0]), self.BH1[1].dot(B[1])]) @@ -250,42 +252,39 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=0., debug=False) # Now apply phaseless approximation. # Use legacy thermal weight update for now. self.update_weight_legacy(walkers, iw, G, cfb, cmf, eshift) - #self.update_weight(walkers, iw, G, cfb, cmf, eshift) + # self.update_weight(walkers, iw, G, cfb, cmf, eshift) self.timer.tupdate += time.time() - start_time - def update_weight(self, walkers, iw, G, cfb, cmf, eshift): - """Update weight for walker `iw`. - """ + """Update weight for walker `iw`.""" M0a = scipy.linalg.det(G[0], check_finite=False) M0b = scipy.linalg.det(G[1], check_finite=False) Mnewa = scipy.linalg.det(walkers.Ga[iw], check_finite=False) Mnewb = scipy.linalg.det(walkers.Gb[iw], check_finite=False) - + # ovlp = det( G^{-1} ) - ovlp_ratio = (M0a * M0b) / (Mnewa * Mnewb) # ovlp_new / ovlp_old - hybrid_energy = -(xp.log(ovlp_ratio) + cfb[iw] + cmf[iw]) / self.dt # Scalar. + ovlp_ratio = (M0a * M0b) / (Mnewa * Mnewb) # ovlp_new / ovlp_old + hybrid_energy = -(xp.log(ovlp_ratio) + cfb[iw] + cmf[iw]) / self.dt # Scalar. hybrid_energy = self.apply_bound_hybrid(hybrid_energy, eshift) importance_function = xp.exp( - -self.dt * (0.5 * (hybrid_energy + walkers.hybrid_energy) - eshift)) + -self.dt * (0.5 * (hybrid_energy + walkers.hybrid_energy) - eshift) + ) - # Splitting w_k = |I(x, \bar{x}, |phi_k>)| e^{i theta_k}, where `k` + # Splitting w_k = |I(x, \bar{x}, |phi_k>)| e^{i theta_k}, where `k` # labels the time slice. magn = xp.abs(importance_function) walkers.hybrid_energy = hybrid_energy - dtheta = (-self.dt * hybrid_energy - cfb[iw]).imag # Scalar. - cosine_fac = xp.amax([0., xp.cos(dtheta)]) + dtheta = (-self.dt * hybrid_energy - cfb[iw]).imag # Scalar. + cosine_fac = xp.amax([0.0, xp.cos(dtheta)]) walkers.weight[iw] *= magn * cosine_fac walkers.M0a[iw] = Mnewa walkers.M0b[iw] = Mnewb - def update_weight_legacy(self, walkers, iw, G, cfb, cmf, eshift): - """Update weight for walker `iw` using legacy code. - """ - #M0a = walkers.M0a[iw] - #M0b = walkers.M0b[iw] + """Update weight for walker `iw` using legacy code.""" + # M0a = walkers.M0a[iw] + # M0b = walkers.M0b[iw] M0a = scipy.linalg.det(G[0], check_finite=False) M0b = scipy.linalg.det(G[1], check_finite=False) Mnewa = scipy.linalg.det(walkers.Ga[iw], check_finite=False) @@ -313,10 +312,10 @@ def update_weight_legacy(self, walkers, iw, G, cfb, cmf, eshift): walkers.M0b[iw] = Mnewb else: - walkers.weight[iw] = 0. + walkers.weight[iw] = 0.0 except ZeroDivisionError: - walkers.weight[iw] = 0. + walkers.weight[iw] = 0.0 def apply_bound_force_bias(self, xbar, max_bound=1.0): absxbar = xp.abs(xbar) @@ -328,7 +327,6 @@ def apply_bound_force_bias(self, xbar, max_bound=1.0): self.nfb_trig += xp.sum(idx_to_rescale) return xbar - def apply_bound_hybrid(self, ehyb, eshift): # Shift is a number but ehyb is not # For initial steps until first estimator communication, `eshift` will be # zero and hybrid energy can be incorrect. So just avoid capping for @@ -340,9 +338,7 @@ def apply_bound_hybrid(self, ehyb, eshift): # Shift is a number but ehyb is not emin = eshift.real - self.ebound return xp.minimum(emax, xp.maximum(ehyb, emin)) - # Form VHS. @abstractmethod def construct_VHS(self, hamiltonian, xshifted): pass - diff --git a/ipie/addons/thermal/propagation/phaseless_generic.py b/ipie/addons/thermal/propagation/phaseless_generic.py index 6f196ed8..b60fd1aa 100644 --- a/ipie/addons/thermal/propagation/phaseless_generic.py +++ b/ipie/addons/thermal/propagation/phaseless_generic.py @@ -28,20 +28,18 @@ class PhaselessGeneric(PhaselessBase): def __init__(self, time_step, mu, exp_nmax=6, lowrank=False, verbose=False): super().__init__(time_step, mu, lowrank=lowrank, exp_nmax=exp_nmax, verbose=verbose) - + @plum.dispatch def construct_VHS(self, hamiltonian: GenericRealChol, xshifted: xp.ndarray): - """Includes `nwalkers`. - """ - nwalkers = xshifted.shape[-1] # Shape (nfields, nwalkers). - VHS = hamiltonian.chol.dot(xshifted) # Shape (nbasis^2, nwalkers). + """Includes `nwalkers`.""" + nwalkers = xshifted.shape[-1] # Shape (nfields, nwalkers). + VHS = hamiltonian.chol.dot(xshifted) # Shape (nbasis^2, nwalkers). VHS = self.isqrt_dt * VHS.T.reshape(nwalkers, hamiltonian.nbasis, hamiltonian.nbasis) - return VHS # Shape (nwalkers, nbasis, nbasis). - + return VHS # Shape (nwalkers, nbasis, nbasis). + @plum.dispatch def construct_VHS(self, hamiltonian: GenericComplexChol, xshifted: xp.ndarray): - """Includes `nwalkers`. - """ + """Includes `nwalkers`.""" nwalkers = xshifted.shape[-1] nchol = hamiltonian.nchol VHS = self.isqrt_dt * ( diff --git a/ipie/addons/thermal/propagation/tests/test_prop_generic.py b/ipie/addons/thermal/propagation/tests/test_prop_generic.py index c2e7377d..6674545c 100644 --- a/ipie/addons/thermal/propagation/tests/test_prop_generic.py +++ b/ipie/addons/thermal/propagation/tests/test_prop_generic.py @@ -22,6 +22,7 @@ try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False except ModuleNotFoundError: @@ -32,7 +33,9 @@ from ipie.addons.thermal.estimators.thermal import one_rdm_from_G from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers -from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky +from ipie.legacy.estimators.generic import ( + local_energy_generic_cholesky as legacy_local_energy_generic_cholesky, +) from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G comm = MPI.COMM_WORLD @@ -44,7 +47,7 @@ nbasis = 10 # Thermal AFQMC params. -mu = -10. +mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 @@ -63,71 +66,129 @@ @pytest.mark.unit def test_mf_shift(): # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - hamiltonian = objs['hamiltonian'] - propagator = objs['propagator'] + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + hamiltonian = objs["hamiltonian"] + propagator = objs["propagator"] # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_propagator = legacy_objs['propagator'] - - numpy.testing.assert_almost_equal(legacy_propagator.propagator.mf_shift, - propagator.mf_shift, decimal=10) + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_propagator = legacy_objs["propagator"] + + numpy.testing.assert_almost_equal( + legacy_propagator.propagator.mf_shift, propagator.mf_shift, decimal=10 + ) @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @pytest.mark.unit def test_BH1(): # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - hamiltonian = objs['hamiltonian'] - propagator = objs['propagator'] + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + hamiltonian = objs["hamiltonian"] + propagator = objs["propagator"] # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_propagator = legacy_objs['propagator'] - - numpy.testing.assert_almost_equal(legacy_propagator.propagator.BH1, - propagator.BH1, decimal=10) + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_propagator = legacy_objs["propagator"] + + numpy.testing.assert_almost_equal(legacy_propagator.propagator.BH1, propagator.BH1, decimal=10) @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @pytest.mark.unit def test_construct_two_body_propagator(): # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] - + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] + cmf, cfb, xshifted, VHS = propagator.construct_two_body_propagator( - walkers, hamiltonian, trial, debug=True) + walkers, hamiltonian, trial, debug=True + ) legacy_cmf = [] legacy_cfb = [] @@ -136,13 +197,13 @@ def test_construct_two_body_propagator(): for iw in range(walkers.nwalkers): _cmf, _cfb, _xshifted, _VHS = legacy_propagator.two_body_propagator( - legacy_walkers.walkers[iw], legacy_hamiltonian, - legacy_trial, xi=propagator.xi[iw]) + legacy_walkers.walkers[iw], legacy_hamiltonian, legacy_trial, xi=propagator.xi[iw] + ) legacy_cmf.append(_cmf) legacy_cfb.append(_cfb) legacy_xshifted.append(_xshifted) legacy_VHS.append(_VHS) - + legacy_xshifted = numpy.array(legacy_xshifted).T numpy.testing.assert_almost_equal(legacy_cmf, cmf, decimal=10) @@ -155,46 +216,72 @@ def test_construct_two_body_propagator(): @pytest.mark.unit def test_phaseless_generic_propagator(): # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] for t in range(walkers.stack[0].nslice): for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) eloc = local_energy_generic_cholesky(hamiltonian, P) legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) legacy_eloc = legacy_local_energy_generic_cholesky( - legacy_system, legacy_hamiltonian, legacy_P) + legacy_system, legacy_hamiltonian, legacy_P + ) numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) numpy.testing.assert_allclose(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw]) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10 + ) numpy.testing.assert_almost_equal(legacy_P, P, decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10 + ) propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) + legacy_hamiltonian, legacy_trial, legacy_walkers, legacy_propagator, xi=propagator.xi + ) if __name__ == "__main__": diff --git a/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py b/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py index c7a788c5..c773890c 100644 --- a/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py +++ b/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py @@ -23,6 +23,7 @@ from ipie.addons.thermal.utils.legacy_testing import build_legacy_ueg_test_case_handlers from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers from ipie.legacy.estimators.ueg import local_energy_ueg as legacy_local_energy_ueg + _no_cython = False except ModuleNotFoundError: @@ -45,62 +46,82 @@ def test_phaseless_ueg_propagator(): nup = 7 ndown = 7 nelec = (nup, ndown) - rs = 1. - ecut = 1. + rs = 1.0 + ecut = 1.0 # Thermal AFQMC params. - mu = -1. + mu = -1.0 beta = 0.1 timestep = 0.01 nwalkers = 1 lowrank = False - + debug = True verbose = False if (comm.rank != 0) else True seed = 7 numpy.random.seed(seed) - + # Test. - objs = build_ueg_test_case_handlers( - nelec, rs, ecut, mu, beta, timestep, nwalkers=nwalkers, - lowrank=lowrank, debug=debug, seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] + objs = build_ueg_test_case_handlers( + nelec, + rs, + ecut, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] # Legacy. legacy_objs = build_legacy_ueg_test_case_handlers( - comm, nelec, rs, ecut, mu, beta, timestep, nwalkers=nwalkers, - lowrank=lowrank, seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] + comm, + nelec, + rs, + ecut, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] h1e = legacy_hamiltonian.H1[0] eri = legacy_hamiltonian.eri_4() for t in range(walkers.stack[0].nslice): for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) eloc = local_energy_generic_cholesky(hamiltonian, P) legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) legacy_eloc = legacy_local_energy_ueg(legacy_system, legacy_hamiltonian, legacy_P) - + legacy_Pa, legacy_Pb = legacy_P legacy_Ptot = legacy_Pa + legacy_Pb - ref_e1 = numpy.einsum('ij,ij->', h1e, legacy_Ptot) + ref_e1 = numpy.einsum("ij,ij->", h1e, legacy_Ptot) Ptot = legacy_Ptot Pa = legacy_Pa Pb = legacy_Pb - ecoul = 0.5 * numpy.einsum('ijkl,ij,kl->', eri, Ptot, Ptot) - exx = -0.5 * numpy.einsum('ijkl,il,kj->', eri, Pa, Pa) - exx -= 0.5 * numpy.einsum('ijkl,il,kj->', eri, Pb, Pb) + ecoul = 0.5 * numpy.einsum("ijkl,ij,kl->", eri, Ptot, Ptot) + exx = -0.5 * numpy.einsum("ijkl,il,kj->", eri, Pa, Pa) + exx -= 0.5 * numpy.einsum("ijkl,il,kj->", eri, Pb, Pb) ref_e2 = ecoul + exx ref_eloc = (ref_e1 + ref_e2, ref_e1, ref_e2) @@ -110,15 +131,23 @@ def test_phaseless_ueg_propagator(): numpy.testing.assert_allclose(legacy_eloc, ref_eloc, atol=1e-10) numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10 + ) propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) + legacy_hamiltonian, legacy_trial, legacy_walkers, legacy_propagator, xi=propagator.xi + ) if __name__ == "__main__": diff --git a/ipie/addons/thermal/qmc/calc.py b/ipie/addons/thermal/qmc/calc.py index 659ce515..fc2adf87 100644 --- a/ipie/addons/thermal/qmc/calc.py +++ b/ipie/addons/thermal/qmc/calc.py @@ -18,18 +18,16 @@ """Helper Routines for setting up a calculation""" -from ipie.config import MPI -from ipie.utils.mpi import MPIHandler -from ipie.utils.io import get_input_value - -from ipie.systems.utils import get_system -from ipie.hamiltonians.utils import get_hamiltonian - -from ipie.addons.thermal.trial.utils import get_trial_density_matrix -from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers from ipie.addons.thermal.propagation.propagator import Propagator from ipie.addons.thermal.qmc.options import ThermalQMCOpts, ThermalQMCParams from ipie.addons.thermal.qmc.thermal_afqmc import ThermalAFQMC +from ipie.addons.thermal.trial.utils import get_trial_density_matrix +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.config import MPI +from ipie.hamiltonians.utils import get_hamiltonian +from ipie.systems.utils import get_system +from ipie.utils.io import get_input_value +from ipie.utils.mpi import MPIHandler def get_driver(options: dict, comm: MPI.COMM_WORLD) -> ThermalAFQMC: @@ -56,7 +54,9 @@ def get_driver(options: dict, comm: MPI.COMM_WORLD) -> ThermalAFQMC: if comm.rank != 0: verbosity = 0 - lowrank = get_input_value(wlk_opts, "lowrank", default=False, alias=["low_rank"], verbose=verbosity) + lowrank = get_input_value( + wlk_opts, "lowrank", default=False, alias=["low_rank"], verbose=verbosity + ) batched = get_input_value(qmc_opts, "batched", default=False, verbose=verbosity) debug = get_input_value(qmc_opts, "debug", default=False, verbose=verbosity) @@ -87,11 +87,19 @@ def get_driver(options: dict, comm: MPI.COMM_WORLD) -> ThermalAFQMC: comm=comm, verbose=verbosity, ) - stack_size = get_input_value(wlk_opts, 'stack_size', default=10, verbose=verbosity) - lowrank_thresh = get_input_value(wlk_opts, 'lowrank_thresh', default=1e-6, alias=["low_rank_thresh"], verbose=verbosity) + stack_size = get_input_value(wlk_opts, "stack_size", default=10, verbose=verbosity) + lowrank_thresh = get_input_value( + wlk_opts, "lowrank_thresh", default=1e-6, alias=["low_rank_thresh"], verbose=verbosity + ) walkers = UHFThermalWalkers( - trial, hamiltonian.nbasis, qmc.nwalkers, stack_size=stack_size, - lowrank=lowrank, lowrank_thresh=lowrank_thresh, verbose=verbosity) + trial, + hamiltonian.nbasis, + qmc.nwalkers, + stack_size=stack_size, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + verbose=verbosity, + ) if (comm.rank == 0) and (qmc.nsteps > 1): print("Only num_steps_per_block = 1 allowed in thermal code! Resetting to value of 1.") @@ -112,7 +120,6 @@ def get_driver(options: dict, comm: MPI.COMM_WORLD) -> ThermalAFQMC: propagator = Propagator[type(hamiltonian)](params.timestep, params.mu) propagator.build(hamiltonian, trial, walkers, mpi_handler) afqmc = ThermalAFQMC( - system, hamiltonian, trial, walkers, @@ -136,7 +143,7 @@ def build_thermal_afqmc_driver( ): if comm.rank != 0: verbosity = 0 - + sys_opts = {"nup": nelec[0], "ndown": nelec[1]} ham_opts = {"integrals": hamiltonian_file} qmc_opts = {"rng_seed": seed} diff --git a/ipie/addons/thermal/qmc/options.py b/ipie/addons/thermal/qmc/options.py index 671bedab..cc2a28c2 100644 --- a/ipie/addons/thermal/qmc/options.py +++ b/ipie/addons/thermal/qmc/options.py @@ -60,6 +60,7 @@ class ThermalQMCOpts(QMCOpts): beta : float Inverse temperature. """ + # pylint: disable=dangerous-default-value # TODO: Remove this class / replace with dataclass def __init__(self, inputs={}, verbose=False): @@ -87,7 +88,7 @@ class ThermalQMCParams(QMCParams): ---------- mu : float Chemical potential. - beta : float + beta : float Inverse temperature. num_walkers : int Number of walkers **per** core / task / computational unit. @@ -104,19 +105,19 @@ class ThermalQMCParams(QMCParams): pop_control_freq : int Frequency at which population control occurs. rng_seed : int - The random number seed. If run in parallel the seeds on other cores / + The random number seed. If run in parallel the seeds on other cores / threads are determined from this. """ + # Due to structure of FT algorithm, `num_steps_per_block` is fixed at 1. # Overide whatever input for backward compatibility. - num_steps_per_block: ClassVar[int] = 1 + num_steps_per_block: ClassVar[int] = 1 mu: Optional[float] = None beta: Optional[float] = None - pop_control_method: str = 'pair_branch' - + pop_control_method: str = "pair_branch" + def __post_init__(self): if self.mu is None: raise TypeError("__init__ missing 1 required argument: 'mu'") if self.beta is None: raise TypeError("__init__ missing 1 required argument: 'beta'") - diff --git a/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py b/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py index 88e138b1..206de2b5 100644 --- a/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py +++ b/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py @@ -18,26 +18,29 @@ import json import tempfile -import h5py import uuid -import pytest -import numpy from typing import Union +import h5py +import numpy +import pytest + try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_driver_generic_test_instance + _no_cython = False except ModuleNotFoundError: _no_cython = True -from ipie.config import MPI -from ipie.analysis.extraction import ( - extract_test_data_hdf5, - extract_data, - extract_observable, - extract_mixed_estimates) from ipie.addons.thermal.utils.testing import build_driver_generic_test_instance +from ipie.analysis.extraction import ( + extract_data, + extract_mixed_estimates, + extract_observable, + extract_test_data_hdf5, +) +from ipie.config import MPI comm = MPI.COMM_WORLD serial_test = comm.size == 1 @@ -69,12 +72,12 @@ def test_thermal_afqmc(): nblocks = 12 stabilize_freq = 10 pop_control_freq = 1 - pop_control_method = 'pair_branch' - #pop_control_method = 'comb' + pop_control_method = "pair_branch" + # pop_control_method = 'comb' lowrank = False verbose = 0 if (comm.rank != 0) else 1 - # Local energy evaluation in legacy code seems wrong. + # Local energy evaluation in legacy code seems wrong. complex_integrals = False debug = True seed = 7 @@ -85,15 +88,28 @@ def test_thermal_afqmc(): # Test. # --------------------------------------------------------------------- afqmc = build_driver_generic_test_instance( - nelec, nbasis, mu, beta, timestep, nblocks, nwalkers=nwalkers, - lowrank=lowrank, pop_control_method=pop_control_method, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - complex_integrals=complex_integrals, debug=debug, seed=seed, - verbose=verbose) + nelec, + nbasis, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) afqmc.finalise() - afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) - + afqmc.estimators.compute_estimators( + hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers + ) + test_energy_data = None test_energy_numer = None test_energy_denom = None @@ -104,17 +120,27 @@ def test_thermal_afqmc(): test_energy_numer = afqmc.estimators["energy"]["ENumer"] test_energy_denom = afqmc.estimators["energy"]["EDenom"] test_number_data = extract_observable(afqmc.estimators.filename, "nav") - + # --------------------------------------------------------------------- # Legacy. # --------------------------------------------------------------------- legacy_afqmc = build_legacy_driver_generic_test_instance( - afqmc.hamiltonian, comm, nelec, mu, beta, timestep, - nblocks, nwalkers=nwalkers, lowrank=lowrank, - stabilize_freq=stabilize_freq, - pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, seed=seed, - estimator_filename=tmpf2.name, verbose=verbose) + afqmc.hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + seed=seed, + estimator_filename=tmpf2.name, + verbose=verbose, + ) legacy_afqmc.run(comm=comm) legacy_afqmc.finalise(verbose=False) legacy_afqmc.estimators.estimators["mixed"].update( @@ -124,7 +150,8 @@ def test_thermal_afqmc(): legacy_afqmc.trial, legacy_afqmc.walk, 0, - legacy_afqmc.propagators.free_projection) + legacy_afqmc.propagators.free_projection, + ) legacy_mixed_data = None enum = None @@ -136,33 +163,41 @@ def test_thermal_afqmc(): enum = legacy_afqmc.estimators.estimators["mixed"].names legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] - + # Check. assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) - + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real) + ) assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real) + ) assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real) + ) assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real) + ) assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real) + ) assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real) + ) assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) - + numpy.mean(legacy_mixed_data.Nav.values[:-1].real) + ) -if __name__ == '__main__': - test_thermal_afqmc() +if __name__ == "__main__": + test_thermal_afqmc() diff --git a/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py b/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py index a8906a84..c52a8511 100644 --- a/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py +++ b/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py @@ -16,32 +16,35 @@ # Joonho Lee # -import os -import sys import json +import os import pprint +import sys import tempfile -import h5py import uuid -import pytest -import numpy from typing import Union +import h5py +import numpy +import pytest + try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_driver_ueg_test_instance + _no_cython = False except ModuleNotFoundError: _no_cython = True -from ipie.config import MPI -from ipie.analysis.extraction import ( - get_metadata, - extract_test_data_hdf5, - extract_data, - extract_observable, - extract_mixed_estimates) from ipie.addons.thermal.utils.testing import build_driver_ueg_test_instance +from ipie.analysis.extraction import ( + extract_data, + extract_mixed_estimates, + extract_observable, + extract_test_data_hdf5, + get_metadata, +) +from ipie.config import MPI comm = MPI.COMM_WORLD serial_test = comm.size == 1 @@ -67,7 +70,7 @@ def compare_test_data(ref_data, test_data): elif k == "EHybrid": alias.append("HybridEnergy") - + err = 0 ref = ref_data[k] @@ -77,7 +80,8 @@ def compare_test_data(ref_data, test_data): comparison[k] = ( numpy.array(ref), numpy.array(test), - numpy.max(numpy.abs(numpy.array(ref) - numpy.array(test))) < 1e-10) + numpy.max(numpy.abs(numpy.array(ref) - numpy.array(test))) < 1e-10, + ) except KeyError: err += 1 @@ -95,40 +99,54 @@ def test_thermal_afqmc_1walker(against_ref=False): nup = 7 ndown = 7 nelec = (nup, ndown) - rs = 1. - ecut = 1. + rs = 1.0 + ecut = 1.0 # Thermal AFQMC params. - mu = -1. + mu = -1.0 beta = 0.1 timestep = 0.01 nwalkers = 1 nblocks = 11 - + stabilize_freq = 10 pop_control_freq = 1 # `pop_control_method` doesn't matter for 1 walker. pop_control_method = "pair_branch" - #pop_control_method = "comb" + # pop_control_method = "comb" lowrank = False - + verbose = False if (comm.rank != 0) else True debug = True seed = 7 numpy.random.seed(seed) - + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: # --------------------------------------------------------------------- # Test. # --------------------------------------------------------------------- afqmc = build_driver_ueg_test_instance( - nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, - lowrank=lowrank, pop_control_method=pop_control_method, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - debug=debug, seed=seed, verbose=verbose) + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + debug=debug, + seed=seed, + verbose=verbose, + ) afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) afqmc.finalise() - afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + afqmc.estimators.compute_estimators( + hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers + ) test_energy_data = None test_energy_numer = None @@ -140,17 +158,28 @@ def test_thermal_afqmc_1walker(against_ref=False): test_energy_numer = afqmc.estimators["energy"]["ENumer"] test_energy_denom = afqmc.estimators["energy"]["EDenom"] test_number_data = extract_observable(afqmc.estimators.filename, "nav") - + # --------------------------------------------------------------------- # Legacy. # --------------------------------------------------------------------- legacy_afqmc = build_legacy_driver_ueg_test_instance( - comm, nelec, rs, ecut, mu, beta, timestep, nblocks, - nwalkers=nwalkers, lowrank=lowrank, - stabilize_freq=stabilize_freq, - pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, seed=seed, - estimator_filename=tmpf2.name, verbose=verbose) + comm, + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + seed=seed, + estimator_filename=tmpf2.name, + verbose=verbose, + ) legacy_afqmc.run(comm=comm) legacy_afqmc.finalise(verbose=False) legacy_afqmc.estimators.estimators["mixed"].update( @@ -160,7 +189,8 @@ def test_thermal_afqmc_1walker(against_ref=False): legacy_afqmc.trial, legacy_afqmc.walk, 0, - legacy_afqmc.propagators.free_projection) + legacy_afqmc.propagators.free_projection, + ) legacy_mixed_data = None enum = None @@ -172,39 +202,50 @@ def test_thermal_afqmc_1walker(against_ref=False): enum = legacy_afqmc.estimators.estimators["mixed"].names legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] - + # Check. assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) - + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real) + ) assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real) + ) assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real) + ) assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real) + ) assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real) + ) assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real) + ) assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) - + numpy.mean(legacy_mixed_data.Nav.values[:-1].real) + ) + # --------------------------------------------------------------------- # Test against reference data. if against_ref: - _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _data_dir = ( + os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + ) _legacy_test_dir = "ueg" _legacy_test = _data_dir + _legacy_test_dir + "/reference_1walker.json" - + test_name = _legacy_test_dir with open(_legacy_test, "r") as f: ref_data = json.load(f) @@ -214,10 +255,10 @@ def test_thermal_afqmc_1walker(against_ref=False): _test_number_data = test_number_data[::skip_val].to_dict(orient="list") energy_comparison = compare_test_data(ref_data, _test_energy_data) number_comparison = compare_test_data(ref_data, _test_number_data) - - print('\nenergy comparison:') + + print("\nenergy comparison:") pprint.pprint(energy_comparison) - print('\nnumber comparison:') + print("\nnumber comparison:") pprint.pprint(number_comparison) local_err_count = 0 @@ -225,13 +266,17 @@ def test_thermal_afqmc_1walker(against_ref=False): for k, v in energy_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") for k, v in number_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") if local_err_count == 0: @@ -245,11 +290,11 @@ def test_thermal_afqmc(against_ref=False): nup = 7 ndown = 7 nelec = (nup, ndown) - rs = 1. - ecut = 1. + rs = 1.0 + ecut = 1.0 # Thermal AFQMC params. - mu = -1. + mu = -1.0 beta = 0.1 timestep = 0.01 nwalkers = 32 @@ -259,26 +304,40 @@ def test_thermal_afqmc(against_ref=False): stabilize_freq = 10 pop_control_freq = 1 pop_control_method = "pair_branch" - #pop_control_method = "comb" + # pop_control_method = "comb" lowrank = False - + verbose = False if (comm.rank != 0) else True debug = True seed = 7 numpy.random.seed(seed) - + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: # --------------------------------------------------------------------- # Test. # --------------------------------------------------------------------- afqmc = build_driver_ueg_test_instance( - nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, - lowrank=lowrank, pop_control_method=pop_control_method, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - debug=debug, seed=seed, verbose=verbose) + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + debug=debug, + seed=seed, + verbose=verbose, + ) afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) afqmc.finalise() - afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + afqmc.estimators.compute_estimators( + hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers + ) test_energy_data = None test_energy_numer = None @@ -290,17 +349,28 @@ def test_thermal_afqmc(against_ref=False): test_energy_numer = afqmc.estimators["energy"]["ENumer"] test_energy_denom = afqmc.estimators["energy"]["EDenom"] test_number_data = extract_observable(afqmc.estimators.filename, "nav") - + # --------------------------------------------------------------------- # Legacy. # --------------------------------------------------------------------- legacy_afqmc = build_legacy_driver_ueg_test_instance( - comm, nelec, rs, ecut, mu, beta, timestep, nblocks, - nwalkers=nwalkers, lowrank=lowrank, - stabilize_freq=stabilize_freq, - pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, seed=seed, - estimator_filename=tmpf2.name, verbose=verbose) + comm, + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + seed=seed, + estimator_filename=tmpf2.name, + verbose=verbose, + ) legacy_afqmc.run(comm=comm) legacy_afqmc.finalise(verbose=False) legacy_afqmc.estimators.estimators["mixed"].update( @@ -310,7 +380,8 @@ def test_thermal_afqmc(against_ref=False): legacy_afqmc.trial, legacy_afqmc.walk, 0, - legacy_afqmc.propagators.free_projection) + legacy_afqmc.propagators.free_projection, + ) legacy_mixed_data = None enum = None @@ -322,39 +393,50 @@ def test_thermal_afqmc(against_ref=False): enum = legacy_afqmc.estimators.estimators["mixed"].names legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] - + # Check. assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) - + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real) + ) assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real) + ) assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real) + ) assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real) + ) assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real) + ) assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real) + ) assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + numpy.mean(legacy_mixed_data.Nav.values[:-1].real) + ) # --------------------------------------------------------------------- # Test against reference data. if against_ref: - _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _data_dir = ( + os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + ) _legacy_test_dir = "ueg" _legacy_test = _data_dir + _legacy_test_dir + "/reference_nompi.json" - + test_name = _legacy_test_dir with open(_legacy_test, "r") as f: ref_data = json.load(f) @@ -364,10 +446,10 @@ def test_thermal_afqmc(against_ref=False): _test_number_data = test_number_data[::skip_val].to_dict(orient="list") energy_comparison = compare_test_data(ref_data, _test_energy_data) number_comparison = compare_test_data(ref_data, _test_number_data) - - print('\nenergy comparison:') + + print("\nenergy comparison:") pprint.pprint(energy_comparison) - print('\nnumber comparison:') + print("\nnumber comparison:") pprint.pprint(number_comparison) local_err_count = 0 @@ -375,13 +457,17 @@ def test_thermal_afqmc(against_ref=False): for k, v in energy_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") for k, v in number_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") if local_err_count == 0: @@ -395,11 +481,11 @@ def test_thermal_afqmc_mpi(against_ref=False): nup = 7 ndown = 7 nelec = (nup, ndown) - rs = 1. - ecut = 1. + rs = 1.0 + ecut = 1.0 # Thermal AFQMC params. - mu = -1. + mu = -1.0 beta = 0.1 timestep = 0.01 nwalkers = 32 // comm.size @@ -409,26 +495,40 @@ def test_thermal_afqmc_mpi(against_ref=False): stabilize_freq = 10 pop_control_freq = 1 pop_control_method = "pair_branch" - #pop_control_method = "comb" + # pop_control_method = "comb" lowrank = False - + verbose = False if (comm.rank != 0) else True debug = True seed = 7 numpy.random.seed(seed) - + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: # --------------------------------------------------------------------- # Test. # --------------------------------------------------------------------- afqmc = build_driver_ueg_test_instance( - nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, - lowrank=lowrank, pop_control_method=pop_control_method, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - debug=debug, seed=seed, verbose=verbose) + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + debug=debug, + seed=seed, + verbose=verbose, + ) afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) afqmc.finalise() - afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + afqmc.estimators.compute_estimators( + hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers + ) test_energy_data = None test_energy_numer = None @@ -440,17 +540,28 @@ def test_thermal_afqmc_mpi(against_ref=False): test_energy_numer = afqmc.estimators["energy"]["ENumer"] test_energy_denom = afqmc.estimators["energy"]["EDenom"] test_number_data = extract_observable(afqmc.estimators.filename, "nav") - + # --------------------------------------------------------------------- # Legacy. # --------------------------------------------------------------------- legacy_afqmc = build_legacy_driver_ueg_test_instance( - comm, nelec, rs, ecut, mu, beta, timestep, nblocks, - nwalkers=nwalkers, lowrank=lowrank, - stabilize_freq=stabilize_freq, - pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, seed=seed, - estimator_filename=tmpf2.name, verbose=verbose) + comm, + nelec, + rs, + ecut, + mu, + beta, + timestep, + nblocks, + nwalkers=nwalkers, + lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + seed=seed, + estimator_filename=tmpf2.name, + verbose=verbose, + ) legacy_afqmc.run(comm=comm) legacy_afqmc.finalise(verbose=False) legacy_afqmc.estimators.estimators["mixed"].update( @@ -460,7 +571,8 @@ def test_thermal_afqmc_mpi(against_ref=False): legacy_afqmc.trial, legacy_afqmc.walk, 0, - legacy_afqmc.propagators.free_projection) + legacy_afqmc.propagators.free_projection, + ) legacy_mixed_data = None enum = None @@ -472,39 +584,50 @@ def test_thermal_afqmc_mpi(against_ref=False): enum = legacy_afqmc.estimators.estimators["mixed"].names legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] - + # Check. assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) - + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real) + ) assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real) + ) assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real) + ) assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real) + ) assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real) + ) assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real) + ) assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real) + ) assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( - numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + numpy.mean(legacy_mixed_data.Nav.values[:-1].real) + ) # --------------------------------------------------------------------- # Test against reference data. if against_ref: - _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _data_dir = ( + os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + ) _legacy_test_dir = "ueg" _legacy_test = _data_dir + _legacy_test_dir + "/reference.json" - + test_name = _legacy_test_dir with open(_legacy_test, "r") as f: ref_data = json.load(f) @@ -514,10 +637,10 @@ def test_thermal_afqmc_mpi(against_ref=False): _test_number_data = test_number_data[::skip_val].to_dict(orient="list") energy_comparison = compare_test_data(ref_data, _test_energy_data) number_comparison = compare_test_data(ref_data, _test_number_data) - - print('\nenergy comparison:') + + print("\nenergy comparison:") pprint.pprint(energy_comparison) - print('\nnumber comparison:') + print("\nnumber comparison:") pprint.pprint(number_comparison) local_err_count = 0 @@ -525,21 +648,24 @@ def test_thermal_afqmc_mpi(against_ref=False): for k, v in energy_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") for k, v in number_comparison.items(): if not v[-1]: local_err_count += 1 - print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print( + f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}" + ) print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") if local_err_count == 0: print(f"\n*** PASSED : {test_name} ***\n") -if __name__ == '__main__': +if __name__ == "__main__": test_thermal_afqmc_1walker(against_ref=True) test_thermal_afqmc(against_ref=True) - #test_thermal_afqmc_mpi(against_ref=True) - + # test_thermal_afqmc_mpi(against_ref=True) diff --git a/ipie/addons/thermal/qmc/thermal_afqmc.py b/ipie/addons/thermal/qmc/thermal_afqmc.py index 3cea31ee..dc682650 100644 --- a/ipie/addons/thermal/qmc/thermal_afqmc.py +++ b/ipie/addons/thermal/qmc/thermal_afqmc.py @@ -15,31 +15,28 @@ # Authors: Fionn Malone # Joonho Lee # - """Driver to perform Thermal AFQMC calculation""" - -import numpy -import time import json +import time from typing import Dict, Optional, Tuple -from ipie.addons.thermal.walkers.pop_controller import ThermalPopController -from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers -from ipie.addons.thermal.propagation.propagator import Propagator +import numpy + from ipie.addons.thermal.estimators.handler import ThermalEstimatorHandler +from ipie.addons.thermal.propagation.propagator import Propagator from ipie.addons.thermal.qmc.options import ThermalQMCParams - -from ipie.utils.io import to_json +from ipie.addons.thermal.walkers.pop_controller import ThermalPopController +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.estimators.estimator_base import EstimatorBase +from ipie.qmc.afqmc import AFQMCBase from ipie.utils.backend import arraylib as xp from ipie.utils.backend import synchronize +from ipie.utils.io import to_json from ipie.utils.mpi import MPIHandler -from ipie.systems.generic import Generic -from ipie.estimators.estimator_base import EstimatorBase from ipie.walkers.base_walkers import WalkerAccumulator -from ipie.qmc.afqmc import AFQMC -class ThermalAFQMC(AFQMC): +class ThermalAFQMC(AFQMCBase): """Thermal AFQMC driver. Parameters @@ -63,48 +60,52 @@ class ThermalAFQMC(AFQMC): Seed deduced from params.rng_seed which is generally different on each MPI process. """ - def __init__(self, - system, # For compatibility with 0T AFQMC code. - hamiltonian, - trial, - walkers, - propagator, - mpi_handler, - params: ThermalQMCParams, - debug: bool = False, - verbose: int = 0): - super().__init__(system, hamiltonian, trial, walkers, propagator, mpi_handler, params, verbose) + + def __init__( + self, + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params: ThermalQMCParams, + debug: bool = False, + verbose: int = 0, + ): + super().__init__( + None, hamiltonian, trial, walkers, propagator, mpi_handler, params, verbose + ) self.debug = debug if self.debug and verbose: - print('# Using legacy `update_weights`.') + print("# Using legacy `update_weights`.") - @staticmethod def build( - nelec: Tuple[int, int], - mu: float, - beta: float, - hamiltonian, - trial, - nwalkers: int = 100, - stack_size: int = 10, - seed: int = None, - nblocks: int = 100, - timestep: float = 0.005, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - debug: bool = False, - verbose: int = 0, - mpi_handler=None,) -> "Thermal AFQMC": + nelec: Tuple[int, int], + mu: float, + beta: float, + hamiltonian, + trial, + nwalkers: int = 100, + stack_size: int = 10, + seed: Optional[int] = None, + nblocks: int = 100, + timestep: float = 0.005, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + debug: bool = False, + verbose: int = 0, + mpi_handler=None, + ) -> "ThermalAFQMC": """Factory method to build thermal AFQMC driver from hamiltonian and trial density matrix. Parameters ---------- - nelec : tuple(int, int) + num_elec : tuple(int, int) Number of alpha and beta electrons. mu : float Chemical potential. @@ -140,47 +141,54 @@ def build( else: comm = mpi_handler.comm - + # pylint: disable = no-value-for-parameter params = ThermalQMCParams( - mu=mu, - beta=beta, - num_walkers=nwalkers, - total_num_walkers=nwalkers * comm.size, - num_blocks=nblocks, - timestep=timestep, - num_stblz=stabilize_freq, - pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, - rng_seed=seed) - - system = Generic(nelec) + mu=mu, + beta=beta, + num_walkers=nwalkers, + total_num_walkers=nwalkers * comm.size, + num_blocks=nblocks, + timestep=timestep, + num_stblz=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + rng_seed=seed, + ) + walkers = UHFThermalWalkers( - trial, hamiltonian.nbasis, nwalkers, stack_size=stack_size, - lowrank=lowrank, lowrank_thresh=lowrank_thresh, - mpi_handler=mpi_handler, verbose=verbose) - propagator = Propagator[type(hamiltonian)]( - timestep, mu, lowrank=lowrank, verbose=verbose) - propagator.build(hamiltonian, trial=trial, walkers=walkers, - mpi_handler=mpi_handler, verbose=verbose) + trial, + hamiltonian.nbasis, + nwalkers, + stack_size=stack_size, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + mpi_handler=mpi_handler, + verbose=verbose, + ) + propagator = Propagator[type(hamiltonian)](timestep, mu, lowrank=lowrank, verbose=verbose) + propagator.build( + hamiltonian, trial=trial, walkers=walkers, mpi_handler=mpi_handler, verbose=verbose + ) return ThermalAFQMC( - system, - hamiltonian, - trial, - walkers, - propagator, - mpi_handler, - params, - debug=debug, - verbose=verbose) - - - def run(self, - walkers = None, - verbose: bool = True, - estimator_filename = None, - additional_estimators: Optional[Dict[str, EstimatorBase]] = None, - print_time_slice: bool = False): + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params, + debug=debug, + verbose=verbose, + ) + + def run( + self, + walkers=None, + estimator_filename=None, + verbose: bool = True, + additional_estimators: Optional[Dict[str, EstimatorBase]] = None, + print_time_slice: bool = False, + ): """Perform Thermal AFQMC simulation on state object using open-ended random walk. Parameters @@ -200,21 +208,22 @@ def run(self, # Setup. self.setup_timers() ft_setup = time.time() - eshift = 0. + eshift = 0.0 if walkers is not None: self.walkers = walkers self.pcontrol = ThermalPopController( - self.params.num_walkers, - self.params.num_steps_per_block, - self.mpi_handler, - self.params.pop_control_method, - verbose=self.verbose) - + self.params.num_walkers, + self.params.num_steps_per_block, + self.mpi_handler, + self.params.pop_control_method, + verbose=self.verbose, + ) + self.get_env_info() self.setup_estimators(estimator_filename, additional_estimators=additional_estimators) - + synchronize() comm = self.mpi_handler.comm self.tsetup += time.time() - ft_setup @@ -237,18 +246,20 @@ def run(self, start = time.time() self.propagator.propagate_walkers( - self.walkers, self.hamiltonian, self.trial, eshift, debug=self.debug) + self.walkers, self.hamiltonian, self.trial, eshift, debug=self.debug + ) self.tprop_fbias = self.propagator.timer.tfbias self.tprop_update = self.propagator.timer.tupdate self.tprop_vhs = self.propagator.timer.tvhs self.tprop_gemm = self.propagator.timer.tgemm - + start_clip = time.time() if t > 0: wbound = self.pcontrol.total_weight * 0.10 - xp.clip(self.walkers.weight, a_min=-wbound, a_max=wbound, - out=self.walkers.weight) # In-place clipping. + xp.clip( + self.walkers.weight, a_min=-wbound, a_max=wbound, out=self.walkers.weight + ) # In-place clipping. synchronize() self.tprop_clip += time.time() - start_clip @@ -259,7 +270,7 @@ def run(self, self.tprop_barrier += time.time() - start_barrier self.tprop += time.time() - start - + if (t > 0) and (t % self.params.pop_control_freq == 0): start = time.time() self.pcontrol.pop_control(self.walkers, comm) @@ -269,11 +280,12 @@ def run(self, self.tpopc_recv = self.pcontrol.timer.recv_time self.tpopc_comm = self.pcontrol.timer.communication_time self.tpopc_non_comm = self.pcontrol.timer.non_communication_time - + # Print estimators at each time slice. if print_time_slice: self.estimators.compute_estimators( - self.hamiltonian, self.trial, self.walkers) + hamiltonian=self.hamiltonian, trial=self.trial, walker_batch=self.walkers + ) self.estimators.print_time_slice(comm, t, self.accumulators) # Accumulate weight, hybrid energy etc. across block. @@ -285,10 +297,12 @@ def run(self, start = time.time() if step % self.params.num_steps_per_block == 0: self.estimators.compute_estimators( - self.hamiltonian, self.trial, self.walkers) + hamiltonian=self.hamiltonian, trial=self.trial, walker_batch=self.walkers + ) self.estimators.print_block( - comm, step // self.params.num_steps_per_block, self.accumulators) + comm, step // self.params.num_steps_per_block, self.accumulators + ) self.accumulators.zero() synchronize() @@ -300,17 +314,17 @@ def run(self, else: eshift += self.accumulators.eshift - eshift - self.walkers.reset(self.trial) # Reset stack, weights, phase. + self.walkers.reset(self.trial) # Reset stack, weights, phase. synchronize() self.tstep += time.time() - start_step def setup_estimators( - self, - filename, - additional_estimators: Optional[Dict[str, EstimatorBase]] = None): + self, filename, additional_estimators: Optional[Dict[str, EstimatorBase]] = None + ): self.accumulators = WalkerAccumulator( - ["Weight", "WeightFactor", "HybridEnergy"], self.params.num_steps_per_block) + ["Weight", "WeightFactor", "HybridEnergy"], self.params.num_steps_per_block + ) comm = self.mpi_handler.comm self.estimators = ThermalEstimatorHandler( self.mpi_handler.comm, @@ -318,7 +332,8 @@ def setup_estimators( self.trial, walker_state=self.accumulators, verbose=(comm.rank == 0 and self.verbose), - filename=filename) + filename=filename, + ) if additional_estimators is not None: for k, v in additional_estimators.items(): @@ -331,9 +346,9 @@ def setup_estimators( self.estimators.initialize(comm) # Calculate estimates for initial distribution of walkers. - self.estimators.compute_estimators(self.hamiltonian, - self.trial, self.walkers) + self.estimators.compute_estimators( + hamiltonian=self.hamiltonian, trial=self.trial, walker_batch=self.walkers + ) self.accumulators.update(self.walkers) self.estimators.print_block(comm, 0, self.accumulators) self.accumulators.zero() - diff --git a/ipie/addons/thermal/trial/chem_pot.py b/ipie/addons/thermal/trial/chem_pot.py index 85244d40..d54c87c5 100644 --- a/ipie/addons/thermal/trial/chem_pot.py +++ b/ipie/addons/thermal/trial/chem_pot.py @@ -23,10 +23,10 @@ from ipie.utils.io import format_fixed_width_floats, format_fixed_width_strings -def find_chemical_potential(alt_convention, rho, beta, num_bins, target, - deps=1e-6, max_it=1000, verbose=False): - """Find the chemical potential to match . - """ +def find_chemical_potential( + alt_convention, rho, beta, num_bins, target, deps=1e-6, max_it=1000, verbose=False +): + """Find the chemical potential to match .""" # TODO: some sort of generic starting point independent of # system/temperature dmu1 = dmu2 = 1 @@ -83,5 +83,4 @@ def delta_nav(dm, nav): def compute_rho(rho, mu, beta, sign=1): - return numpy.einsum( - "ijk,k->ijk", rho, numpy.exp(sign * beta * mu * numpy.ones(rho.shape[-1]))) + return numpy.einsum("ijk,k->ijk", rho, numpy.exp(sign * beta * mu * numpy.ones(rho.shape[-1]))) diff --git a/ipie/addons/thermal/trial/mean_field.py b/ipie/addons/thermal/trial/mean_field.py index 4ba6fa78..d205f536 100644 --- a/ipie/addons/thermal/trial/mean_field.py +++ b/ipie/addons/thermal/trial/mean_field.py @@ -20,19 +20,38 @@ import scipy.linalg from ipie.addons.thermal.estimators.generic import fock_generic +from ipie.addons.thermal.estimators.greens_function import greens_function from ipie.addons.thermal.estimators.particle_number import particle_number from ipie.addons.thermal.estimators.thermal import one_rdm_stable -from ipie.addons.thermal.estimators.greens_function import greens_function from ipie.addons.thermal.trial.chem_pot import compute_rho, find_chemical_potential from ipie.addons.thermal.trial.one_body import OneBody + class MeanField(OneBody): - def __init__(self, hamiltonian, nelec, beta, dt, options=None, alt_convention=False, H1=None, verbose=False): + def __init__( + self, + hamiltonian, + nelec, + beta, + dt, + options=None, + alt_convention=False, + H1=None, + verbose=False, + ): if options is None: options = {} - super().__init__(hamiltonian, nelec, beta, dt, options=options, - alt_convention=alt_convention, H1=H1, verbose=verbose) + super().__init__( + hamiltonian, + nelec, + beta, + dt, + options=options, + alt_convention=alt_convention, + H1=H1, + verbose=verbose, + ) if verbose: print("# Building THF density matrix.") @@ -42,12 +61,16 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, alt_convention=Fa self.find_mu = options.get("find_mu", True) self.P, HMF, self.mu = self.thermal_hartree_fock(hamiltonian, beta) muN = self.mu * numpy.eye(hamiltonian.nbasis, dtype=self.G.dtype) - self.dmat = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), - scipy.linalg.expm(-dt * (HMF[1] - muN))]) - self.dmat_inv = numpy.array([scipy.linalg.inv(self.dmat[0], check_finite=False), - scipy.linalg.inv(self.dmat[1], check_finite=False)]) - self.G = numpy.array([greens_function(self.dmat[0]), - greens_function(self.dmat[1])]) + self.dmat = numpy.array( + [scipy.linalg.expm(-dt * (HMF[0] - muN)), scipy.linalg.expm(-dt * (HMF[1] - muN))] + ) + self.dmat_inv = numpy.array( + [ + scipy.linalg.inv(self.dmat[0], check_finite=False), + scipy.linalg.inv(self.dmat[1], check_finite=False), + ] + ) + self.G = numpy.array([greens_function(self.dmat[0]), greens_function(self.dmat[1])]) self.nav = particle_number(self.P).real def thermal_hartree_fock(self, hamiltonian, beta): @@ -63,12 +86,18 @@ def thermal_hartree_fock(self, hamiltonian, beta): print(f"\n# Macro iteration: {it}") HMF = self.scf(hamiltonian, beta, mu_old, P) - rho = numpy.array([scipy.linalg.expm(-dt * HMF[0]), - scipy.linalg.expm(-dt * HMF[1])]) + rho = numpy.array([scipy.linalg.expm(-dt * HMF[0]), scipy.linalg.expm(-dt * HMF[1])]) if self.find_mu: mu = find_chemical_potential( - self.alt_convention, rho, dt, self.nstack, self.nav, - deps=self.deps, max_it=self.max_it, verbose=self.verbose) + self.alt_convention, + rho, + dt, + self.nstack, + self.nav, + deps=self.deps, + max_it=self.max_it, + verbose=self.verbose, + ) else: mu = self.mu @@ -92,17 +121,19 @@ def scf(self, hamiltonian, beta, mu, P): HMF = fock_generic(hamiltonian, P) dt = self.dtau muN = mu * numpy.eye(hamiltonian.nbasis, dtype=self.G.dtype) - rho = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), - scipy.linalg.expm(-dt * (HMF[1] - muN))]) + rho = numpy.array( + [scipy.linalg.expm(-dt * (HMF[0] - muN)), scipy.linalg.expm(-dt * (HMF[1] - muN))] + ) Pold = one_rdm_stable(rho, self.nstack) if self.verbose: print("# Running Thermal SCF.") - for it in range(self.max_scf_it): + for _ in range(self.max_scf_it): HMF = fock_generic(hamiltonian, Pold) - rho = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), - scipy.linalg.expm(-dt * (HMF[1] - muN))]) + rho = numpy.array( + [scipy.linalg.expm(-dt * (HMF[0] - muN)), scipy.linalg.expm(-dt * (HMF[1] - muN))] + ) Pnew = (1 - self.alpha) * one_rdm_stable(rho, self.nstack) + self.alpha * Pold change = numpy.linalg.norm(Pnew - Pold) diff --git a/ipie/addons/thermal/trial/one_body.py b/ipie/addons/thermal/trial/one_body.py index c982d452..50d15b44 100644 --- a/ipie/addons/thermal/trial/one_body.py +++ b/ipie/addons/thermal/trial/one_body.py @@ -19,16 +19,25 @@ import numpy import scipy.linalg +from ipie.addons.thermal.estimators.greens_function import greens_function from ipie.addons.thermal.estimators.particle_number import particle_number from ipie.addons.thermal.estimators.thermal import one_rdm_stable -from ipie.addons.thermal.estimators.greens_function import greens_function from ipie.addons.thermal.trial.chem_pot import compute_rho, find_chemical_potential from ipie.utils.misc import update_stack class OneBody: - def __init__(self, hamiltonian, nelec, beta, dt, options=None, - alt_convention=False, H1=None, verbose=False): + def __init__( + self, + hamiltonian, + nelec, + beta, + dt, + options=None, + alt_convention=False, + H1=None, + verbose=False, + ): if options is None: options = {} @@ -43,7 +52,7 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, except AttributeError: self.H1 = hamiltonian.h1e - + else: self.H1 = H1 @@ -59,7 +68,7 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, if verbose: print(f"# condition number of BT: {cond: 10e}") - + self.nelec = nelec self.nav = options.get("nav", None) @@ -89,8 +98,7 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, self.stack_size = min(self.nslice, int(3.0 / numpy.log10(self.cond))) if verbose: - print("# Initial stack size, # of slices: {}, {}".format( - self.stack_size, self.nslice)) + print(f"# Initial stack size, # of slices: {self.stack_size}, {self.nslice}") # Adjust stack size self.stack_size = update_stack(self.stack_size, self.nslice, verbose=verbose) @@ -109,15 +117,30 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, self.dtau = self.stack_size * dt if self.mu is None: - self.rho = numpy.array([scipy.linalg.expm(-self.dtau * (self.H1[0])), - scipy.linalg.expm(-self.dtau * (self.H1[1]))]) + self.rho = numpy.array( + [ + scipy.linalg.expm(-self.dtau * (self.H1[0])), + scipy.linalg.expm(-self.dtau * (self.H1[1])), + ] + ) self.mu = find_chemical_potential( - self.alt_convention, self.rho, self.dtau, self.nstack, - self.nav, deps=self.deps, max_it=self.max_it, verbose=verbose) - + self.alt_convention, + self.rho, + self.dtau, + self.nstack, + self.nav, + deps=self.deps, + max_it=self.max_it, + verbose=verbose, + ) + else: - self.rho = numpy.array([scipy.linalg.expm(-self.dtau * (self.H1[0])), - scipy.linalg.expm(-self.dtau * (self.H1[1]))]) + self.rho = numpy.array( + [ + scipy.linalg.expm(-self.dtau * (self.H1[0])), + scipy.linalg.expm(-self.dtau * (self.H1[1])), + ] + ) if self.verbose: print(f"# Chemical potential in trial density matrix: {self.mu: .10e}") @@ -129,8 +152,12 @@ def __init__(self, hamiltonian, nelec, beta, dt, options=None, print(f"# Average particle number in trial density matrix: {self.nav}") self.dmat = compute_rho(self.dmat, self.mu, dt, sign=sign) - self.dmat_inv = numpy.array([scipy.linalg.inv(self.dmat[0], check_finite=False), - scipy.linalg.inv(self.dmat[1], check_finite=False)]) + self.dmat_inv = numpy.array( + [ + scipy.linalg.inv(self.dmat[0], check_finite=False), + scipy.linalg.inv(self.dmat[1], check_finite=False), + ] + ) self.G = numpy.array([greens_function(self.dmat[0]), greens_function(self.dmat[1])]) self.error = False diff --git a/ipie/addons/thermal/trial/tests/test_chem_pot.py b/ipie/addons/thermal/trial/tests/test_chem_pot.py index 7fcbcff9..b6b74e7b 100644 --- a/ipie/addons/thermal/trial/tests/test_chem_pot.py +++ b/ipie/addons/thermal/trial/tests/test_chem_pot.py @@ -21,7 +21,9 @@ import pytest from ipie.addons.thermal.trial.chem_pot import find_chemical_potential -from ipie.legacy.trial_density_matrices.chem_pot import find_chemical_potential as legacy_find_chemical_potential +from ipie.legacy.trial_density_matrices.chem_pot import ( + find_chemical_potential as legacy_find_chemical_potential, +) @pytest.mark.unit @@ -36,8 +38,7 @@ def test_find_chemical_potential(): dtau = dt * stack_size h1e = numpy.random.random((nbsf, nbsf)) - rho = numpy.array([scipy.linalg.expm(-dtau * h1e), - scipy.linalg.expm(-dtau * h1e)]) + rho = numpy.array([scipy.linalg.expm(-dtau * h1e), scipy.linalg.expm(-dtau * h1e)]) mu = find_chemical_potential(alt_convention, rho, dt, nstack, nav) legacy_mu = legacy_find_chemical_potential(alt_convention, rho, dt, nstack, nav) @@ -45,8 +46,5 @@ def test_find_chemical_potential(): numpy.testing.assert_allclose(mu, legacy_mu) -if __name__ == '__main__': +if __name__ == "__main__": test_find_chemical_potential() - - - diff --git a/ipie/addons/thermal/trial/tests/test_mean_field.py b/ipie/addons/thermal/trial/tests/test_mean_field.py index 8c3aaf6d..57971904 100644 --- a/ipie/addons/thermal/trial/tests/test_mean_field.py +++ b/ipie/addons/thermal/trial/tests/test_mean_field.py @@ -21,6 +21,7 @@ try: from ipie.legacy.trial_density_matrices.mean_field import MeanField as LegacyMeanField + _no_cython = False except ModuleNotFoundError: @@ -41,42 +42,45 @@ def test_mean_field(): nelec = (nup, ndown) nbasis = 10 - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 - + alt_convention = False sparse = False complex_integrals = True verbose = True sym = 8 - if complex_integrals: sym = 4 - + if complex_integrals: + sym = 4 + # Test. system = Generic(nelec) - h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, - sym=sym, tol=1e-10) - hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), - chol=chol.reshape((-1, nbasis**2)).T.copy(), - ecore=0) + h1e, chol, _, eri = generate_hamiltonian( + nbasis, nelec, cplx=complex_integrals, sym=sym, tol=1e-10 + ) + hamiltonian = HamGeneric( + h1e=numpy.array([h1e, h1e]), chol=chol.reshape((-1, nbasis**2)).T.copy(), ecore=0 + ) trial = MeanField(hamiltonian, nelec, beta, timestep, verbose=verbose) # Lgeacy. legacy_system = Generic(nelec, verbose=verbose) legacy_system.mu = mu legacy_hamiltonian = LegacyHamGeneric( - h1e=hamiltonian.H1, - chol=hamiltonian.chol, - ecore=hamiltonian.ecore, verbose=verbose) + h1e=hamiltonian.H1, chol=hamiltonian.chol, ecore=hamiltonian.ecore, verbose=verbose + ) legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( - (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis) + ) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention legacy_hamiltonian.sparse = sparse - legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) + legacy_trial = LegacyMeanField( + legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose + ) assert trial.nelec == nelec numpy.testing.assert_almost_equal(trial.nav, numpy.sum(nelec), decimal=5) @@ -93,8 +97,5 @@ def test_mean_field(): numpy.testing.assert_allclose(trial.dmat_inv, legacy_trial.dmat_inv) -if __name__ == '__main__': +if __name__ == "__main__": test_mean_field() - - - diff --git a/ipie/addons/thermal/trial/tests/test_one_body.py b/ipie/addons/thermal/trial/tests/test_one_body.py index bdc7dce8..5dfe69cc 100644 --- a/ipie/addons/thermal/trial/tests/test_one_body.py +++ b/ipie/addons/thermal/trial/tests/test_one_body.py @@ -32,7 +32,7 @@ def test_one_body(): nelec = (nup, ndown) nbasis = 10 - mu = -1. + mu = -1.0 beta = 0.1 timestep = 0.01 @@ -40,15 +40,17 @@ def test_one_body(): verbose = True sym = 8 - if complex_integrals: sym = 4 - + if complex_integrals: + sym = 4 + # Test. system = Generic(nelec) - h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, - sym=sym, tol=1e-10) - hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), - chol=chol.reshape((-1, nbasis**2)).T.copy(), - ecore=0) + h1e, chol, _, eri = generate_hamiltonian( + nbasis, nelec, cplx=complex_integrals, sym=sym, tol=1e-10 + ) + hamiltonian = HamGeneric( + h1e=numpy.array([h1e, h1e]), chol=chol.reshape((-1, nbasis**2)).T.copy(), ecore=0 + ) trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) assert trial.nelec == nelec @@ -59,8 +61,5 @@ def test_one_body(): assert trial.G.shape == (2, nbasis, nbasis) -if __name__ == '__main__': +if __name__ == "__main__": test_one_body() - - - diff --git a/ipie/addons/thermal/trial/utils.py b/ipie/addons/thermal/trial/utils.py index 60808a15..0ac49644 100644 --- a/ipie/addons/thermal/trial/utils.py +++ b/ipie/addons/thermal/trial/utils.py @@ -20,8 +20,7 @@ from ipie.addons.thermal.trial.one_body import OneBody -def get_trial_density_matrix(hamiltonian, nelec, beta, dt, options=None, - comm=None, verbose=False): +def get_trial_density_matrix(hamiltonian, nelec, beta, dt, options=None, comm=None, verbose=False): """Wrapper to select trial wavefunction class. Parameters @@ -50,12 +49,26 @@ def get_trial_density_matrix(hamiltonian, nelec, beta, dt, options=None, ) elif trial_type == "one_body": - trial = OneBody(hamiltonian, nelec, beta, dt, options=options, - alt_convention=alt_convention, verbose=verbose) + trial = OneBody( + hamiltonian, + nelec, + beta, + dt, + options=options, + alt_convention=alt_convention, + verbose=verbose, + ) elif trial_type == "thermal_hartree_fock": - trial = MeanField(hamiltonian, nelec, beta, dt, options=options, - alt_convention=alt_convention, verbose=verbose) + trial = MeanField( + hamiltonian, + nelec, + beta, + dt, + options=options, + alt_convention=alt_convention, + verbose=verbose, + ) else: trial = None diff --git a/ipie/addons/thermal/utils/legacy_testing.py b/ipie/addons/thermal/utils/legacy_testing.py index 2e2c829d..feb5d28b 100644 --- a/ipie/addons/thermal/utils/legacy_testing.py +++ b/ipie/addons/thermal/utils/legacy_testing.py @@ -16,94 +16,92 @@ # Joonho Lee # -import numpy from typing import Tuple, Union -from ipie.systems.generic import Generic -from ipie.utils.mpi import MPIHandler +import numpy from ipie.addons.thermal.qmc.options import ThermalQMCOpts - -from ipie.legacy.systems.ueg import UEG as LegacyUEG -from ipie.legacy.hamiltonians.ueg import UEG as LegacyHamUEG from ipie.legacy.hamiltonians._generic import Generic as LegacyHamGeneric -from ipie.legacy.trial_density_matrices.onebody import OneBody as LegacyOneBody -from ipie.legacy.trial_density_matrices.mean_field import MeanField as LegacyMeanField -from ipie.legacy.walkers.handler import Walkers +from ipie.legacy.hamiltonians.ueg import UEG as LegacyHamUEG +from ipie.legacy.qmc.thermal_afqmc import ThermalAFQMC as LegacyThermalAFQMC +from ipie.legacy.systems.ueg import UEG as LegacyUEG from ipie.legacy.thermal_propagation.continuous import Continuous from ipie.legacy.thermal_propagation.planewave import PlaneWave -from ipie.legacy.qmc.thermal_afqmc import ThermalAFQMC as LegacyThermalAFQMC +from ipie.legacy.trial_density_matrices.mean_field import MeanField as LegacyMeanField +from ipie.legacy.trial_density_matrices.onebody import OneBody as LegacyOneBody +from ipie.legacy.walkers.handler import Walkers +from ipie.systems.generic import Generic +from ipie.utils.mpi import MPIHandler -def legacy_propagate_walkers(legacy_hamiltonian, legacy_trial, legacy_walkers, legacy_propagator, xi=None): +def legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, legacy_propagator, xi=None +): if xi is None: xi = [None] * len(legacy_walkers) for iw, walker in enumerate(legacy_walkers.walkers): - legacy_propagator.propagate_walker( - legacy_hamiltonian, walker, legacy_trial, xi=xi[iw]) + legacy_propagator.propagate_walker(legacy_hamiltonian, walker, legacy_trial, xi=xi[iw]) return legacy_walkers def build_legacy_generic_test_case_handlers( - hamiltonian, - comm, - nelec: Tuple[int, int], - mu: float, - beta: float, - timestep: float, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - alt_convention: bool = False, - sparse: bool = False, - mf_trial: bool = True, - propagate: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + hamiltonian, + comm, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + alt_convention: bool = False, + sparse: bool = False, + mf_trial: bool = True, + propagate: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): numpy.random.seed(seed) - + legacy_options = { "walkers": { "stack_size": stack_size, "low_rank": lowrank, "low_rank_thresh": lowrank_thresh, - "pop_control": pop_control_method - }, - - "propagator": { - "optimised": False + "pop_control": pop_control_method, }, + "propagator": {"optimised": False}, } # 1. Build system. legacy_system = Generic(nelec, verbose=verbose) legacy_system.mu = mu - + # 2. Build Hamiltonian. legacy_hamiltonian = LegacyHamGeneric( - h1e=hamiltonian.H1, - chol=hamiltonian.chol, - ecore=hamiltonian.ecore) + h1e=hamiltonian.H1, chol=hamiltonian.chol, ecore=hamiltonian.ecore + ) legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( - (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis) + ) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention legacy_hamiltonian.sparse = sparse - + # 3. Build trial. - legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose) if mf_trial: - legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) - # 4. Build walkers. + legacy_trial = LegacyMeanField( + legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose + ) + # 4. Build walkers. qmc_opts = ThermalQMCOpts() qmc_opts.nwalkers = nwalkers qmc_opts.ntot_walkers = nwalkers @@ -115,50 +113,63 @@ def build_legacy_generic_test_case_handlers( qmc_opts.pop_control_method = pop_control_method qmc_opts.seed = seed - legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, - qmc_opts, walker_opts=legacy_options['walkers'], - verbose=verbose, comm=comm) + legacy_walkers = Walkers( + legacy_system, + legacy_hamiltonian, + legacy_trial, + qmc_opts, + walker_opts=legacy_options["walkers"], + verbose=verbose, + comm=comm, + ) # 5. Build propagator. legacy_propagator = Continuous( - legacy_options["propagator"], qmc_opts, legacy_system, - legacy_hamiltonian, legacy_trial, verbose=verbose, - lowrank=lowrank) + legacy_options["propagator"], + qmc_opts, + legacy_system, + legacy_hamiltonian, + legacy_trial, + verbose=verbose, + lowrank=lowrank, + ) if propagate: - for t in range(legacy_walkers[0].stack.ntime_slices): - for iw, walker in enumerate(legacy_walkers): - legacy_propagator.propagate_walker( - legacy_hamiltonian, walker, legacy_trial) - - legacy_objs = {'system': legacy_system, - 'trial': legacy_trial, - 'hamiltonian': legacy_hamiltonian, - 'walkers': legacy_walkers, - 'propagator': legacy_propagator} + for _ in range(legacy_walkers[0].stack.ntime_slices): + for _, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker(legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = { + "system": legacy_system, + "trial": legacy_trial, + "hamiltonian": legacy_hamiltonian, + "walkers": legacy_walkers, + "propagator": legacy_propagator, + } return legacy_objs - + def build_legacy_generic_test_case_handlers_mpi( - hamiltonian, - mpi_handler: MPIHandler, - nelec: Tuple[int, int], - mu: float, - beta: float, - timestep: float, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - alt_convention: bool = False, - sparse: bool = False, - mf_trial: bool = True, - propagate: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + hamiltonian, + mpi_handler: MPIHandler, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + alt_convention: bool = False, + sparse: bool = False, + mf_trial: bool = True, + propagate: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): numpy.random.seed(seed) comm = mpi_handler.comm @@ -167,37 +178,34 @@ def build_legacy_generic_test_case_handlers_mpi( "stack_size": stack_size, "low_rank": lowrank, "low_rank_thresh": lowrank_thresh, - "pop_control": pop_control_method - }, - - "propagator": { - "optimised": False + "pop_control": pop_control_method, }, + "propagator": {"optimised": False}, } - + # 1. Build system. legacy_system = Generic(nelec, verbose=verbose) legacy_system.mu = mu - + # 2. Build Hamiltonian. legacy_hamiltonian = LegacyHamGeneric( - h1e=hamiltonian.H1, - chol=hamiltonian.chol, - ecore=hamiltonian.ecore) + h1e=hamiltonian.H1, chol=hamiltonian.chol, ecore=hamiltonian.ecore + ) legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( - (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis) + ) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention legacy_hamiltonian.sparse = sparse - + # 3. Build trial. - legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose) if mf_trial: - legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) - # 4. Build walkers. + legacy_trial = LegacyMeanField( + legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose + ) + # 4. Build walkers. qmc_opts = ThermalQMCOpts() qmc_opts.nwalkers = nwalkers qmc_opts.ntot_walkers = nwalkers * comm.size @@ -209,53 +217,66 @@ def build_legacy_generic_test_case_handlers_mpi( qmc_opts.pop_control_method = pop_control_method qmc_opts.seed = seed - legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, - qmc_opts, walker_opts=legacy_options['walkers'], - verbose=verbose, comm=comm) + legacy_walkers = Walkers( + legacy_system, + legacy_hamiltonian, + legacy_trial, + qmc_opts, + walker_opts=legacy_options["walkers"], + verbose=verbose, + comm=comm, + ) # 5. Build propagator. legacy_propagator = Continuous( - legacy_options["propagator"], qmc_opts, legacy_system, - legacy_hamiltonian, legacy_trial, verbose=verbose, - lowrank=lowrank) + legacy_options["propagator"], + qmc_opts, + legacy_system, + legacy_hamiltonian, + legacy_trial, + verbose=verbose, + lowrank=lowrank, + ) if propagate: - for t in range(legacy_walkers[0].stack.ntime_slices): - for iw, walker in enumerate(legacy_walkers): - legacy_propagator.propagate_walker( - legacy_hamiltonian, walker, legacy_trial) - - legacy_objs = {'system': legacy_system, - 'trial': legacy_trial, - 'hamiltonian': legacy_hamiltonian, - 'walkers': legacy_walkers, - 'propagator': legacy_propagator} + for _ in range(legacy_walkers[0].stack.ntime_slices): + for _, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker(legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = { + "system": legacy_system, + "trial": legacy_trial, + "hamiltonian": legacy_hamiltonian, + "walkers": legacy_walkers, + "propagator": legacy_propagator, + } return legacy_objs - + def build_legacy_driver_generic_test_instance( - hamiltonian, - comm, - nelec: Tuple[int, int], - mu: float, - beta: float, - timestep: float, - nblocks: int, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - alt_convention: bool = False, - sparse: bool = False, - seed: Union[int, None] = None, - estimator_filename: Union[str, None] = None, - verbose: int = 0): + hamiltonian, + comm, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + estimator_filename: Union[str, None] = None, + verbose: int = 0, +): nup, ndown = nelec numpy.random.seed(seed) - + legacy_options = { "qmc": { "dt": timestep, @@ -269,27 +290,16 @@ def build_legacy_driver_generic_test_instance( "pop_control_freq": pop_control_freq, "pop_control_method": pop_control_method, "rng_seed": seed, - "batched": False + "batched": False, }, - - "propagator": { - "optimised": False - }, - + "propagator": {"optimised": False}, "walkers": { "stack_size": stack_size, "low_rank": lowrank, "low_rank_thresh": lowrank_thresh, - "pop_control": pop_control_method + "pop_control": pop_control_method, }, - - "system": { - "name": "Generic", - "nup": nup, - "ndown": ndown, - "mu": mu - }, - + "system": {"name": "Generic", "nup": nup, "ndown": ndown, "mu": mu}, "estimators": { "filename": estimator_filename, }, @@ -298,42 +308,44 @@ def build_legacy_driver_generic_test_instance( legacy_system = Generic(nelec) legacy_system.mu = mu legacy_hamiltonian = LegacyHamGeneric( - h1e=hamiltonian.H1, - chol=hamiltonian.chol, - ecore=hamiltonian.ecore) + h1e=hamiltonian.H1, chol=hamiltonian.chol, ecore=hamiltonian.ecore + ) legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( - (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis) + ) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention legacy_hamiltonian.sparse = sparse legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, timestep) - - afqmc = LegacyThermalAFQMC(comm, legacy_options, legacy_system, - legacy_hamiltonian, legacy_trial, verbose=verbose) + + afqmc = LegacyThermalAFQMC( + comm, legacy_options, legacy_system, legacy_hamiltonian, legacy_trial, verbose=verbose + ) return afqmc def build_legacy_ueg_test_case_handlers( - comm, - nelec: Tuple[int, int], - rs: float, - ecut: float, - mu: float, - beta: float, - timestep: float, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - propagate: bool = False, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - alt_convention: bool = False, - sparse: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + comm, + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + propagate: bool = False, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): numpy.random.seed(seed) nup, ndown = nelec legacy_options = { @@ -344,35 +356,30 @@ def build_legacy_ueg_test_case_handlers( "ecut": ecut, "thermal": True, "write_integrals": False, - "low_rank": lowrank - }, - - "propagator": { - "optimised": False + "low_rank": lowrank, }, - + "propagator": {"optimised": False}, "walkers": { "stack_size": stack_size, "low_rank": lowrank, "low_rank_thresh": lowrank_thresh, - "pop_control": pop_control_method + "pop_control": pop_control_method, }, } # 1. Build out system. - legacy_system = LegacyUEG(options=legacy_options['ueg']) + legacy_system = LegacyUEG(options=legacy_options["ueg"]) legacy_system.mu = mu - + # 2. Build Hamiltonian. - legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options['ueg']) + legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options["ueg"]) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention # 3. Build trial. - legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) - - # 4. Build walkers. + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose) + + # 4. Build walkers. qmc_opts = ThermalQMCOpts() qmc_opts.nwalkers = nwalkers qmc_opts.ntot_walkers = nwalkers * comm.size @@ -384,50 +391,64 @@ def build_legacy_ueg_test_case_handlers( qmc_opts.pop_control_method = pop_control_method qmc_opts.seed = seed - legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, - qmc_opts, walker_opts=legacy_options['walkers'], - verbose=verbose, comm=comm) + legacy_walkers = Walkers( + legacy_system, + legacy_hamiltonian, + legacy_trial, + qmc_opts, + walker_opts=legacy_options["walkers"], + verbose=verbose, + comm=comm, + ) # 5. Build propagator. - legacy_propagator = PlaneWave(legacy_system, legacy_hamiltonian, legacy_trial, - qmc_opts, options=legacy_options["propagator"], - lowrank=lowrank, verbose=verbose) + legacy_propagator = PlaneWave( + legacy_system, + legacy_hamiltonian, + legacy_trial, + qmc_opts, + options=legacy_options["propagator"], + lowrank=lowrank, + verbose=verbose, + ) if propagate: - for t in range(legacy_walkers[0].stack.ntime_slices): - for iw, walker in enumerate(legacy_walkers): - legacy_propagator.propagate_walker( - legacy_hamiltonian, walker, legacy_trial) - - legacy_objs = {'system': legacy_system, - 'trial': legacy_trial, - 'hamiltonian': legacy_hamiltonian, - 'walkers': legacy_walkers, - 'propagator': legacy_propagator} + for _ in range(legacy_walkers[0].stack.ntime_slices): + for _, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker(legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = { + "system": legacy_system, + "trial": legacy_trial, + "hamiltonian": legacy_hamiltonian, + "walkers": legacy_walkers, + "propagator": legacy_propagator, + } return legacy_objs - + def build_legacy_driver_ueg_test_instance( - comm, - nelec: Tuple[int, int], - rs: float, - ecut: float, - mu: float, - beta: float, - timestep: float, - nblocks: int, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - alt_convention: bool = False, - sparse: bool = False, - seed: Union[int, None] = None, - estimator_filename: Union[str, None] = None, - verbose: int = 0): + comm, + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + estimator_filename: Union[str, None] = None, + verbose: int = 0, +): numpy.random.seed(seed) nup, ndown = nelec legacy_options = { @@ -438,9 +459,8 @@ def build_legacy_driver_ueg_test_instance( "ecut": ecut, "thermal": True, "write_integrals": False, - "low_rank": lowrank + "low_rank": lowrank, }, - "qmc": { "dt": timestep, # Input of `nwalkers` refers to the total number of walkers in @@ -453,40 +473,34 @@ def build_legacy_driver_ueg_test_instance( "pop_control_freq": pop_control_freq, "pop_control_method": pop_control_method, "rng_seed": seed, - "batched": False - }, - - "propagator": { - "optimised": False + "batched": False, }, - + "propagator": {"optimised": False}, "walkers": { "stack_size": stack_size, "low_rank": lowrank, "low_rank_thresh": lowrank_thresh, - "pop_control": pop_control_method + "pop_control": pop_control_method, }, - "estimators": { "filename": estimator_filename, }, } # 1. Build out system. - legacy_system = LegacyUEG(options=legacy_options['ueg']) + legacy_system = LegacyUEG(options=legacy_options["ueg"]) legacy_system.mu = mu - + # 2. Build Hamiltonian. - legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options['ueg']) + legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options["ueg"]) legacy_hamiltonian.mu = mu legacy_hamiltonian._alt_convention = alt_convention # 3. Build trial. - legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, - timestep, verbose=verbose) - + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, timestep, verbose=verbose) + # 4. Build Thermal AFQMC. - afqmc = LegacyThermalAFQMC(comm, legacy_options, legacy_system, - legacy_hamiltonian, legacy_trial, verbose=verbose) + afqmc = LegacyThermalAFQMC( + comm, legacy_options, legacy_system, legacy_hamiltonian, legacy_trial, verbose=verbose + ) return afqmc - diff --git a/ipie/addons/thermal/utils/testing.py b/ipie/addons/thermal/utils/testing.py index 0e25b695..f828c155 100644 --- a/ipie/addons/thermal/utils/testing.py +++ b/ipie/addons/thermal/utils/testing.py @@ -16,55 +16,58 @@ # Joonho Lee # -import numpy from typing import Tuple, Union -from ipie.utils.mpi import MPIHandler -from ipie.utils.testing import generate_hamiltonian -from ipie.hamiltonians.generic import Generic as HamGeneric +import numpy -from ipie.addons.thermal.utils.ueg import UEG -from ipie.addons.thermal.trial.one_body import OneBody -from ipie.addons.thermal.trial.mean_field import MeanField -from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers from ipie.addons.thermal.propagation.phaseless_generic import PhaselessGeneric from ipie.addons.thermal.qmc.thermal_afqmc import ThermalAFQMC +from ipie.addons.thermal.trial.mean_field import MeanField +from ipie.addons.thermal.trial.one_body import OneBody +from ipie.addons.thermal.utils.ueg import UEG +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.utils.mpi import MPIHandler +from ipie.utils.testing import generate_hamiltonian def build_generic_test_case_handlers( - nelec: Tuple[int, int], - nbasis: int, - mu: float, - beta: float, - timestep: float, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - diagonal: bool = False, - mf_trial: bool = True, - propagate: bool = False, - complex_integrals: bool = False, - debug: bool = False, - with_eri: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + diagonal: bool = False, + mf_trial: bool = True, + propagate: bool = False, + complex_integrals: bool = False, + debug: bool = False, + with_eri: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): sym = 8 - if complex_integrals: sym = 4 + if complex_integrals: + sym = 4 numpy.random.seed(seed) - + # 1. Generate random integrals. - h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, - sym=sym, tol=1e-10) + h1e, chol, _, eri = generate_hamiltonian( + nbasis, nelec, cplx=complex_integrals, sym=sym, tol=1e-10 + ) if diagonal: h1e = numpy.diag(numpy.diag(h1e)) - + # 2. Build Hamiltonian. - hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), - chol=chol.reshape((-1, nbasis**2)).T.copy(), - ecore=0) - + hamiltonian = HamGeneric( + h1e=numpy.array([h1e, h1e]), chol=chol.reshape((-1, nbasis**2)).T.copy(), ecore=0 + ) + # 3. Build trial. trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) @@ -73,62 +76,73 @@ def build_generic_test_case_handlers( # 4. Build walkers. walkers = UHFThermalWalkers( - trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, - lowrank_thresh=lowrank_thresh, verbose=verbose) - + trial, + nbasis, + nwalkers, + stack_size=stack_size, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + verbose=verbose, + ) + # 5. Build propagator. propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) propagator.build(hamiltonian, trial=trial, walkers=walkers, verbose=verbose) if propagate: - for t in range(walkers.stack[0].nslice): + for _ in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) - objs = {'trial': trial, - 'hamiltonian': hamiltonian, - 'walkers': walkers, - 'propagator': propagator} + objs = { + "trial": trial, + "hamiltonian": hamiltonian, + "walkers": walkers, + "propagator": propagator, + } if with_eri: - objs['eri'] = eri + objs["eri"] = eri return objs def build_generic_test_case_handlers_mpi( - nelec: Tuple[int, int], - nbasis: int, - mu: float, - beta: float, - timestep: float, - mpi_handler: MPIHandler, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - diagonal: bool = False, - mf_trial: bool = True, - propagate: bool = False, - complex_integrals: bool = False, - debug: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + mpi_handler: MPIHandler, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + diagonal: bool = False, + mf_trial: bool = True, + propagate: bool = False, + complex_integrals: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): sym = 8 - if complex_integrals: sym = 4 + if complex_integrals: + sym = 4 numpy.random.seed(seed) - + # 1. Generate random integrals. - h1e, chol, _, _ = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, - sym=sym, tol=1e-10) + h1e, chol, _, _ = generate_hamiltonian( + nbasis, nelec, cplx=complex_integrals, sym=sym, tol=1e-10 + ) if diagonal: h1e = numpy.diag(numpy.diag(h1e)) - + # 2. Build Hamiltonian. - hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), - chol=chol.reshape((-1, nbasis**2)).T.copy(), - ecore=0) - + hamiltonian = HamGeneric( + h1e=numpy.array([h1e, h1e]), chol=chol.reshape((-1, nbasis**2)).T.copy(), ecore=0 + ) + # 3. Build trial. trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) @@ -137,98 +151,125 @@ def build_generic_test_case_handlers_mpi( # 4. Build walkers. walkers = UHFThermalWalkers( - trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, - lowrank_thresh=lowrank_thresh, mpi_handler=mpi_handler, verbose=verbose) - + trial, + nbasis, + nwalkers, + stack_size=stack_size, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + mpi_handler=mpi_handler, + verbose=verbose, + ) + # 5. Build propagator. propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) - propagator.build(hamiltonian, trial=trial, walkers=walkers, - mpi_handler=mpi_handler, verbose=verbose) + propagator.build( + hamiltonian, trial=trial, walkers=walkers, mpi_handler=mpi_handler, verbose=verbose + ) if propagate: - for t in range(walkers.stack[0].nslice): + for _ in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) - objs = {'trial': trial, - 'hamiltonian': hamiltonian, - 'walkers': walkers, - 'propagator': propagator} + objs = { + "trial": trial, + "hamiltonian": hamiltonian, + "walkers": walkers, + "propagator": propagator, + } return objs def build_driver_generic_test_instance( - nelec: Tuple[int, int], - nbasis: int, - mu: float, - beta: float, - timestep: float, - nblocks: int, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - diagonal: bool = False, - complex_integrals: bool = False, - debug: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + diagonal: bool = False, + complex_integrals: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): sym = 8 - if complex_integrals: sym = 4 + if complex_integrals: + sym = 4 numpy.random.seed(seed) # 1. Generate random integrals. - h1e, chol, _, _ = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, - sym=sym, tol=1e-10) - + h1e, chol, _, _ = generate_hamiltonian( + nbasis, nelec, cplx=complex_integrals, sym=sym, tol=1e-10 + ) + if diagonal: h1e = numpy.diag(numpy.diag(h1e)) - + # 2. Build Hamiltonian. - hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), - chol=chol.reshape((-1, nbasis**2)).T.copy(), - ecore=0) - + hamiltonian = HamGeneric( + h1e=numpy.array([h1e, h1e]), chol=chol.reshape((-1, nbasis**2)).T.copy(), ecore=0 + ) + # 3. Build trial. trial = MeanField(hamiltonian, nelec, beta, timestep) # 4. Build Thermal AFQMC driver. afqmc = ThermalAFQMC.build( - nelec, mu, beta, hamiltonian, trial, nwalkers=nwalkers, - stack_size=stack_size, seed=seed, nblocks=nblocks, timestep=timestep, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, lowrank=lowrank, - lowrank_thresh=lowrank_thresh, debug=debug, verbose=verbose) + nelec, + mu, + beta, + hamiltonian, + trial, + nwalkers=nwalkers, + stack_size=stack_size, + seed=seed, + nblocks=nblocks, + timestep=timestep, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + debug=debug, + verbose=verbose, + ) return afqmc def build_ueg_test_case_handlers( - nelec: Tuple[int, int], - rs: float, - ecut: float, - mu: float, - beta: float, - timestep: float, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - propagate: bool = False, - debug: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + propagate: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): nup, ndown = nelec ueg_opts = { - "nup": nup, - "ndown": ndown, - "rs": rs, - "ecut": ecut, - "thermal": True, - "write_integrals": False, - "low_rank": lowrank - } + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank, + } numpy.random.seed(seed) @@ -243,69 +284,79 @@ def build_ueg_test_case_handlers( print(f"# nchol = {nchol}") print(f"# nup = {nup}") print(f"# ndown = {ndown}") - + h1 = ueg.H1[0] - chol = 2. * ueg.chol_vecs.toarray().copy() - ecore = 0. + chol = 2.0 * ueg.chol_vecs.toarray().copy() + ecore = 0.0 # 2. Build Hamiltonian. hamiltonian = HamGeneric( - numpy.array([h1, h1], dtype=numpy.complex128), - numpy.array(chol, dtype=numpy.complex128), - ecore, - verbose=verbose) + numpy.array([h1, h1], dtype=numpy.complex128), + numpy.array(chol, dtype=numpy.complex128), + ecore, + verbose=verbose, + ) # 3. Build trial. trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) # 4. Build walkers. walkers = UHFThermalWalkers( - trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, - lowrank_thresh=lowrank_thresh, verbose=verbose) - + trial, + nbasis, + nwalkers, + stack_size=stack_size, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + verbose=verbose, + ) + # 5. Build propagator. propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) propagator.build(hamiltonian, trial=trial, walkers=walkers, verbose=verbose) if propagate: - for t in range(walkers.stack[0].nslice): + for _ in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) - objs = {'trial': trial, - 'hamiltonian': hamiltonian, - 'walkers': walkers, - 'propagator': propagator} + objs = { + "trial": trial, + "hamiltonian": hamiltonian, + "walkers": walkers, + "propagator": propagator, + } return objs def build_driver_ueg_test_instance( - nelec: Tuple[int, int], - rs: float, - ecut: float, - mu: float, - beta: float, - timestep: float, - nblocks: int, - nwalkers: int = 100, - stack_size: int = 10, - lowrank: bool = False, - lowrank_thresh: float = 1e-6, - stabilize_freq: int = 5, - pop_control_freq: int = 5, - pop_control_method: str = 'pair_branch', - debug: bool = False, - seed: Union[int, None] = None, - verbose: int = 0): + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = "pair_branch", + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0, +): nup, ndown = nelec ueg_opts = { - "nup": nup, - "ndown": ndown, - "rs": rs, - "ecut": ecut, - "thermal": True, - "write_integrals": False, - "low_rank": lowrank - } + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank, + } numpy.random.seed(seed) @@ -320,27 +371,40 @@ def build_driver_ueg_test_instance( print(f"# nchol = {nchol}") print(f"# nup = {nup}") print(f"# ndown = {ndown}") - + h1 = ueg.H1[0] - chol = 2. * ueg.chol_vecs.toarray().copy() - ecore = 0. + chol = 2.0 * ueg.chol_vecs.toarray().copy() + ecore = 0.0 # 2. Build Hamiltonian. hamiltonian = HamGeneric( - numpy.array([h1, h1], dtype=numpy.complex128), - numpy.array(chol, dtype=numpy.complex128), - ecore, - verbose=verbose) + numpy.array([h1, h1], dtype=numpy.complex128), + numpy.array(chol, dtype=numpy.complex128), + ecore, + verbose=verbose, + ) # 3. Build trial. trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) # 4. Build Thermal AFQMC driver. afqmc = ThermalAFQMC.build( - nelec, mu, beta, hamiltonian, trial, nwalkers=nwalkers, - stack_size=stack_size, seed=seed, nblocks=nblocks, timestep=timestep, - stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, - pop_control_method=pop_control_method, lowrank=lowrank, - lowrank_thresh=lowrank_thresh, debug=debug, verbose=verbose) + nelec, + mu, + beta, + hamiltonian, + trial, + nwalkers=nwalkers, + stack_size=stack_size, + seed=seed, + nblocks=nblocks, + timestep=timestep, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + lowrank=lowrank, + lowrank_thresh=lowrank_thresh, + debug=debug, + verbose=verbose, + ) return afqmc - diff --git a/ipie/addons/thermal/utils/ueg.py b/ipie/addons/thermal/utils/ueg.py index 43286bfa..a1a9e1e3 100644 --- a/ipie/addons/thermal/utils/ueg.py +++ b/ipie/addons/thermal/utils/ueg.py @@ -20,6 +20,7 @@ import scipy.sparse from ipie.utils.io import write_qmcpack_sparse + class UEG: """UEG system class (integrals read from fcidump) @@ -79,7 +80,7 @@ def __init__(self, options, verbose=False): self.thermal = options.get("thermal", False) self._alt_convention = options.get("alt_convention", False) self.write_ints = options.get("write_integrals", False) - + self.sparse = True self.control_variate = False self.diagH1 = True @@ -116,11 +117,9 @@ def __init__(self, options, verbose=False): print(f"# Volume: {self.vol:13.8e}") print(f"# k-space factor (2pi/L): {self.kfac:13.8e}") - def build(self, verbose=False): # Get plane wave basis vectors and corresponding eigenvalues. - self.sp_eigv, self.basis, self.nmax = self.sp_energies( - self.ktwist, self.kfac, self.ecut) + self.sp_eigv, self.basis, self.nmax = self.sp_energies(self.ktwist, self.kfac, self.ecut) self.shifted_nmax = 2 * self.nmax self.imax_sq = numpy.dot(self.basis[-1], self.basis[-1]) self.create_lookup_table() @@ -134,11 +133,11 @@ def build(self, verbose=False): self.ncore = 0 self.nfv = 0 self.mo_coeff = None - + # --------------------------------------------------------------------- T = numpy.diag(self.sp_eigv) h1e_mod = self.mod_one_body(T) - self.H1 = numpy.array([T, T]) # Making alpha and beta. + self.H1 = numpy.array([T, T]) # Making alpha and beta. self.h1e_mod = numpy.array([h1e_mod, h1e_mod]) # --------------------------------------------------------------------- @@ -168,12 +167,13 @@ def build(self, verbose=False): self.write_integrals() if verbose: - print("# Approximate memory required for " - "two-body potentials: {:13.8e} GB.".format((3 * self.iA.nnz * 16 / (1024**3)))) + print( + "# Approximate memory required for " + "two-body potentials: {:13.8e} GB.".format((3 * self.iA.nnz * 16 / (1024**3))) + ) print("# Finished constructing two-body potentials.") print("# Finished building UEG object.") - def sp_energies(self, ks, kfac, ecut): """Calculate the allowed kvectors and resulting single particle eigenvalues (basically kinetic energy) which can fit in the sphere in kspace determined by ecut. @@ -223,7 +223,6 @@ def sp_energies(self, ks, kfac, ecut): kval = numpy.array(kval)[ix] return spval, kval, nmax - def create_lookup_table(self): basis_ix = [] for k in self.basis: @@ -236,7 +235,6 @@ def create_lookup_table(self): self.max_ix = max(basis_ix) - def lookup_basis(self, vec): if numpy.dot(vec, vec) <= self.imax_sq: ix = self.map_basis_to_index(vec) @@ -252,12 +250,12 @@ def lookup_basis(self, vec): else: ib = None - def map_basis_to_index(self, k): - return ((k[0] + self.nmax) - + self.shifted_nmax * (k[1] + self.nmax) - + self.shifted_nmax * self.shifted_nmax * (k[2] + self.nmax)) - + return ( + (k[0] + self.nmax) + + self.shifted_nmax * (k[1] + self.nmax) + + self.shifted_nmax * self.shifted_nmax * (k[2] + self.nmax) + ) def get_momentum_transfers(self): """Get arrays of plane wave basis vectors connected by momentum transfers Q.""" @@ -282,7 +280,7 @@ def get_momentum_transfers(self): self.ikpq_i += [idxkpq_list_i] self.ikpq_kpq += [idxkpq_list_kpq] - + # --------------------------------------------------------------------- self.ipmq_i = [] self.ipmq_pmq = [] @@ -308,7 +306,6 @@ def get_momentum_transfers(self): self.ipmq_i[iq] = numpy.array(self.ipmq_i[iq], dtype=numpy.int64) self.ipmq_pmq[iq] = numpy.array(self.ipmq_pmq[iq], dtype=numpy.int64) - def madelung(self): """Use expression in Schoof et al. (PhysRevLett.115.130402) for the Madelung contribution to the total energy fitted to L.M. Fraser et al. @@ -331,7 +328,6 @@ def madelung(self): c2 = (3.0 / (4.0 * numpy.pi)) ** (1.0 / 3.0) return c1 * c2 / (self.ne ** (1.0 / 3.0) * self.rs) - def mod_one_body(self, T): """Absorb the diagonal term of the two-body Hamiltonian to the one-body term. Essentially adding the third term in Eq.(11b) of Phys. Rev. B 75, 245123. @@ -357,7 +353,6 @@ def mod_one_body(self, T): return h1e_mod - def vq(self, q): """The typical 3D Coulomb kernel @@ -373,7 +368,6 @@ def vq(self, q): """ return 4 * numpy.pi / numpy.dot(q, q) - def scaled_density_operator_incore(self, transpose): """Density operator as defined in Eq.(6) of PRB(75)245123 @@ -445,10 +439,10 @@ def scaled_density_operator_incore(self, transpose): rho_q = scipy.sparse.csc_matrix( (values, (row_index, col_index)), shape=(self.nbasis * self.nbasis, nq), - dtype=numpy.complex128) + dtype=numpy.complex128, + ) return rho_q - def two_body_potentials_incore(self): """Calculate A and B of Eq.(13) of PRB(75)245123 for a given plane-wave vector q @@ -466,7 +460,6 @@ def two_body_potentials_incore(self): iB = -(rho_q - rho_qH) return (rho_q, iA, iB) - def hijkl(self, i, j, k, l): """Compute = (ik|jl) = 1/Omega * 4pi/(kk-ki)**2 @@ -492,7 +485,6 @@ def hijkl(self, i, j, k, l): else: return 0.0 - def compute_real_transformation(self): U22 = numpy.zeros((2, 2), dtype=numpy.complex128) U22[0, 0] = 1.0 / numpy.sqrt(2.0) @@ -525,17 +517,16 @@ def compute_real_transformation(self): U = U.T.copy() return U - def eri_4(self): eri_chol = 4 * self.chol_vecs.dot(self.chol_vecs.T) eri_chol = ( - eri_chol.toarray().reshape((self.nbasis, self.nbasis, self.nbasis, self.nbasis)).real) + eri_chol.toarray().reshape((self.nbasis, self.nbasis, self.nbasis, self.nbasis)).real + ) eri_chol = eri_chol.transpose(0, 1, 3, 2) return eri_chol - def eri_8(self): - """Compute 8-fold symmetric integrals. Useful for running standard + """Compute 8-fold symmetric integrals. Useful for running standard quantum chemistry methods,""" eri = self.eri_4() U = self.compute_real_transformation() @@ -544,7 +535,6 @@ def eri_8(self): eri2 = numpy.einsum("lr,pqls->pqrs", U.conj(), eri1, optimize=True) eri3 = numpy.einsum("st,pqrs->pqrt", U, eri2, optimize=True).real return eri3 - def write_integrals(self, filename="ueg_integrals.h5"): write_qmcpack_sparse( @@ -553,5 +543,5 @@ def write_integrals(self, filename="ueg_integrals.h5"): self.nelec, self.nbasis, enuc=0.0, - filename=filename) - + filename=filename, + ) diff --git a/ipie/addons/thermal/walkers/pop_controller.py b/ipie/addons/thermal/walkers/pop_controller.py index a0aed19a..8725044f 100644 --- a/ipie/addons/thermal/walkers/pop_controller.py +++ b/ipie/addons/thermal/walkers/pop_controller.py @@ -36,8 +36,15 @@ def __init__( verbose=False, ): super().__init__( - num_walkers_local, num_steps, mpi_handler, pop_control_method, - min_weight, max_weight, reconfiguration_freq, verbose) + num_walkers_local, + num_steps, + mpi_handler, + pop_control_method, + min_weight, + max_weight, + reconfiguration_freq, + verbose, + ) def pop_control(self, walkers, comm): self.timer.start_time() @@ -102,7 +109,9 @@ def get_buffer(walkers, iw): buff = xp.zeros(walkers.buff_size, dtype=numpy.complex128) for d in walkers.buff_names: data = walkers.__dict__[d] - if (data is None) or isinstance(data, (int, float, complex, numpy.float64, numpy.complex128)): + if (data is None) or isinstance( + data, (int, float, complex, numpy.float64, numpy.complex128) + ): continue assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated if isinstance(data[iw], (xp.ndarray)): @@ -135,18 +144,22 @@ def set_buffer(walkers, iw, buff): s = 0 for d in walkers.buff_names: data = walkers.__dict__[d] - if (data is None) or isinstance(data, (int, float, complex, numpy.float64, numpy.complex128)): + if (data is None) or isinstance( + data, (int, float, complex, numpy.float64, numpy.complex128) + ): continue assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated if isinstance(data[iw], xp.ndarray): walkers.__dict__[d][iw] = xp.array( - buff[s : s + data[iw].size].reshape(data[iw].shape).copy()) + buff[s : s + data[iw].size].reshape(data[iw].shape).copy() + ) s += data[iw].size elif isinstance(data[iw], list): for ix, l in enumerate(data[iw]): if isinstance(l, (xp.ndarray)): walkers.__dict__[d][iw][ix] = xp.array( - buff[s : s + l.size].reshape(l.shape).copy()) + buff[s : s + l.size].reshape(l.shape).copy() + ) s += l.size elif isinstance(l, (int, float, complex)): walkers.__dict__[d][iw][ix] = buff[s] @@ -159,8 +172,8 @@ def set_buffer(walkers, iw, buff): else: walkers.__dict__[d][iw] = buff[s] s += 1 - - walkers.stack[iw].set_buffer(buff[walkers.buff_size:]) + + walkers.stack[iw].set_buffer(buff[walkers.buff_size :]) def comb(walkers, comm, weights, target_weight, timer=PopControllerTimer()): @@ -286,11 +299,11 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( glob_inf_1 = numpy.empty([comm.size, walkers.nwalkers], dtype=numpy.int64) glob_inf_1.fill(1) glob_inf_2 = numpy.array( - [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64) + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], dtype=numpy.int64 + ) glob_inf_3 = numpy.array( - [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64) + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], dtype=numpy.int64 + ) timer.add_non_communication() @@ -308,9 +321,15 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( # Rescale weights. glob_inf = numpy.zeros((walkers.nwalkers * comm.size, 4), dtype=numpy.float64) glob_inf[:, 0] = glob_inf_0.ravel() # contains walker |w_i| - glob_inf[:, 1] = glob_inf_1.ravel() # all initialized to 1 when it becomes 2 then it will be "branched" - glob_inf[:, 2] = glob_inf_2.ravel() # contain processor+walker indices (initial) (i.e., where walkers live) - glob_inf[:, 3] = glob_inf_3.ravel() # contain processor+walker indices (final) (i.e., where walkers live) + glob_inf[:, 1] = ( + glob_inf_1.ravel() + ) # all initialized to 1 when it becomes 2 then it will be "branched" + glob_inf[:, 2] = ( + glob_inf_2.ravel() + ) # contain processor+walker indices (initial) (i.e., where walkers live) + glob_inf[:, 3] = ( + glob_inf_3.ravel() + ) # contain processor+walker indices (final) (i.e., where walkers live) sort = numpy.argsort(glob_inf[:, 0], kind="mergesort") isort = numpy.argsort(sort, kind="mergesort") glob_inf = glob_inf[sort] diff --git a/ipie/addons/thermal/walkers/stack.py b/ipie/addons/thermal/walkers/stack.py index af048b77..4476cd34 100644 --- a/ipie/addons/thermal/walkers/stack.py +++ b/ipie/addons/thermal/walkers/stack.py @@ -24,6 +24,7 @@ from ipie.utils.misc import get_numeric_names + class PropagatorStack: def __init__( self, @@ -67,8 +68,9 @@ def __init__( self.left = numpy.zeros((self.nstack, 2, nbasis, nbasis), dtype=dtype) self.right = numpy.zeros((self.nstack, 2, nbasis, nbasis), dtype=dtype) - self.G = numpy.asarray([numpy.eye(self.nbasis, dtype=dtype), # Ga - numpy.eye(self.nbasis, dtype=dtype)]) # Gb + self.G = numpy.asarray( + [numpy.eye(self.nbasis, dtype=dtype), numpy.eye(self.nbasis, dtype=dtype)] # Ga + ) # Gb if self.lowrank: self.update_new = self.update_low_rank @@ -329,8 +331,9 @@ def update_low_rank(self, B): self.theta[s][:, :] = 0.0 self.theta[s][:mT, :] = Qlcr_pad[:, :mT].dot(numpy.diag(Dlcr[:mT])).T # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.CT[s][:,:mT].dot(self.theta[s][:mT,:]) - self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) -\ - self.theta[s][:mT, :].T.dot(self.CT[s][:, :mT].T.conj()) + self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.theta[s][:mT, :].T.dot( + self.CT[s][:, :mT].T.conj() + ) # self.CT[s][:,:mT] = self.CT[s][:,:mT].conj() # print("# mL, mR, mT = {}, {}, {}".format(mL, mR, mT)) @@ -397,8 +400,9 @@ def update_low_rank(self, B): self.theta[s][:, :] = 0.0 self.theta[s][:mT, :] = Qlcr_pad[:, :mT].dot(numpy.diag(Dlcr[:mT])).T # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.CT[s][:,:mT].dot(self.theta[s][:mT,:]) - self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) -\ - self.theta[s][:mT, :].T.dot(self.CT[s][:, :mT].T.conj()) + self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.theta[s][:mT, :].T.dot( + self.CT[s][:, :mT].T.conj() + ) # self.CT = numpy.zeros(shape=(2, nbasis, nbasis),dtype=dtype) # self.theta = numpy.zeros(shape=(2, nbasis, nbasis),dtype=dtype) diff --git a/ipie/addons/thermal/walkers/tests/__init__.py b/ipie/addons/thermal/walkers/tests/__init__.py index 381108b3..e2aed039 100644 --- a/ipie/addons/thermal/walkers/tests/__init__.py +++ b/ipie/addons/thermal/walkers/tests/__init__.py @@ -11,4 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - diff --git a/ipie/addons/thermal/walkers/tests/test_population_control.py b/ipie/addons/thermal/walkers/tests/test_population_control.py index 9250e5d5..2dd6edf8 100644 --- a/ipie/addons/thermal/walkers/tests/test_population_control.py +++ b/ipie/addons/thermal/walkers/tests/test_population_control.py @@ -23,6 +23,7 @@ try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers_mpi from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False except ModuleNotFoundError: @@ -35,6 +36,7 @@ comm = MPI.COMM_WORLD + @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @pytest.mark.unit def test_pair_branch_batch(): @@ -47,14 +49,14 @@ def test_pair_branch_batch(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 nblocks = 3 # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! nsteps_per_block = 1 - pop_control_method = 'pair_branch' + pop_control_method = "pair_branch" lowrank = False mf_trial = True @@ -65,51 +67,79 @@ def test_pair_branch_batch(): numpy.random.seed(seed) # Test. - objs = build_generic_test_case_handlers_mpi( - nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, - debug=debug, seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] - pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, - pop_control_method, verbose=verbose) - + objs = build_generic_test_case_handlers_mpi( + nelec, + nbasis, + mu, + beta, + timestep, + mpi_handler, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] + pcontrol = ThermalPopController( + nwalkers, nsteps_per_block, mpi_handler, pop_control_method, verbose=verbose + ) + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers_mpi( - hamiltonian, mpi_handler, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] - - for block in range(nblocks): + hamiltonian, + mpi_handler, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] + + for block in range(nblocks): for t in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) - + legacy_hamiltonian, + legacy_trial, + legacy_walkers, + legacy_propagator, + xi=propagator.xi, + ) + if t > 0: pcontrol.pop_control(walkers, mpi_handler.comm) legacy_walkers.pop_control(mpi_handler.comm) - walkers.reset(trial) # Reset stack, weights, phase. + walkers.reset(trial) # Reset stack, weights, phase. legacy_walkers.reset(legacy_trial) for iw in range(walkers.nwalkers): assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) - assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + assert numpy.allclose( + walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight + ) + # TODO: Lowrank code is WIP. See: https://github.com/JoonhoLee-Group/ipie/issues/302 -#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") -#@pytest.mark.unit +# @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +# @pytest.mark.unit def test_pair_branch_batch_lowrank(): mpi_handler = MPIHandler() @@ -120,14 +150,14 @@ def test_pair_branch_batch_lowrank(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 nblocks = 3 # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! nsteps_per_block = 1 - pop_control_method = 'pair_branch' + pop_control_method = "pair_branch" lowrank = True mf_trial = False @@ -139,67 +169,99 @@ def test_pair_branch_batch_lowrank(): numpy.random.seed(seed) options = { - 'nelec': nelec, - 'nbasis': nbasis, - 'mu': mu, - 'beta': beta, - 'timestep': timestep, - 'nwalkers': nwalkers, - 'seed': seed, - 'nsteps_per_block': nsteps_per_block, - 'nblocks': nblocks, - 'stabilize_freq': stabilize_freq, - 'pop_control_freq': pop_control_freq, - 'pop_control_method': pop_control_method, - 'lowrank': lowrank, - 'complex_integrals': complex_integrals, - 'mf_trial': mf_trial, - 'propagate': propagate, - 'diagonal': diagonal, - } - + "nelec": nelec, + "nbasis": nbasis, + "mu": mu, + "beta": beta, + "timestep": timestep, + "nwalkers": nwalkers, + "seed": seed, + "nsteps_per_block": nsteps_per_block, + "nblocks": nblocks, + "stabilize_freq": stabilize_freq, + "pop_control_freq": pop_control_freq, + "pop_control_method": pop_control_method, + "lowrank": lowrank, + "complex_integrals": complex_integrals, + "mf_trial": mf_trial, + "propagate": propagate, + "diagonal": diagonal, + } + # Test. - objs = build_generic_test_case_handlers_mpi( - nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, - diagonal=diagonal, debug=debug, seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] - pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, - pop_control_method=pop_control_method, verbose=verbose) - + objs = build_generic_test_case_handlers_mpi( + nelec, + nbasis, + mu, + beta, + timestep, + mpi_handler, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + diagonal=diagonal, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] + pcontrol = ThermalPopController( + nwalkers, + nsteps_per_block, + mpi_handler, + pop_control_method=pop_control_method, + verbose=verbose, + ) + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers_mpi( - hamiltonian, mpi_handler, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] - - for block in range(nblocks): + hamiltonian, + mpi_handler, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] + + for block in range(nblocks): for t in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) - + legacy_hamiltonian, + legacy_trial, + legacy_walkers, + legacy_propagator, + xi=propagator.xi, + ) + if t > 0: pcontrol.pop_control(walkers, mpi_handler.comm) legacy_walkers.pop_control(mpi_handler.comm) - walkers.reset(trial) # Reset stack, weights, phase. + walkers.reset(trial) # Reset stack, weights, phase. legacy_walkers.reset(legacy_trial) for iw in range(walkers.nwalkers): assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) - assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + assert numpy.allclose( + walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight + ) @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @@ -213,14 +275,14 @@ def test_comb_batch(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 nblocks = 3 # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! nsteps_per_block = 1 - pop_control_method = 'comb' + pop_control_method = "comb" lowrank = False mf_trial = True @@ -231,56 +293,87 @@ def test_comb_batch(): numpy.random.seed(seed) # Test. - objs = build_generic_test_case_handlers_mpi( - nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, - debug=debug, seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] - pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, - pop_control_method=pop_control_method, verbose=verbose) - + objs = build_generic_test_case_handlers_mpi( + nelec, + nbasis, + mu, + beta, + timestep, + mpi_handler, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] + pcontrol = ThermalPopController( + nwalkers, + nsteps_per_block, + mpi_handler, + pop_control_method=pop_control_method, + verbose=verbose, + ) + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers_mpi( - hamiltonian, mpi_handler, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - pop_control_method=pop_control_method, seed=seed, - verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] - - for block in range(nblocks): + hamiltonian, + mpi_handler, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + pop_control_method=pop_control_method, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] + + for block in range(nblocks): for t in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) - + legacy_hamiltonian, + legacy_trial, + legacy_walkers, + legacy_propagator, + xi=propagator.xi, + ) + if t > 0: pcontrol.pop_control(walkers, mpi_handler.comm) legacy_walkers.pop_control(mpi_handler.comm) - walkers.reset(trial) # Reset stack, weights, phase. + walkers.reset(trial) # Reset stack, weights, phase. legacy_walkers.reset(legacy_trial) for iw in range(walkers.nwalkers): assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) - assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + assert numpy.allclose( + walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight + ) # TODO: Lowrank code is WIP. See: https://github.com/JoonhoLee-Group/ipie/issues/302 -#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") -#@pytest.mark.unit +# @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +# @pytest.mark.unit def test_comb_batch_lowrank(): mpi_handler = MPIHandler() - + # System params. nup = 5 ndown = 5 @@ -288,14 +381,14 @@ def test_comb_batch_lowrank(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 12 nblocks = 3 # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! nsteps_per_block = 1 - pop_control_method = 'comb' + pop_control_method = "comb" lowrank = True mf_trial = False @@ -305,54 +398,86 @@ def test_comb_batch_lowrank(): verbose = False if (comm.rank != 0) else True seed = 7 numpy.random.seed(seed) - + # Test. - objs = build_generic_test_case_handlers_mpi( - nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, - diagonal=diagonal, debug=debug, seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - propagator = objs['propagator'] - pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, - pop_control_method=pop_control_method, verbose=verbose) - + objs = build_generic_test_case_handlers_mpi( + nelec, + nbasis, + mu, + beta, + timestep, + mpi_handler, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + diagonal=diagonal, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + propagator = objs["propagator"] + pcontrol = ThermalPopController( + nwalkers, + nsteps_per_block, + mpi_handler, + pop_control_method=pop_control_method, + verbose=verbose, + ) + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers_mpi( - hamiltonian, mpi_handler, nelec, mu, beta, timestep, - nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, - seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - legacy_propagator = legacy_objs['propagator'] - - for block in range(nblocks): + hamiltonian, + mpi_handler, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + legacy_propagator = legacy_objs["propagator"] + + for block in range(nblocks): for t in range(walkers.stack[0].nslice): propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) legacy_walkers = legacy_propagate_walkers( - legacy_hamiltonian, legacy_trial, legacy_walkers, - legacy_propagator, xi=propagator.xi) - + legacy_hamiltonian, + legacy_trial, + legacy_walkers, + legacy_propagator, + xi=propagator.xi, + ) + if t > 0: pcontrol.pop_control(walkers, mpi_handler.comm) legacy_walkers.pop_control(mpi_handler.comm) - walkers.reset(trial) # Reset stack, weights, phase. + walkers.reset(trial) # Reset stack, weights, phase. legacy_walkers.reset(legacy_trial) for iw in range(walkers.nwalkers): assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) - assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + assert numpy.allclose( + walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight + ) if __name__ == "__main__": test_pair_branch_batch() test_comb_batch() - - #test_pair_branch_batch_lowrank() - #test_comb_batch_lowrank() + + # test_pair_branch_batch_lowrank() + # test_comb_batch_lowrank() diff --git a/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py b/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py index a02da99c..f082e670 100644 --- a/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py +++ b/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py @@ -23,6 +23,7 @@ try: from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False except ModuleNotFoundError: @@ -33,11 +34,14 @@ from ipie.addons.thermal.estimators.thermal import one_rdm_from_G from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers -from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky +from ipie.legacy.estimators.generic import ( + local_energy_generic_cholesky as legacy_local_energy_generic_cholesky, +) from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G comm = MPI.COMM_WORLD + @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") @pytest.mark.unit def test_thermal_walkers_fullrank(): @@ -48,7 +52,7 @@ def test_thermal_walkers_fullrank(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 10 @@ -60,43 +64,72 @@ def test_thermal_walkers_fullrank(): verbose = True seed = 7 numpy.random.seed(seed) - + # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) eloc = local_energy_generic_cholesky(hamiltonian, P) legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) legacy_eloc = legacy_local_energy_generic_cholesky( - legacy_system, legacy_hamiltonian, legacy_P) + legacy_system, legacy_hamiltonian, legacy_P + ) numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10 + ) # TODO: Lowrank code is WIP. -#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") -#@pytest.mark.unit +# @pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +# @pytest.mark.unit def test_thermal_walkers_lowrank(): # System params. nup = 5 @@ -105,7 +138,7 @@ def test_thermal_walkers_lowrank(): nbasis = 10 # Thermal AFQMC params. - mu = -10. + mu = -10.0 beta = 0.1 timestep = 0.01 nwalkers = 10 @@ -118,39 +151,69 @@ def test_thermal_walkers_lowrank(): verbose = True seed = 7 numpy.random.seed(seed) - + # Test. - objs = build_generic_test_case_handlers( - nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, - mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, - seed=seed, verbose=verbose) - trial = objs['trial'] - hamiltonian = objs['hamiltonian'] - walkers = objs['walkers'] - + objs = build_generic_test_case_handlers( + nelec, + nbasis, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + complex_integrals=complex_integrals, + debug=debug, + seed=seed, + verbose=verbose, + ) + trial = objs["trial"] + hamiltonian = objs["hamiltonian"] + walkers = objs["walkers"] + # Legacy. legacy_objs = build_legacy_generic_test_case_handlers( - hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, - lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) - legacy_system = legacy_objs['system'] - legacy_trial = legacy_objs['trial'] - legacy_hamiltonian = legacy_objs['hamiltonian'] - legacy_walkers = legacy_objs['walkers'] - + hamiltonian, + comm, + nelec, + mu, + beta, + timestep, + nwalkers=nwalkers, + lowrank=lowrank, + mf_trial=mf_trial, + seed=seed, + verbose=verbose, + ) + legacy_system = legacy_objs["system"] + legacy_trial = legacy_objs["trial"] + legacy_hamiltonian = legacy_objs["hamiltonian"] + legacy_walkers = legacy_objs["walkers"] + for iw in range(walkers.nwalkers): - P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) eloc = local_energy_generic_cholesky(hamiltonian, P) legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) legacy_eloc = legacy_local_energy_generic_cholesky( - legacy_system, legacy_hamiltonian, legacy_P) + legacy_system, legacy_hamiltonian, legacy_P + ) numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) - numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10 + ) + numpy.testing.assert_almost_equal( + legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10 + ) + if __name__ == "__main__": test_thermal_walkers_fullrank() - #test_thermal_walkers_lowrank() + # test_thermal_walkers_lowrank() diff --git a/ipie/addons/thermal/walkers/uhf_walkers.py b/ipie/addons/thermal/walkers/uhf_walkers.py index e3cac522..9348f422 100644 --- a/ipie/addons/thermal/walkers/uhf_walkers.py +++ b/ipie/addons/thermal/walkers/uhf_walkers.py @@ -27,20 +27,20 @@ from ipie.walkers.base_walkers import BaseWalkers from ipie.addons.thermal.trial.one_body import OneBody + class UHFThermalWalkers(BaseWalkers): def __init__( self, trial: OneBody, nbasis: int, nwalkers: int, - stack_size = None, + stack_size=None, lowrank: bool = False, lowrank_thresh: float = 1e-6, - mpi_handler = None, + mpi_handler=None, verbose: bool = False, ): - """UHF style walker. - """ + """UHF style walker.""" assert isinstance(trial, OneBody) super().__init__(nwalkers, verbose=verbose) @@ -67,16 +67,14 @@ def __init__( self.lowrank_thresh = lowrank_thresh self.Ga = numpy.zeros( - shape=(self.nwalkers, self.nbasis, self.nbasis), - dtype=numpy.complex128) + shape=(self.nwalkers, self.nbasis, self.nbasis), dtype=numpy.complex128 + ) self.Gb = numpy.zeros( - shape=(self.nwalkers, self.nbasis, self.nbasis), - dtype=numpy.complex128) + shape=(self.nwalkers, self.nbasis, self.nbasis), dtype=numpy.complex128 + ) self.Ghalf = None - max_diff_diag = numpy.linalg.norm( - (numpy.diag( - trial.dmat[0].diagonal()) - trial.dmat[0])) + max_diff_diag = numpy.linalg.norm((numpy.diag(trial.dmat[0].diagonal()) - trial.dmat[0])) if max_diff_diag < 1e-10: self.diagonal_trial = True @@ -91,17 +89,20 @@ def __init__( print(f"# Walker stack size: {self.stack_size}") print(f"# Using low rank trick: {self.lowrank}") - self.stack = [PropagatorStack( - self.stack_size, - self.nslice, - self.nbasis, - numpy.complex128, - trial.dmat, - trial.dmat_inv, - diagonal=self.diagonal_trial, - lowrank=self.lowrank, - thresh=self.lowrank_thresh, - ) for iw in range(self.nwalkers)] + self.stack = [ + PropagatorStack( + self.stack_size, + self.nslice, + self.nbasis, + numpy.complex128, + trial.dmat, + trial.dmat_inv, + diagonal=self.diagonal_trial, + lowrank=self.lowrank, + thresh=self.lowrank_thresh, + ) + for iw in range(self.nwalkers) + ] # Initialise all propagators to the trial density matrix. for iw in range(self.nwalkers): @@ -109,12 +110,14 @@ def __init__( greens_function_qr_strat(self, iw) self.stack[iw].G[0] = self.Ga[iw] self.stack[iw].G[1] = self.Gb[iw] - + # Shape (nwalkers,). - self.M0a = numpy.array([ - scipy.linalg.det(self.Ga[iw], check_finite=False) for iw in range(self.nwalkers)]) - self.M0b = numpy.array([ - scipy.linalg.det(self.Gb[iw], check_finite=False) for iw in range(self.nwalkers)]) + self.M0a = numpy.array( + [scipy.linalg.det(self.Ga[iw], check_finite=False) for iw in range(self.nwalkers)] + ) + self.M0b = numpy.array( + [scipy.linalg.det(self.Gb[iw], check_finite=False) for iw in range(self.nwalkers)] + ) for iw in range(self.nwalkers): self.stack[iw].ovlp = numpy.array([1.0 / self.M0a[iw], 1.0 / self.M0b[iw]]) @@ -132,14 +135,12 @@ def __init__( self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) def calc_greens_function(self, iw, slice_ix=None, inplace=True): - """Return the Green's function for walker `iw`. - """ + """Return the Green's function for walker `iw`.""" if self.lowrank: - return self.stack[iw].G # G[0] = Ga, G[1] = Gb + return self.stack[iw].G # G[0] = Ga, G[1] = Gb else: return greens_function_qr_strat(self, iw, slice_ix=slice_ix, inplace=inplace) - def reset(self, trial): self.weight = numpy.ones(self.nwalkers) @@ -149,7 +150,7 @@ def reset(self, trial): self.stack[iw].reset() self.stack[iw].set_all(trial.dmat) self.calc_greens_function(iw) - + # For compatibiltiy with BaseWalkers class. def reortho(self): pass diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 23b5652b..8a5eb13f 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -15,6 +15,8 @@ # Author: Fionn Malone # +from typing import Union + import plum from ipie.estimators.estimator_base import EstimatorBase @@ -42,7 +44,6 @@ from ipie.trial_wavefunction.single_det import SingleDet from ipie.utils.backend import arraylib as xp from ipie.walkers.uhf_walkers import UHFWalkers -from typing import Union @plum.dispatch @@ -133,7 +134,7 @@ def __init__( self.print_to_stdout = True self.ascii_filename = filename - def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): + def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None): trial.calc_greens_function(walkers) # Need to be able to dispatch here energy = local_energy(system, hamiltonian, walkers, trial) diff --git a/ipie/estimators/estimator_base.py b/ipie/estimators/estimator_base.py index c78c9322..0cedfda5 100644 --- a/ipie/estimators/estimator_base.py +++ b/ipie/estimators/estimator_base.py @@ -15,13 +15,13 @@ # Author: Fionn Malone # -from abc import abstractmethod, ABCMeta +from abc import ABCMeta, abstractmethod import numpy as np -from ipie.utils.io import format_fixed_width_strings, format_fixed_width_floats from ipie.utils.backend import arraylib as xp from ipie.utils.backend import to_host +from ipie.utils.io import format_fixed_width_floats, format_fixed_width_strings class EstimatorBase(metaclass=ABCMeta): @@ -78,7 +78,7 @@ def size(self) -> int: size = 0 for _, v in self._data.items(): if isinstance(v, np.ndarray): - size += np.prod(v.shape) + size += int(np.prod(v.shape)) else: size += 1 return size @@ -89,8 +89,9 @@ def shape(self, shape) -> tuple: self._shape = shape @abstractmethod - def compute_estimator(self, system, walkers, hamiltonian, trial) -> np.ndarray: - ... + def compute_estimator( + self, system=None, walkers=None, hamiltonian=None, trial=None + ) -> np.ndarray: ... @property def names(self): @@ -142,5 +143,4 @@ def zero(self): else: self._data[k] = 0.0j - def post_reduce_hook(self, data) -> None: - ... + def post_reduce_hook(self, data) -> None: ... diff --git a/ipie/estimators/generic.py b/ipie/estimators/generic.py index 6f69d085..469abc3c 100644 --- a/ipie/estimators/generic.py +++ b/ipie/estimators/generic.py @@ -33,9 +33,7 @@ def local_energy_generic_opt(system, G, Ghalf=None, eri=None): assert eri is not None vipjq_aa = eri[0, : na**2 * M**2].reshape((na, M, na, M)) - vipjq_bb = eri[0, na**2 * M**2 : na**2 * M**2 + nb**2 * M**2].reshape( - (nb, M, nb, M) - ) + vipjq_bb = eri[0, na**2 * M**2 : na**2 * M**2 + nb**2 * M**2].reshape((nb, M, nb, M)) vipjq_ab = eri[0, na**2 * M**2 + nb**2 * M**2 :].reshape((na, M, nb, M)) Ga, Gb = Ghalf[0], Ghalf[1] diff --git a/ipie/estimators/handler.py b/ipie/estimators/handler.py index a35dbfe1..cbe01fe9 100644 --- a/ipie/estimators/handler.py +++ b/ipie/estimators/handler.py @@ -190,7 +190,7 @@ def increment_file_number(self): self.index = self.index + 1 self.filename = self.basename + f".{self.index}.h5" - def compute_estimators(self, comm, system, hamiltonian, trial, walker_batch): + def compute_estimators(self, system=None, hamiltonian=None, trial=None, walker_batch=None): """Update estimators with bached psi Parameters @@ -201,7 +201,9 @@ def compute_estimators(self, comm, system, hamiltonian, trial, walker_batch): # TODO: generalize for different block groups (loop over groups) offset = self.num_walker_props for k, e in self.items(): - e.compute_estimator(system, walker_batch, hamiltonian, trial) + e.compute_estimator( + system=system, walkers=walker_batch, hamiltonian=hamiltonian, trial=trial + ) start = offset + self.get_offset(k) end = start + int(self[k].size) self.local_estimates[start:end] += e.data diff --git a/ipie/estimators/tests/test_estimators.py b/ipie/estimators/tests/test_estimators.py index 59adc89f..0bd480fd 100644 --- a/ipie/estimators/tests/test_estimators.py +++ b/ipie/estimators/tests/test_estimators.py @@ -72,5 +72,5 @@ def test_estimator_handler(): handler["energy1"] = estim handler.json_string = "" handler.initialize(comm) - handler.compute_estimators(comm, system, ham, trial, walker_batch) - handler.compute_estimators(comm, system, ham, trial, walker_batch) + handler.compute_estimators(system, ham, trial, walker_batch) + handler.compute_estimators(system, ham, trial, walker_batch) diff --git a/ipie/qmc/afqmc.py b/ipie/qmc/afqmc.py index 9ce27e89..e71edbc4 100644 --- a/ipie/qmc/afqmc.py +++ b/ipie/qmc/afqmc.py @@ -17,6 +17,7 @@ # """Driver to perform AFQMC calculation""" +import abc import json import time import uuid @@ -41,8 +42,216 @@ from ipie.walkers.walkers_dispatch import get_initial_walker, UHFWalkersTrial -class AFQMC(object): - """AFQMC driver. +class AFQMCBase(metaclass=abc.ABCMeta): + """Basic interface for AFQMC Driver""" + + def __init__( + self, + system, + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params: QMCParams, + verbose: int = 0, + ): + self.system = system + self.hamiltonian = hamiltonian + self.trial = trial + self.walkers = walkers + self.propagator = propagator + self.mpi_handler = mpi_handler # mpi_handler should be passed into here + self.shared_comm = self.mpi_handler.shared_comm + self.verbose = verbose + self.verbosity = int(verbose) + self.params = params + self._init_time = time.time() + self._parallel_rng_seed = set_rng_seed(params.rng_seed, self.mpi_handler.comm) + + @abc.abstractmethod + def run( + self, + walkers=None, + estimator_filename=None, + verbose=True, + additional_estimators: Optional[Dict[str, EstimatorBase]] = None, + ): + """Code to run the AFQMC calculation.""" + raise NotImplementedError + + def distribute_hamiltonian(self): + if self.mpi_handler.nmembers > 1: + if self.mpi_handler.comm.rank == 0: + print("# Chunking hamiltonian.") + # self.hamiltonian.chunk(self.mpi_handler) + if self.mpi_handler.comm.rank == 0: + print("# Chunking trial.") + self.trial.chunk(self.mpi_handler) + + def copy_to_gpu(self): + comm = self.mpi_handler.comm + if config.get_option("use_gpu"): + ngpus = xp.cuda.runtime.getDeviceCount() + _ = xp.cuda.runtime.getDeviceProperties(0) + # xp.cuda.runtime.setDevice(self.shared_comm.rank % 4) + xp.cuda.runtime.setDevice(comm.rank % ngpus) + if comm.rank == 0: + if ngpus > comm.size: + print( + f"# There are unused GPUs ({comm.size} MPI tasks but {ngpus} GPUs). " + " Check if this is really what you wanted." + ) + self.propagator.cast_to_cupy(self.verbose and comm.rank == 0) + self.hamiltonian.cast_to_cupy(self.verbose and comm.rank == 0) + self.trial.cast_to_cupy(self.verbose and comm.rank == 0) + self.walkers.cast_to_cupy(self.verbose and comm.rank == 0) + + def get_env_info(self): + # TODO: Move this somewhere else. + this_uuid = str(uuid.uuid1()) + try: + sha1, branch, local_mods = get_git_info() + except: + sha1 = "None" + branch = "None" + local_mods = [] + if self.verbose: + self.sys_info = print_env_info( + sha1, branch, local_mods, this_uuid, self.mpi_handler.size + ) + mem_avail = get_host_memory() + print(f"# MPI communicator : {type(self.mpi_handler.comm)}") + print(f"# Available memory on the node is {mem_avail:4.3f} GB") + + def setup_timers(self): + # TODO: Better timer + self.tsetup = 0 + self.tortho = 0 + self.tprop = 0 + + self.tprop_fbias = 0.0 + self.tprop_ovlp = 0.0 + self.tprop_update = 0.0 + self.tprop_gf = 0.0 + self.tprop_vhs = 0.0 + self.tprop_gemm = 0.0 + self.tprop_clip = 0.0 + self.tprop_barrier = 0.0 + + self.testim = 0 + self.tpopc = 0 + self.tpopc_comm = 0 + self.tpopc_non_comm = 0 + self.tstep = 0 + + def finalise(self, verbose=False): + """Tidy up. + + Parameters + ---------- + verbose : bool + If true print out some information to stdout. + """ + nsteps = max(self.params.num_steps_per_block, 1) + nblocks = max(self.params.num_blocks, 1) + nstblz = max(nsteps // self.params.num_stblz, 1) + npcon = max(nsteps // self.params.pop_control_freq, 1) + if self.mpi_handler.rank == 0: + if verbose: + print(f"# End Time: {time.asctime():s}") + print(f"# Running time : {time.time() - self._init_time:.6f} seconds") + print("# Timing breakdown (per call, total calls per block, total blocks):") + print(f"# - Setup: {self.tsetup:.6f} s") + print( + "# - Block: {:.6f} s / block for {} total blocks".format( + self.tstep / (nblocks), nblocks + ) + ) + print( + "# - Propagation: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Force bias: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_fbias / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - VHS: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_vhs / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Green's Function: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_gf / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Overlap: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_ovlp / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Weights Update: {:.6f} s / call for {} call(s) in each of {} blocks".format( + (self.tprop_update + self.tprop_clip) / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - GEMM operations: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_gemm / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Barrier: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tprop_barrier / (nblocks * nsteps), nsteps, nblocks + ) + ) + print( + "# - Estimators: {:.6f} s / call for {} call(s)".format( + self.testim / nblocks, nblocks + ) + ) + print( + "# - Orthogonalisation: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tortho / (nstblz * nblocks), nstblz, nblocks + ) + ) + print( + "# - Population control: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tpopc / (npcon * nblocks), npcon, nblocks + ) + ) + print( + "# - Commnication: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tpopc_comm / (npcon * nblocks), npcon, nblocks + ) + ) + print( + "# - Non-Commnication: {:.6f} s / call for {} call(s) in each of {} blocks".format( + self.tpopc_non_comm / (npcon * nblocks), npcon, nblocks + ) + ) + + def determine_dtype(self, propagator, system): + """Determine dtype for trial wavefunction and walkers. + + Parameters + ---------- + propagator : dict + Propagator input options. + system : object + system object. + """ + hs_type = propagator.get("hubbard_stratonovich", "discrete") + continuous = "continuous" in hs_type + twist = system.ktwist.all() is not None + return continuous or twist + + +class AFQMC(AFQMCBase): + """AFQMC driver for zero temperature open ended random walk. Parameters ---------- @@ -79,18 +288,9 @@ def __init__( params: QMCParams, verbose: int = 0, ): - self.system = system - self.hamiltonian = hamiltonian - self.trial = trial - self.walkers = walkers - self.propagator = propagator - self.mpi_handler = mpi_handler # mpi_handler should be passed into here - self.shared_comm = self.mpi_handler.shared_comm - self.verbose = verbose - self.verbosity = int(verbose) - self.params = params - self._init_time = time.time() - self._parallel_rng_seed = set_rng_seed(params.rng_seed, self.mpi_handler.comm) + super().__init__( + system, hamiltonian, trial, walkers, propagator, mpi_handler, params, verbose + ) @staticmethod # TODO: wavefunction type, trial type, hamiltonian type @@ -100,7 +300,7 @@ def build( trial_wavefunction, walkers=None, num_walkers: int = 100, - seed: int = None, + seed: Optional[int] = None, num_steps_per_block: int = 25, num_blocks: int = 100, timestep: float = 0.005, @@ -117,7 +317,7 @@ def build( Number of alpha and beta electrons. hamiltonian : Hamiltonian describing the system. - trial_wavefunction : + trial_wavefunction: Trial wavefunction num_walkers : int Number of walkers per MPI process used in the simulation. The TOTAL @@ -176,7 +376,7 @@ def build( ) walkers.build( trial_wavefunction - ) # any intermediates that require information from trial + ) # any intermediates that require information from trial_wavefunction # TODO: this is a factory not a class propagator = Propagator[type(hamiltonian)](params.timestep) propagator.build(hamiltonian, trial_wavefunction, walkers, mpi_handler) @@ -198,7 +398,7 @@ def build_from_hdf5( ham_file, wfn_file, num_walkers: int = 100, - seed: int = None, + seed: Optional[int] = None, num_steps_per_block: int = 25, num_blocks: int = 100, timestep: float = 0.005, @@ -274,50 +474,6 @@ def build_from_hdf5( mpi_handler=mpi_handler, ) - def distribute_hamiltonian(self): - if self.mpi_handler.nmembers > 1: - if self.mpi_handler.comm.rank == 0: - print("# Chunking hamiltonian.") - # self.hamiltonian.chunk(self.mpi_handler) - if self.mpi_handler.comm.rank == 0: - print("# Chunking trial.") - self.trial.chunk(self.mpi_handler) - - def copy_to_gpu(self): - comm = self.mpi_handler.comm - if config.get_option("use_gpu"): - ngpus = xp.cuda.runtime.getDeviceCount() - _ = xp.cuda.runtime.getDeviceProperties(0) - # xp.cuda.runtime.setDevice(self.shared_comm.rank % 4) - xp.cuda.runtime.setDevice(comm.rank % ngpus) - if comm.rank == 0: - if ngpus > comm.size: - print( - f"# There are unused GPUs ({comm.size} MPI tasks but {ngpus} GPUs). " - " Check if this is really what you wanted." - ) - self.propagator.cast_to_cupy(self.verbose and comm.rank == 0) - self.hamiltonian.cast_to_cupy(self.verbose and comm.rank == 0) - self.trial.cast_to_cupy(self.verbose and comm.rank == 0) - self.walkers.cast_to_cupy(self.verbose and comm.rank == 0) - - def get_env_info(self): - # TODO: Move this somewhere else. - this_uuid = str(uuid.uuid1()) - try: - sha1, branch, local_mods = get_git_info() - except: - sha1 = "None" - branch = "None" - local_mods = [] - if self.verbose: - self.sys_info = print_env_info( - sha1, branch, local_mods, this_uuid, self.mpi_handler.size - ) - mem_avail = get_host_memory() - print(f"# MPI communicator : {type(self.mpi_handler.comm)}") - print(f"# Available memory on the node is {mem_avail:4.3f} GB") - def setup_estimators( self, filename, additional_estimators: Optional[Dict[str, EstimatorBase]] = None ): @@ -344,37 +500,14 @@ def setup_estimators( self.estimators.initialize(comm) # Calculate estimates for initial distribution of walkers. - self.estimators.compute_estimators( - comm, self.system, self.hamiltonian, self.trial, self.walkers - ) + self.estimators.compute_estimators(self.system, self.hamiltonian, self.trial, self.walkers) self.accumulators.update(self.walkers) self.estimators.print_block(comm, 0, self.accumulators) self.accumulators.zero() - def setup_timers(self): - # TODO: Better timer - self.tsetup = 0 - self.tortho = 0 - self.tprop = 0 - - self.tprop_fbias = 0.0 - self.tprop_ovlp = 0.0 - self.tprop_update = 0.0 - self.tprop_gf = 0.0 - self.tprop_vhs = 0.0 - self.tprop_gemm = 0.0 - self.tprop_clip = 0.0 - self.tprop_barrier = 0.0 - - self.testim = 0 - self.tpopc = 0 - self.tpopc_comm = 0 - self.tpopc_non_comm = 0 - self.tstep = 0 - def run( self, - psi=None, + walkers=None, estimator_filename=None, verbose=True, additional_estimators: Optional[Dict[str, EstimatorBase]] = None, @@ -383,7 +516,7 @@ def run( Parameters ---------- - psi : :class:`pie.walker.Walkers` object + walkers : :class:`pie.walker.Walkers` object Initial wavefunction / distribution of walkers. Default None. estimator_filename : str File to write estimates to. @@ -392,8 +525,8 @@ def run( """ self.setup_timers() tzero_setup = time.time() - if psi is not None: - self.walkers = psi + if walkers is not None: + self.walkers = walkers self.setup_timers() eshift = 0.0 self.walkers.orthogonalise() @@ -479,7 +612,7 @@ def run( start = time.time() if step % self.params.num_steps_per_block == 0: self.estimators.compute_estimators( - comm, self.system, self.hamiltonian, self.trial, self.walkers + self.system, self.hamiltonian, self.trial, self.walkers ) self.estimators.print_block( comm, step // self.params.num_steps_per_block, self.accumulators @@ -498,107 +631,3 @@ def run( eshift += self.accumulators.eshift - eshift synchronize() self.tstep += time.time() - start_step - - def finalise(self, verbose=False): - """Tidy up. - - Parameters - ---------- - verbose : bool - If true print out some information to stdout. - """ - nsteps = max(self.params.num_steps_per_block, 1) - nblocks = max(self.params.num_blocks, 1) - nstblz = max(nsteps // self.params.num_stblz, 1) - npcon = max(nsteps // self.params.pop_control_freq, 1) - if self.mpi_handler.rank == 0: - if verbose: - print(f"# End Time: {time.asctime():s}") - print(f"# Running time : {time.time() - self._init_time:.6f} seconds") - print("# Timing breakdown (per call, total calls per block, total blocks):") - print(f"# - Setup: {self.tsetup:.6f} s") - print( - "# - Block: {:.6f} s / block for {} total blocks".format( - self.tstep / (nblocks), nblocks - ) - ) - print( - "# - Propagation: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Force bias: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_fbias / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - VHS: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_vhs / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Green's Function: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_gf / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Overlap: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_ovlp / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Weights Update: {:.6f} s / call for {} call(s) in each of {} blocks".format( - (self.tprop_update + self.tprop_clip) / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - GEMM operations: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_gemm / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Barrier: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tprop_barrier / (nblocks * nsteps), nsteps, nblocks - ) - ) - print( - "# - Estimators: {:.6f} s / call for {} call(s)".format( - self.testim / nblocks, nblocks - ) - ) - print( - "# - Orthogonalisation: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tortho / (nstblz * nblocks), nstblz, nblocks - ) - ) - print( - "# - Population control: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tpopc / (npcon * nblocks), npcon, nblocks - ) - ) - print( - "# - Commnication: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tpopc_comm / (npcon * nblocks), npcon, nblocks - ) - ) - print( - "# - Non-Commnication: {:.6f} s / call for {} call(s) in each of {} blocks".format( - self.tpopc_non_comm / (npcon * nblocks), npcon, nblocks - ) - ) - - def determine_dtype(self, propagator, system): - """Determine dtype for trial wavefunction and walkers. - - Parameters - ---------- - propagator : dict - Propagator input options. - system : object - system object. - """ - hs_type = propagator.get("hubbard_stratonovich", "discrete") - continuous = "continuous" in hs_type - twist = system.ktwist.all() is not None - return continuous or twist diff --git a/ipie/qmc/options.py b/ipie/qmc/options.py index 437a5422..faa9a1a6 100644 --- a/ipie/qmc/options.py +++ b/ipie/qmc/options.py @@ -156,7 +156,7 @@ class QMCParams: pop_control_freq : int Frequency at which population control occurs. rng_seed : int - The random number seed. If run in parallel the seeds on other cores / + The random number seed. If run in parallel the seeds on other cores / threads are determined from this. """ diff --git a/ipie/qmc/tests/test_afqmc_multi_det_batch.py b/ipie/qmc/tests/test_afqmc_multi_det_batch.py index f37d4320..b9f0aae6 100644 --- a/ipie/qmc/tests/test_afqmc_multi_det_batch.py +++ b/ipie/qmc/tests/test_afqmc_multi_det_batch.py @@ -75,7 +75,7 @@ def test_generic_multi_det_batch(): afqmc.run(verbose=0, estimator_filename=tmpf.name) afqmc.finalise(verbose=0) afqmc.estimators.compute_estimators( - comm, afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers + afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers ) numer_batch = afqmc.estimators["energy"]["ENumer"] denom_batch = afqmc.estimators["energy"]["EDenom"] diff --git a/ipie/qmc/tests/test_afqmc_single_det_batch.py b/ipie/qmc/tests/test_afqmc_single_det_batch.py index 8a426314..c11317dd 100644 --- a/ipie/qmc/tests/test_afqmc_single_det_batch.py +++ b/ipie/qmc/tests/test_afqmc_single_det_batch.py @@ -78,7 +78,7 @@ def test_generic_single_det_batch(): afqmc.run(verbose=False, estimator_filename=tmpf.name) afqmc.finalise(verbose=0) afqmc.estimators.compute_estimators( - comm, afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers + afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers ) numer_batch = afqmc.estimators["energy"]["ENumer"] denom_batch = afqmc.estimators["energy"]["EDenom"] @@ -165,7 +165,7 @@ def test_generic_single_det_batch_density_diff(): afqmc.run(verbose=False, estimator_filename=tmpf.name) afqmc.finalise(verbose=0) afqmc.estimators.compute_estimators( - comm, afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers + afqmc.system, afqmc.hamiltonian, afqmc.trial, afqmc.walkers ) numer_batch = afqmc.estimators["energy"]["ENumer"] diff --git a/ipie/trial_wavefunction/single_det.py b/ipie/trial_wavefunction/single_det.py index 3af4e335..abc043fd 100644 --- a/ipie/trial_wavefunction/single_det.py +++ b/ipie/trial_wavefunction/single_det.py @@ -40,7 +40,7 @@ def __init__(self, wavefunction, num_elec, num_basis, handler=MPIHandler(), verb self._max_num_dets = 1 imag_norm = numpy.sum(self.psi.imag.ravel() * self.psi.imag.ravel()) if imag_norm <= 1e-8: - #print("# making trial wavefunction MO coefficient real") + # print("# making trial wavefunction MO coefficient real") self.psi = numpy.array(self.psi.real, dtype=numpy.float64) self.psi0a = self.psi[:, : self.nalpha] diff --git a/ipie/trial_wavefunction/tests/test_noci.py b/ipie/trial_wavefunction/tests/test_noci.py index d7acf4d2..a8a96c12 100644 --- a/ipie/trial_wavefunction/tests/test_noci.py +++ b/ipie/trial_wavefunction/tests/test_noci.py @@ -35,5 +35,5 @@ def test_noci(): trial.calculate_energy(sys, ham) -if __name__ == '__main__': +if __name__ == "__main__": test_noci() diff --git a/ipie/trial_wavefunction/tests/test_wavefunction_base.py b/ipie/trial_wavefunction/tests/test_wavefunction_base.py index 94e47cff..43b6f20b 100644 --- a/ipie/trial_wavefunction/tests/test_wavefunction_base.py +++ b/ipie/trial_wavefunction/tests/test_wavefunction_base.py @@ -18,5 +18,5 @@ def test_wavefunction_base(): assert trial.nbasis == num_basis -if __name__ == '__main__': +if __name__ == "__main__": test_wavefunction_base() diff --git a/ipie/utils/misc.py b/ipie/utils/misc.py index 4c956d98..b8195f9d 100644 --- a/ipie/utils/misc.py +++ b/ipie/utils/misc.py @@ -25,6 +25,7 @@ import time import types from functools import reduce +from typing import Dict import numpy import scipy.sparse @@ -294,6 +295,37 @@ def get_node_mem(): return 0.0 +def get_numpy_blas_info() -> Dict[str, str]: + """Get useful numpy blas / lapack info.""" + info = {} + try: + config = numpy.show_config(mode="dicts") + blas_config = config["Build Dependencies"]["blas"] + info["BLAS"] = { + "lib": blas_config["name"], + "version": blas_config["version"], + "include directory": blas_config["lib directory"], + } + for k, v in info["BLAS"].items(): + print(f"# - BLAS {k}: {v}") + except TypeError: + try: + np_lib = numpy.__config__.blas_opt_info["libraries"] # pylint: disable=no-member + lib_dir = numpy.__config__.blas_opt_info["library_dirs"] # pylint: disable=no-member + except AttributeError: + np_lib = numpy.__config__.blas_ilp64_opt_info["libraries"] # pylint: disable=no-member + lib_dir = numpy.__config__.blas_ilp64_opt_info[ # pylint: disable=no-member + "library_dirs" + ] + print(f"# - BLAS lib: {' '.join(np_lib):s}") + print(f"# - BLAS dir: {' '.join(lib_dir):s}") + info["BLAS"] = { + "lib": " ".join(np_lib), + "path": " ".join(lib_dir), + } + return info + + def print_env_info(sha1, branch, local_mods, uuid, nranks): import ipie @@ -329,20 +361,7 @@ def print_env_info(sha1, branch, local_mods, uuid, nranks): print(f"# Using {lib:s} v{vers:s} from: {path:s}.") info[f"{lib:s}"] = {"version": vers, "path": path} if lib == "numpy": - try: - np_lib = l.__config__.blas_opt_info["libraries"] - except AttributeError: - np_lib = l.__config__.blas_ilp64_opt_info["libraries"] - print(f"# - BLAS lib: {' '.join(np_lib):s}") - try: - lib_dir = l.__config__.blas_opt_info["library_dirs"] - except AttributeError: - lib_dir = l.__config__.blas_ilp64_opt_info["library_dirs"] - print(f"# - BLAS dir: {' '.join(lib_dir):s}") - info[f"{lib:s}"]["BLAS"] = { - "lib": " ".join(np_lib), - "path": " ".join(lib_dir), - } + info[f"{lib:s}"] = get_numpy_blas_info() elif lib == "mpi4py": mpicc = l.get_config().get("mpicc", "none") print(f"# - mpicc: {mpicc:s}") diff --git a/ipie/walkers/pop_controller.py b/ipie/walkers/pop_controller.py index 15f7efbf..65e2cdb8 100644 --- a/ipie/walkers/pop_controller.py +++ b/ipie/walkers/pop_controller.py @@ -5,6 +5,7 @@ from ipie.config import MPI from ipie.utils.backend import arraylib as xp + class PopControllerTimer: def __init__(self): self.start_time_const = 0.0 @@ -173,13 +174,15 @@ def set_buffer(walkers, iw, buff): assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated if isinstance(data[iw], xp.ndarray): walkers.__dict__[d][iw] = xp.array( - buff[s : s + data[iw].size].reshape(data[iw].shape).copy()) + buff[s : s + data[iw].size].reshape(data[iw].shape).copy() + ) s += data[iw].size elif isinstance(data[iw], list): for ix, l in enumerate(data[iw]): if isinstance(l, (xp.ndarray)): walkers.__dict__[d][iw][ix] = xp.array( - buff[s : s + l.size].reshape(l.shape).copy()) + buff[s : s + l.size].reshape(l.shape).copy() + ) s += l.size elif isinstance(l, (int, float, complex)): walkers.__dict__[d][iw][ix] = buff[s] @@ -315,11 +318,11 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( glob_inf_1 = numpy.empty([comm.size, walkers.nwalkers], dtype=numpy.int64) glob_inf_1.fill(1) glob_inf_2 = numpy.array( - [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64) + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], dtype=numpy.int64 + ) glob_inf_3 = numpy.array( - [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64) + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], dtype=numpy.int64 + ) timer.add_non_communication() @@ -337,9 +340,15 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( # Rescale weights. glob_inf = numpy.zeros((walkers.nwalkers * comm.size, 4), dtype=numpy.float64) glob_inf[:, 0] = glob_inf_0.ravel() # contains walker |w_i| - glob_inf[:, 1] = glob_inf_1.ravel() # all initialized to 1 when it becomes 2 then it will be "branched" - glob_inf[:, 2] = glob_inf_2.ravel() # contain processor+walker indices (initial) (i.e., where walkers live) - glob_inf[:, 3] = glob_inf_3.ravel() # contain processor+walker indices (final) (i.e., where walkers live) + glob_inf[:, 1] = ( + glob_inf_1.ravel() + ) # all initialized to 1 when it becomes 2 then it will be "branched" + glob_inf[:, 2] = ( + glob_inf_2.ravel() + ) # contain processor+walker indices (initial) (i.e., where walkers live) + glob_inf[:, 3] = ( + glob_inf_3.ravel() + ) # contain processor+walker indices (final) (i.e., where walkers live) sort = numpy.argsort(glob_inf[:, 0], kind="mergesort") isort = numpy.argsort(sort, kind="mergesort") glob_inf = glob_inf[sort] diff --git a/requirements.txt b/requirements.txt index 3d9432a0..41dbc0a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ cython >= 0.29.0 h5py >= 3.0.0 -numpy >= 1.20.0, < 1.26.0 -scipy >= 1.3.0, <=1.10.1 +numpy >= 1.20.0 +scipy >= 1.3.0 pytest -pandas == 1.5.1 # issue with pyblock and pandas > 2 +pandas numba plum-dispatch diff --git a/setup.py b/setup.py index c6857748..0cb2f5f4 100644 --- a/setup.py +++ b/setup.py @@ -66,7 +66,7 @@ def main() -> None: packages=find_packages(exclude=["examples", "docs", "tests", "tools", "setup.py"]), license="Apache 2.0", description="Python implementations of Imaginary-time Evolution algorithms", - python_requires=">=3.7.0,<3.12.0", + python_requires=">=3.8.0", scripts=[ "bin/ipie", "tools/extract_dice.py",