Skip to content

Commit

Permalink
cleanup ThermalAFQMC class for only Generic case, by inheriting ZT AF…
Browse files Browse the repository at this point in the history
…QMC class
  • Loading branch information
jiangtong1000 committed Sep 1, 2023
1 parent c68f886 commit 4c1e2c6
Show file tree
Hide file tree
Showing 5 changed files with 276 additions and 42 deletions.
2 changes: 1 addition & 1 deletion ipie/thermal/propagation/continuous.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def __init__(self, options, qmc, system, hamiltonian, trial, verbose=False, lowr
system, hamiltonian, trial, qmc, options=options, verbose=verbose
)
# Mean field shifted one-body propagator
self.mu = system.mu
self.mu = hamiltonian.mu
self.propagator.construct_one_body_propagator(hamiltonian, qmc.dt)

self.BH1 = self.propagator.BH1
Expand Down
193 changes: 193 additions & 0 deletions ipie/thermal/qmc/thermal_afqmc_clean.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
"""Driver to perform AFQMC calculation"""
import time
import uuid
from ipie.thermal.propagation.utils import get_propagator
from ipie.thermal.walkers.handler import Walkers
from ipie.qmc.options import QMCOpts
from ipie.qmc.utils import set_rng_seed
from ipie.utils.misc import get_git_info, print_env_info
from ipie.qmc.afqmc import AFQMC
from ipie.qmc.options import QMCParams
from ipie.systems.generic import Generic
from ipie.utils.mpi import MPIHandler


## This is now only applicable to the Generic case!
## See test_generic.py for example.

class ThermalAFQMC(AFQMC):
def __init__(
self,
num_elec,
mu,
beta,
hamiltonian,
trial,
num_walkers: int = 100,
seed: int = None,
num_steps_per_block: int = 25,
num_blocks: int = 100,
timestep: float = 0.005,
stabilize_freq=5,
pop_control_freq=5,
verbose=True,
):

mpi_handler = MPIHandler()
comm = mpi_handler.comm
params = QMCParams(
num_walkers=num_walkers,
total_num_walkers=num_walkers * comm.size,
num_blocks=num_blocks,
num_steps_per_block=num_steps_per_block,
timestep=timestep,
num_stblz=stabilize_freq,
pop_control_freq=pop_control_freq,
rng_seed=seed,
)

system = Generic(num_elec, verbose)
system.mu = mu

qmc = QMCOpts() # should be removed later after walker is cleaned up
qmc.nwalkers = num_walkers
qmc.ntot_walkers = num_walkers * comm.size
qmc.beta = beta
qmc.total_steps = num_blocks
qmc.nsteps = 1
qmc.ntime_slices = int(round(beta / timestep))
qmc.rng_seed = set_rng_seed(seed, comm)
qmc.dt = timestep

if comm.rank == 0:
self.uuid = str(uuid.uuid1())
get_sha1 = True
if get_sha1:
self.sha1, self.branch, self.local_mods = get_git_info()
else:
self.sha1 = "None"
if verbose:
self.sys_info = print_env_info(
self.sha1, self.branch, self.local_mods, self.uuid, comm.size
)

self.qmc = qmc
self.qmc.nstblz = 10
self.qmc.npop_control = 1

wlk_opts = {} # should be removed later after walker is cleaned up
walkers = Walkers(
system,
hamiltonian,
trial,
qmc,
walker_opts=wlk_opts,
verbose=verbose)

prop_opts = {} # should be removed later after walker is cleaned up

propagator = get_propagator(
prop_opts,
qmc,
system,
hamiltonian,
trial,
verbose=verbose
)
self.propagators = propagator
self.root = True

super().__init__(system, hamiltonian, trial, walkers, propagator, params, verbose)

