bugfix of gpu msd-afqmc; fix tutorial

examples/02-multi_determinant/run_afqmc.py

@@ -1,41 +1,63 @@
+"""
+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
 
 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)

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

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()