Skip to content

Commit 0286eaa

Browse files
MatthewRHermesMatthewRHermes
authored
Merge pull request #189 from ankit76/afqmc
AFQMC
2 parents 919f7a3 + 5a58e11 commit 0286eaa

77 files changed

Lines changed: 16755 additions & 0 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/run_ci.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ python -m pip install --upgrade pip
1212
pip install "scipy<1.16"
1313
pip install pytest
1414
pip install .
15+
pip install -U jax
1516

1617
pip install trexio
1718
pip install mcfun

‎examples/afqmc/00-rhf.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Requires a separate JAX installation.
2+
# Phaseless AFQMC with an RHF trial state.
3+
4+
from pyscf import afqmc, gto, scf
5+
6+
mol = gto.M(
7+
atom="""
8+
O 0.0000000000 0.0000000000 0.0000000000
9+
H 0.9562300000 0.0000000000 0.0000000000
10+
H -0.2353791634 0.9268076728 0.0000000000
11+
""",
12+
basis="6-31g",
13+
verbose=3,
14+
)
15+
16+
mf = scf.RHF(mol)
17+
mf.kernel()
18+
19+
af = afqmc.AFQMC(mf)
20+
af.verbose = 4
21+
mean, err = af.kernel()

‎examples/afqmc/01-cisd.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Requires a separate JAX installation.
2+
# Phaseless AFQMC with a CISD trial state.
3+
4+
from pyscf import afqmc, cc, gto, scf
5+
6+
mol = gto.M(
7+
atom="""
8+
O 0.0000000000 0.0000000000 0.0000000000
9+
H 0.9562300000 0.0000000000 0.0000000000
10+
H -0.2353791634 0.9268076728 0.0000000000
11+
""",
12+
basis="6-31g",
13+
verbose=3,
14+
)
15+
16+
# RHF
17+
mf = scf.RHF(mol)
18+
mf.kernel()
19+
20+
mycc = cc.CCSD(mf)
21+
mycc.frozen = 1 # freeze O 1s core
22+
mycc.kernel()
23+
et = mycc.ccsd_t() # for comparison
24+
print(f"CCSD(T) total energy: {mycc.e_tot + et}")
25+
26+
af = afqmc.AFQMC(mycc)
27+
af.n_walkers = 100
28+
af.n_eql_blocks = 10
29+
af.n_blocks = 100
30+
af.seed = 548280
31+
af.verbose = 4
32+
mean, err = af.kernel()

‎examples/afqmc/02-rhf-fp.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
# Requires a separate JAX installation.
2+
# Free projection AFQMC with an RHF state on both sides.
3+
4+
from pyscf import afqmc, gto, scf
5+
6+
mol = gto.M(
7+
atom="""
8+
O 0.0000000000 0.0000000000 0.0000000000
9+
H 0.9562300000 0.0000000000 0.0000000000
10+
H -0.2353791634 0.9268076728 0.0000000000
11+
""",
12+
basis="6-31g",
13+
verbose=3,
14+
)
15+
16+
mf = scf.RHF(mol)
17+
mf.kernel()
18+
19+
af = afqmc.AFQMCFP(mf)
20+
af.dt = 0.1
21+
af.n_prop_steps = 100
22+
af.n_blocks = 5
23+
af.ene0 = mf.e_tot
24+
af.n_traj = 10
25+
af.seed = 6
26+
af.n_walkers = 200
27+
af.walker_kind = "restricted"
28+
mean, err = af.kernel()

‎examples/afqmc/03-lno-rhfafqmc.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
# Requires a separate JAX installation.
2+
3+
import os
4+
5+
os.environ["OMP_NUM_THREADS"] = (
6+
"1" # For reproducibility. LNO makes slightly different orbitals every time.
7+
)
8+
import sys
9+
10+
import numpy as np
11+
from pyscf.data.elements import chemcore
12+
13+
from pyscf import gto, mp, scf
14+
from pyscf.afqmc.lnoafqmc import LNOAFQMC, prep_local_orbitals
15+
from pyscf.lib import logger
16+
17+
log = logger.Logger(sys.stdout, 6)
18+
19+
20+
# water dimer
21+
atom = """
22+
O -1.485163346097 -0.114724564047 0.000000000000
23+
H -1.868415346097 0.762298435953 0.000000000000
24+
H -0.533833346097 0.040507435953 0.000000000000
25+
O 1.416468653903 0.111264435953 0.000000000000
26+
H 1.746241653903 -0.373945564047 -0.758561000000
27+
H 1.746241653903 -0.373945564047 0.758561000000
28+
"""
29+
30+
basis = "631g"
31+
32+
mol = gto.M(atom=atom, basis=basis)
33+
mol.verbose = 3
34+
frozen = chemcore(mol)
35+
36+
mf = scf.RHF(mol).density_fit()
37+
mf.kernel() # type: ignore
38+
39+
# canonical
40+
mmp = mp.MP2(mf, frozen=frozen)
41+
mmp.kernel()
42+
efull_mp2 = mmp.e_corr
43+
44+
lo_coeff, frag_lolist = prep_local_orbitals(mf, frozen=frozen)
45+
# frag_lolist = [[0]] # One can run a particular fragment like this
46+
# LNO-AFQMC calculation: here we can scan over a list of thresholds
47+
mcc = LNOAFQMC(mf, lo_coeff, frag_lolist, frozen=frozen).set(verbose=5)
48+
mcc.seed = 1234
49+
mcc.n_blocks = 100
50+
mcc.n_walkers = 200
51+
52+
gamma = 10 # thresh_occ / thresh_vir
53+
threshs = np.asarray([1e-5])
54+
elno_afqmc_uncorr = np.zeros_like(threshs)
55+
stocherr_afqmc = np.zeros_like(threshs)
56+
elno_mp2 = np.zeros_like(threshs)
57+
elno_afqmc = np.zeros_like(threshs)
58+
for i, thresh in enumerate(threshs):
59+
mcc.lno_thresh = [thresh * gamma, thresh]
60+
mcc.kernel()
61+
stocherr_afqmc[i] = mcc.afqmc_error_ecorr
62+
elno_afqmc_uncorr[i] = mcc.e_corr_afqmc
63+
elno_mp2[i] = mcc.e_corr_pt2
64+
elno_afqmc[i] = mcc.e_corr_afqmc - elno_mp2[i] + efull_mp2
65+
66+
log.info("")
67+
for i, thresh in enumerate(threshs):
68+
e0 = elno_afqmc_uncorr[i]
69+
err = stocherr_afqmc[i]
70+
e1 = elno_afqmc[i]
71+
log.info(
72+
"thresh = %.3e E_corr(LNO-AFQMC) = %.15g +/- %.15g E_corr(LNO-AFQMC+∆PT2) = %.15g +/- %.15g",
73+
thresh,
74+
e0,
75+
err,
76+
e1,
77+
err,
78+
)

