diff --git a/examples/02-multi_determinant/run_afqmc.py b/examples/02-multi_determinant/run_afqmc.py index 0d1753ba..ac7d774c 100644 --- a/examples/02-multi_determinant/run_afqmc.py +++ b/examples/02-multi_determinant/run_afqmc.py @@ -1,3 +1,7 @@ +""" +Run AFQMC with multi-determinant trial wavefunction. +Both GPU and MPI can be enabled in this script. +""" import h5py import numpy from pyscf import fci, gto, mcscf, scf @@ -5,37 +9,55 @@ from ipie.hamiltonians.generic import Generic as HamGeneric from ipie.qmc.afqmc import AFQMC from ipie.systems.generic import Generic -from ipie.trial_wavefunction.particle_hole import ParticleHoleNonChunked +from ipie.trial_wavefunction.particle_hole import ParticleHole from ipie.utils.from_pyscf import gen_ipie_input_from_pyscf_chk +from ipie.walkers.uhf_walkers import UHFWalkersParticleHole +from ipie.utils.mpi import MPIHandler -nocca = 4 -noccb = 2 +try: + from mpi4py import MPI + comm = MPI.COMM_WORLD +except ImportError: + comm = None -mol = gto.M( - atom=[("N", 0, 0, 0), ("N", (0, 0, 3.0))], - basis="ccpvdz", - verbose=3, - spin=nocca - noccb, - unit="Bohr", -) -mf = scf.RHF(mol) -mf.chkfile = "scf.chk" -ehf = mf.kernel() -M = 6 -N = 6 -mc = mcscf.CASSCF(mf, M, N) -mc.chkfile = "scf.chk" -e_tot, e_cas, fcivec, mo, mo_energy = mc.kernel() -coeff, occa, occb = zip( - *fci.addons.large_ci(fcivec, M, (nocca, noccb), tol=1e-8, return_strs=False) -) -# Need to write wavefunction to checkpoint file. -with h5py.File("scf.chk", "r+") as fh5: - fh5["mcscf/ci_coeffs"] = coeff - fh5["mcscf/occs_alpha"] = occa - fh5["mcscf/occs_beta"] = occb +from ipie.config import config +# config.update_option("use_gpu", True) # enable GPU +config.update_option("use_gpu", False) # disable GPU + +comm_size = comm.size if comm is not None else 1 +num_walkers = 640 // comm_size + + +if comm is None or comm.rank == 0: + nocca = 4 + noccb = 2 + + mol = gto.M( + atom=[("N", 0, 0, 0), ("N", (0, 0, 3.0))], + basis="ccpvdz", + verbose=3, + spin=nocca - noccb, + unit="Bohr", + ) + mf = scf.RHF(mol) + mf.chkfile = "scf.chk" + ehf = mf.kernel() + M = 6 + N = 6 + mc = mcscf.CASSCF(mf, M, N) + mc.chkfile = "scf.chk" + e_tot, e_cas, fcivec, mo, mo_energy = mc.kernel() + coeff, occa, occb = zip( + *fci.addons.large_ci(fcivec, M, (nocca, noccb), tol=1e-8, return_strs=False) + ) + # Need to write wavefunction to checkpoint file. + with h5py.File("scf.chk", "r+") as fh5: + fh5["mcscf/ci_coeffs"] = coeff + fh5["mcscf/occs_alpha"] = occa + fh5["mcscf/occs_beta"] = occb + + gen_ipie_input_from_pyscf_chk("scf.chk", mcscf=True) -gen_ipie_input_from_pyscf_chk("scf.chk", mcscf=True) mol_nelec = [8, 6] with h5py.File("hamiltonian.h5") as fa: @@ -59,7 +81,7 @@ occa = fh5["occ_alpha"][:] occb = fh5["occ_beta"][:] wavefunction = (coeff, occa, occb) -trial = ParticleHoleNonChunked( +trial = ParticleHole( wavefunction, mol_nelec, num_basis, @@ -70,12 +92,26 @@ trial.build() trial.half_rotate(ham) +initial_walker = numpy.hstack([trial.psi0a, trial.psi0b]) +random_perturbation = numpy.random.random(initial_walker.shape) +initial_walker = initial_walker + random_perturbation +initial_walker, _ = numpy.linalg.qr(initial_walker) +walkers = UHFWalkersParticleHole( + initial_walker, + mol_nelec[0], + mol_nelec[1], + num_basis, + num_walkers, + MPIHandler(), +) +walkers.build(trial) afqmc_msd = AFQMC.build( mol_nelec, ham, trial, - num_walkers=10, + walkers=walkers, + num_walkers=num_walkers, num_steps_per_block=25, num_blocks=10, timestep=0.005, @@ -84,5 +120,5 @@ pop_control_freq=5, verbose=True, ) -# afqmc_msd.run() -# afqmc_msd.finalise(verbose=True) +afqmc_msd.run() +afqmc_msd.finalise(verbose=True) diff --git a/ipie/estimators/greens_function_multi_det.py b/ipie/estimators/greens_function_multi_det.py index 8c381d8a..0c7b7e00 100644 --- a/ipie/estimators/greens_function_multi_det.py +++ b/ipie/estimators/greens_function_multi_det.py @@ -1272,8 +1272,8 @@ def greens_function_multi_det_wicks_opt_gpu(walker_batch, trial, build_full=Fals nbasis = walker_batch.Ga.shape[-1] - walker_batch.Ga.fill(0.0 + 0.0j) - walker_batch.Gb.fill(0.0 + 0.0j) + walker_batch.Ga = xp.zeros_like(walker_batch.Ga) + walker_batch.Gb = xp.zeros_like(walker_batch.Gb) # Build reference Green's functions and overlaps # Note abuse of naming convention this is really theta for the reference @@ -1317,8 +1317,8 @@ def greens_function_multi_det_wicks_opt_gpu(walker_batch, trial, build_full=Fals walker_batch.G0b = G0b walker_batch.Q0a = xp.eye(nbasis)[None, :] - G0a walker_batch.Q0b = xp.eye(nbasis)[None, :] - G0b - walker_batch.CIa.fill(0.0 + 0.0j) - walker_batch.CIb.fill(0.0 + 0.0j) + walker_batch.CIa = xp.zeros_like(walker_batch.CIa) + walker_batch.CIb = xp.zeros_like(walker_batch.CIb) dets_a_full, dets_b_full = compute_determinants_batched( walker_batch.Ghalfa, walker_batch.Ghalfb, trial diff --git a/ipie/walkers/walkers_dispatch.py b/ipie/walkers/walkers_dispatch.py index 1fc3a628..606d63d9 100644 --- a/ipie/walkers/walkers_dispatch.py +++ b/ipie/walkers/walkers_dispatch.py @@ -45,9 +45,15 @@ def get_initial_walker(trial: TrialWavefunctionBase) -> numpy.ndarray: num_dets = 1 elif isinstance(trial, ParticleHole): initial_walker = numpy.hstack([trial.psi0a, trial.psi0b]) + # random_walker = numpy.random.random(initial_walker.shape) + # initial_walker = initial_walker + random_walker + # initial_walker, _ = numpy.linalg.qr(initial_walker) num_dets = trial.num_dets elif isinstance(trial, ParticleHoleNonChunked): initial_walker = numpy.hstack([trial.psi0a, trial.psi0b]) + # random_walker = numpy.random.random(initial_walker.shape) + # initial_walker = initial_walker + random_walker + # initial_walker, _ = numpy.linalg.qr(initial_walker) num_dets = trial.num_dets elif isinstance(trial, NOCI): initial_walker = trial.psi[0].copy()