bugfix of gpu msd-afqmc; fix tutorial

examples/02-multi_determinant/run_afqmc.py

@@ -1,41 +1,61 @@
+"""
+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
 
-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 +79,7 @@
     occa = fh5["occ_alpha"][:]
     occb = fh5["occ_beta"][:]
 wavefunction = (coeff, occa, occb)
-trial = ParticleHoleNonChunked(
+trial = ParticleHole(
     wavefunction,
     mol_nelec,
     num_basis,
@@ -75,7 +95,7 @@
     mol_nelec,
     ham,
     trial,
-    num_walkers=10,
+    num_walkers=num_walkers,
     num_steps_per_block=25,
     num_blocks=10,
     timestep=0.005,
@@ -84,5 +104,5 @@
     pop_control_freq=5,
     verbose=True,
 )
-# afqmc_msd.run()
-# afqmc_msd.finalise(verbose=True)
+afqmc_msd.run()
+afqmc_msd.finalise(verbose=True)

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