‎examples/afqmc/04-ucisd.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Requires a separate JAX installation.
2+
# Phaseless AFQMC with a UCISD trial state.
3+
4+
from pyscf import afqmc, cc, gto, scf
5+
6+
mol = gto.M(
7+
atom="""
8+
N -1.67119571 -1.44021737 0.00000000
9+
H -2.12619571 -0.65213425 0.00000000
10+
H -0.76119571 -1.44021737 0.00000000
11+
""",
12+
spin=1,
13+
basis="6-31g",
14+
verbose=3,
15+
)
16+
17+
mf = scf.UHF(mol)
18+
mf.kernel()
19+
20+
mo1 = mf.stability()[0]
21+
dm1 = mf.make_rdm1(mo1, mf.mo_occ)
22+
mf = mf.run(dm1)
23+
mf.stability()
24+
25+
mycc = cc.UCCSD(mf)
26+
mycc.kernel()
27+
28+
af = afqmc.AFQMC(mycc)
29+
af.n_walkers = 100
30+
af.n_blocks = 100
31+
af.verbose = 4
32+
mean, err = af.kernel()

‎examples/afqmc/05-cisd-fp.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Requires a separate JAX installation.
2+
# Free projection AFQMC with a CISD initial state.
3+
4+
from pyscf import afqmc, cc, gto, scf
5+
6+
mol = gto.M(
7+
atom="""
8+
O 0.0000000000 0.0000000000 0.0000000000
9+
H 0.9562300000 0.0000000000 0.0000000000
10+
H -0.2353791634 0.9268076728 0.0000000000
11+
""",
12+
basis="6-31g",
13+
verbose=3,
14+
)
15+
16+
# RHF
17+
mf = scf.RHF(mol)
18+
mf.kernel()
19+
20+
mycc = cc.CCSD(mf)
21+
mycc.kernel()
22+
23+
af = afqmc.AFQMCFP(mycc)
24+
af.dt = 0.1
25+
af.n_prop_steps = 10
26+
af.n_blocks = 5
27+
af.ene0 = mycc.e_tot
28+
af.n_traj = 10
29+
af.seed = 6
30+
af.n_walkers = 200
31+
af.walker_kind = "restricted"
32+
af.verbose = 4
33+
mean, err = af.kernel()

‎examples/afqmc/README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# AFQMC Examples
2+
3+
AFQMC currently requires Python 3.10 or newer and [JAX](https://github.com/jax-ml/jax).
4+
For CPU runs, install JAX with `pip install -U jax`.
5+
For NVIDIA GPU runs, install `pip install -U "jax[cuda12]"` or
6+
`pip install -U "jax[cuda13]"`.
7+
If preferred, we also provide convenience extras:
8+
`pyscf-forge[afqmc-cpu]`, `pyscf-forge[afqmc-cuda12]`, and
9+
`pyscf-forge[afqmc-cuda13]`.

‎examples/afqmc/staging/00-stage-cisd.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Requires a separate JAX installation.
2+
3+
from pyscf import cc, gto, scf
4+
5+
from pyscf import afqmc
6+
7+
8+
STAGED_PATH = "h2o_afqmc.h5"
9+
10+
mol = gto.M(
11+
atom="""
12+
O 0.0000000000 0.0000000000 0.0000000000
13+
H 0.9562300000 0.0000000000 0.0000000000
14+
H -0.2353791634 0.9268076728 0.0000000000
15+
""",
16+
basis="6-31g",
17+
verbose=3,
18+
)
19+
20+
mf = scf.RHF(mol)
21+
mf.kernel()
22+
23+
mycc = cc.CCSD(mf)
24+
mycc.frozen = 1
25+
mycc.kernel()
26+
27+
af = afqmc.AFQMC(mycc)
28+
af.save_staged(STAGED_PATH)
29+
30+
print(f"Wrote staged AFQMC inputs with a CISD trial to {STAGED_PATH}.")

‎examples/afqmc/staging/01-run-staged.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Requires a separate JAX installation.
2+
3+
from pyscf import afqmc
4+
5+
6+
STAGED_PATH = "h2o_afqmc.h5"
7+
8+
af = afqmc.AFQMC.from_staged(STAGED_PATH)
9+
# The staged file can come from an expensive CPU-side CCSD calculation.
10+
af.n_walkers = 20
11+
af.n_eql_blocks = 10
12+
af.n_blocks = 200
13+
af.seed = 42
14+
15+
mean, err = af.kernel()
16+
print(f"AFQMC energy: {mean:.10f} +/- {err:.10f}")

0 commit comments

Comments
 (0)