def run(self, walkers=None, comm=None, verbose=None):
"""Perform AFQMC simulation on state object using open-ended random walk.
Parameters
----------
state : :class:`pie.state.State` object
Model and qmc parameters.
walk: :class:`pie.walker.Walkers` object
Initial wavefunction / distribution of walkers.
comm : MPI communicator
"""
if walkers is not None:
self.walkers = walkers
self.setup_timers()
# Calculate estimates for initial distribution of walkers.
self.estimators.estimators["mixed"].update(
self.qmc,
self.system,
self.hamiltonian,
self.trial,
self.walkers,
0,
self.propagators.free_projection,
)
# Print out zeroth step for convenience.
self.estimators.estimators["mixed"].print_step(comm, comm.size, 0, 1)

for step in range(1, self.qmc.total_steps + 1):
start_path = time.time()
for ts in range(0, self.qmc.ntime_slices):
if self.verbosity >= 2 and comm.rank == 0:
print(" # Timeslice %d of %d." % (ts, self.qmc.ntime_slices))
start = time.time()
for w in self.walkers.walkers:
self.propagators.propagate_walker(self.hamiltonian, w, ts, 0)
if (abs(w.weight) > w.total_weight * 0.10) and ts > 0:
w.weight = w.total_weight * 0.10
self.tprop += time.time() - start
start = time.time()
if ts % self.qmc.npop_control == 0 and ts != 0:
self.walkers.pop_control(comm)
self.tpopc += time.time() - start
self.tpath += time.time() - start_path
start = time.time()
self.estimators.update(
self.qmc,
self.system,
self.hamiltonian,
self.trial,
self.walkers,
step,
self.propagators.free_projection,
)
self.testim += time.time() - start
self.estimators.print_step(
comm,
comm.size,
step,
free_projection=self.propagators.free_projection,
)
self.walkers.reset(self.trial)

def finalise(self, verbose):
"""Tidy up.
Parameters
----------
verbose : bool
If true print out some information to stdout.
"""
pass

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

def setup_timers(self):
self.tpath = 0
self.tprop = 0
self.testim = 0
self.tpopc = 0
Binary file not shown.
110 changes: 82 additions & 28 deletions ipie/thermal/tests/test_generic.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
from ipie.systems.generic import Generic
from pyscf import gto, scf, lo
from ipie.utils.from_pyscf import generate_hamiltonian
from ipie.hamiltonians.generic import Generic as HamGeneric
from ipie.thermal.qmc.thermal_afqmc import ThermalAFQMC
from ipie.thermal.trial.utils import get_trial_density_matrix
import numpy
import mpi4py.MPI as MPI
import h5py
from ipie.thermal.estimators.handler import Estimators
from ipie.qmc.options import QMCOpts
from ipie.utils.io import to_json

comm = MPI.COMM_WORLD

Expand All @@ -31,37 +33,48 @@
s1e = mol.intor("int1e_ovlp_sph")
ao_coeff = lo.orth.lowdin(s1e)

system = Generic(nelec=mol.nelec)

integrals = generate_hamiltonian(
mol,
mf.mo_coeff,
mf.get_hcore(),
ao_coeff,
ortho_ao=True,
chol_cut=1.e-8
)
# integrals = generate_hamiltonian(
# mol,
# mf.mo_coeff,
# mf.get_hcore(),
# ao_coeff,
# ortho_ao=True,
# chol_cut=1.e-8
# )
# num_chol = integrals.chol.shape[1]
# num_basis = integrals.nbasis
# with h5py.File("reference_data/generic_integrals.h5", "w") as fa:
# fa["hcore"] = integrals.H1[0]
# fa["LXmn"] = integrals.chol.T.reshape(num_chol, num_basis, num_basis)
# fa["e0"] = integrals.ecore

num_chol = integrals.chol.shape[1]
num_basis = integrals.nbasis
with h5py.File("reference_data/generic_integrals.h5", "r") as fa:
Lxmn = fa["LXmn"][:]
num_chol = Lxmn.shape[0]
num_basis = Lxmn.shape[1]

with h5py.File("generic_integrals.h5", "w") as fa:
fa["hcore"] = integrals.H1[0]
fa["LXmn"] = integrals.chol.T.reshape(num_chol, num_basis, num_basis)
fa["e0"] = integrals.ecore
system = Generic(nelec=mol.nelec)

mu = -10.0
beta = 0.1
dt = 0.01
num_walkers = 1
seed = 7
num_steps_per_block = 1
blocks = 10
stabilise_freq = 10
pop_control_freq = 1

options = {
"qmc": {
"dt": 0.01,
"nwalkers": 1,
"blocks": 10,
"dt": dt,
"nwalkers": num_walkers,
"blocks": blocks,
"nsteps": 10,
"beta": 0.1,
"rng_seed": 7,
"pop_control_freq": 1,
"stabilise_freq": 10,
"beta": beta,
"rng_seed": seed,
"pop_control_freq": pop_control_freq,
"stabilise_freq": stabilise_freq,
"batched": False
},
"walkers": {
Expand All @@ -77,7 +90,7 @@
},
"hamiltonian": {
"name": "Generic",
"integrals": "generic_integrals.h5",
"integrals": "reference_data/generic_integrals.h5",
"_alt_convention": False,
"sparse": False,
"mu": mu
Expand All @@ -90,7 +103,6 @@
}
}


def compare_test_data(ref, test):
comparison = {}
for k, v in ref.items():
Expand All @@ -106,13 +118,55 @@ def compare_test_data(ref, test):
print(f"# Issue with test data key {k}")
return comparison


from ipie.thermal.qmc.thermal_afqmc_clean import ThermalAFQMC
from ipie.hamiltonians.utils import get_hamiltonian
system = Generic(mol.nelec)
system.mu = mu
hamiltonian = get_hamiltonian(system, options["hamiltonian"])
trial = get_trial_density_matrix(system, hamiltonian, beta, dt, comm=comm, options=options["trial"])

qmc = QMCOpts() # should be removed later after walker is cleaned up
qmc.nwalkers = num_walkers
qmc.ntot_walkers = num_walkers * comm.size
qmc.beta = beta
qmc.nsteps = 1


afqmc = ThermalAFQMC(mol.nelec, mu, beta,
hamiltonian,
trial,
num_walkers,
seed,
num_steps_per_block,
blocks,
dt,
stabilise_freq,
pop_control_freq
)

test_name = "generic"
import tempfile
import json
from ipie.analysis.extraction import extract_test_data_hdf5
with tempfile.NamedTemporaryFile() as tmpf:
options["estimators"]["filename"] = tmpf.name
afqmc = ThermalAFQMC(comm, options=options, parallel=comm.size > 1, verbose=1)
estimators = Estimators(
options["estimators"],
comm.rank == 0,
qmc,
system,
hamiltonian,
trial,
None
)
afqmc.estimators = estimators
if comm.rank == 0:
json_string = to_json(afqmc)
afqmc.estimators.json_string = json_string
afqmc.estimators.dump_metadata()
afqmc.estimators.estimators["mixed"].print_key()
afqmc.estimators.estimators["mixed"].print_header()
afqmc.run(comm=comm)
afqmc.finalise(comm)
test_data = extract_test_data_hdf5(tmpf.name)
Expand Down
13 changes: 0 additions & 13 deletions ipie/thermal/walkers/handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,19 +73,6 @@ def __init__(
self.buff_size = self.walkers[0].buff_size + self.walkers[0].stack.buff_size
self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128)
stack_size = self.walkers[0].stack_size
if hamiltonian.name == "Hubbard":
if stack_size % qmc.nstblz != 0 or qmc.nstblz < stack_size:
if verbose:
print("# Stabilisation frequency is not commensurate " "with stack size.")
print("# Determining a better value.")
if qmc.nstblz < stack_size:
qmc.nstblz = stack_size
if verbose:
print(f"# Updated stabilization frequency: {qmc.nstblz}")
else:
qmc.nstblz = update_stack(
qmc.nstblz, stack_size, name="nstblz", verbose=verbose
)
else:
raise NotImplementedError
if hamiltonian.name == "Generic" or system.name == "UEG":
Expand Down

1 comment on commit 4c1e2c6

@jiangtong1000
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I modified ThermalAFQMC by letting it inherit the Zero Temperature AFQMC, for only generic case. But it needs to be further modified once the walker is cleaned up.

Please sign in to comment.