add lno module (#91)

Author: fishjojo

.github/workflows/install_pyscf.sh

@@ -4,5 +4,6 @@ python -m pip cache purge
 pip install wheel
 pip install pytest
 pip install pytest-cov
-pip install jax
-pip install pyscf
+pip install 'numpy<2.4'
+pip install 'jax<0.9'
+pip install 'pyscf<2.12.0'

.github/workflows/run_test.sh

@@ -7,7 +7,7 @@ export OMP_NUM_THREADS=1
 
 coverage erase
 
-MODULES=("scipy" "gto" "cc" "fci" "gw" "mp" "tdscf" "lo" "pbc")
+MODULES=("scipy" "gto" "cc" "fci" "gw" "mp" "tdscf" "lo" "pbc" "lno")
 
 FAILED=0
 

examples/lno/00-lno_mp2.py

@@ -0,0 +1,50 @@
+'''LNO-MP2 energy and gradient.
+'''
+import jax
+from pyscfad import gto, scf
+from pyscfad.mp import dfmp2
+from pyscfad import config
+from pyscfad.lno import LNOMP2
+
+# use optimized backpropagation
+config.update('pyscfad_moleintor_opt', True)
+config.update('pyscfad_scf_implicit_diff', True)
+config.update('pyscfad_ccsd_implicit_diff', True)
+
+atom = 'water_dimer.xyz'
+basis = 'ccpvdz'
+frozen = None
+
+mol = gto.Mole(atom=atom, basis=basis)
+mol.verbose = 4
+mol.build(trace_exp=False, trace_ctr_coeff=False)
+
+# canonical MP2
+def mp2_energy(mol):
+    mf = scf.RHF(mol).density_fit()
+    mf.kernel()
+
+    mymp = dfmp2.MP2(mf, frozen=frozen)
+    mymp.kernel()
+    return mymp.e_tot
+
+e_mp2, jac_mp2 = jax.value_and_grad(mp2_energy)(mol)
+
+# LNO-MP2
+thresh = 1e-4
+def lno_mp2_energy(mol):
+    mf = scf.RHF(mol).density_fit()
+    ehf = mf.kernel()
+
+    mfcc = LNOMP2(mf, thresh=thresh, frozen=frozen)
+    mfcc.thresh_occ = thresh
+    mfcc.thresh_vir = thresh
+    mfcc.lo_type = 'iao'
+    mfcc.kernel(frag_lolist=None)
+    return ehf + mfcc.e_corr
+
+e_lno_mp2, jac_lno_mp2 = jax.value_and_grad(lno_mp2_energy)(mol)
+
+print(e_mp2, e_lno_mp2, e_mp2-e_lno_mp2)
+print(jac_lno_mp2.coords)
+print(abs(jac_lno_mp2.coords - jac_mp2.coords).max())

examples/lno/01-lno_ccsd_t.py

@@ -0,0 +1,51 @@
+'''LNO-CCSD(T) energy and gradient
+'''
+import jax
+from pyscfad import gto, scf, mp, cc
+from pyscfad.cc import dfccsd
+from pyscfad import config
+from pyscfad.lno import LNOCCSD
+
+config.update('pyscfad_moleintor_opt', True)
+config.update('pyscfad_scf_implicit_diff', True)
+config.update('pyscfad_ccsd_implicit_diff', True)
+
+atom = 'water_dimer.xyz'
+basis = 'ccpvdz'
+frozen = 2
+
+mol = gto.Mole(atom=atom, basis=basis)
+mol.verbose = 4
+mol.build(trace_exp=False, trace_ctr_coeff=False)
+
+# canonical CCSD(T)
+mf = scf.RHF(mol).density_fit()
+mf.kernel()
+mycc = dfccsd.RCCSD(mf, frozen=frozen)
+eris = mycc.ao2mo()
+mycc.kernel(eris=eris)
+et = mycc.ccsd_t(eris=eris)
+
+# LNO-CCSD(T)
+thresh = 1e-4
+def energy(mol):
+    mf = scf.RHF(mol).density_fit()
+    ehf = mf.kernel()
+
+    mmp = mp.dfmp2.MP2(mf, frozen=frozen)
+    mmp.kernel(with_t2=False)
+
+    mfcc = LNOCCSD(mf, thresh=thresh, frozen=frozen)
+    mfcc.thresh_occ = thresh
+    mfcc.thresh_vir = thresh
+    mfcc.lo_type = 'iao'
+    mfcc.no_type = 'ie'
+    mfcc.ccsd_t = True
+    mfcc.kernel(frag_lolist=None)
+
+    ecc_pt2corrected = mfcc.e_corr_pt2corrected(mmp.e_corr)
+    return ehf + ecc_pt2corrected
+
+e, jac = jax.value_and_grad(energy)(mol)
+print(e, mycc.e_tot+et, e-(mycc.e_tot+et))
+print(jac.coords)

examples/lno/11-mpi_lno_ccsd_t.py

@@ -0,0 +1,69 @@
+'''LNO-CCSD(T) with MPI parallelization.
+
+run with:
+mpirun -n 2 python 11-mpi_lno_ccsd_t.py
+'''
+from mpi4py import MPI
+import jax
+import numpy
+from pyscfad import gto, scf, mp
+from pyscfad.cc import dfccsd
+from pyscfad import config
+from pyscfad.lno.ccsd_mpi import LNOCCSD
+
+config.update('pyscfad_moleintor_opt', True)
+config.update('pyscfad_scf_implicit_diff', True)
+config.update('pyscfad_ccsd_implicit_diff', True)
+
+atom = 'water_dimer.xyz'
+basis = 'ccpvdz'
+
+mol = gto.Mole(atom=atom, basis=basis)
+mol.verbose = 4
+mol.build(trace_exp=False, trace_ctr_coeff=False)
+
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+frozen = 2
+thresh_occ = 1e-3
+thresh_vir = 1e-4
+def energy(mol):
+    mf = scf.RHF(mol).density_fit()
+    ehf = mf.kernel()
+
+    mfcc = LNOCCSD(mf, frozen=frozen)
+    mfcc.thresh_occ = thresh_occ
+    mfcc.thresh_vir = thresh_vir
+    mfcc.lo_type = 'iao'
+    mfcc.ccsd_t = True
+    mfcc.kernel(frag_lolist=None)
+
+    if rank == 0:
+        mmp = mp.dfmp2.MP2(mf, frozen=frozen)
+        mmp.kernel(with_t2=False)
+        ecc_pt2corrected = mfcc.e_corr_pt2corrected(mmp.e_corr)
+        etot = ehf + ecc_pt2corrected
+    else:
+        etot = mfcc.e_corr - mfcc.e_corr_pt2
+    return etot
+
+e, jac = jax.value_and_grad(energy)(mol)
+e = numpy.asarray(e)
+grad = numpy.asarray(jac.coords)
+
+if rank == 0:
+    etot = numpy.zeros_like(e)
+    grad_tot = numpy.zeros_like(grad)
+else:
+    etot = None
+    grad_tot = None
+
+comm.Reduce([e, MPI.DOUBLE], [etot, MPI.DOUBLE],
+            op=MPI.SUM, root=0)
+
+comm.Reduce([grad, MPI.DOUBLE], [grad_tot, MPI.DOUBLE],
+            op=MPI.SUM, root=0)
+
+if rank == 0:
+    print(f'LNO-CCSD(T) energy: {etot}')
+    print(f'LNO-CCSD(T) gradient:\n{grad_tot}')

examples/lno/water_dimer.xyz

@@ -0,0 +1,8 @@
+6
+
+O   -1.485163346097   -0.114724564047    0.000000000000
+H   -1.868415346097    0.762298435953    0.000000000000
+H   -0.533833346097    0.040507435953    0.000000000000
+O    1.416468653903    0.111264435953    0.000000000000
+H    1.746241653903   -0.373945564047   -0.758561000000
+H    1.746241653903   -0.373945564047    0.758561000000

pyscfad/lno/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2023-2026 The PySCFAD Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+LNO methods
+"""
+from pyscfad.lno.mp2 import LNOMP2
+from pyscfad.lno.ccsd import LNOCCSD, LNOCCSD_T

pyscfad/lno/_checkpointed.py

@@ -0,0 +1,45 @@
+# Copyright 2023-2026 The PySCFAD Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+from pyscfad import numpy as np
+
+def make_mp2_rdm1_ie(Lia, Ljb, eia, ejb):
+    naux1, nocc1, nvir1 = Lia.shape
+    naux2, nocc2, nvir2 = Ljb.shape
+    assert naux1 == naux2
+    assert nvir1 == nvir2
+    naux = naux1
+    nvir = nvir1
+
+    @jax.checkpoint
+    def fn(carry, x):
+        dmvv, dmoo = carry
+        La, ea = x
+        buf = np.dot(La.T, Ljb.reshape(naux,-1)).reshape(nvir, nocc2, nvir)
+        t2i = buf / (ea[:,None,None] + ejb[None,:,:])
+        dmvv +=  np.dot(t2i.reshape(nvir, -1), t2i.reshape(nvir, -1).T)
+        dmvv -= .5 * np.einsum('ajc,cjb->ab', t2i, t2i)
+        dmvv += np.dot(t2i.reshape(-1,nvir).T, t2i.reshape(-1,nvir))
+        dmvv -= .5 * np.einsum('cja,bjc->ab', t2i, t2i)
+
+        dmoo += np.einsum('aib,ajb->ij', t2i, t2i)
+        dmoo -= .5 * np.einsum('aib,bja->ij', t2i, t2i)
+        dmoo += np.einsum('bia,bja->ij', t2i, t2i)
+        dmoo -= .5 * np.einsum('bia,ajb->ij', t2i, t2i)
+        return (dmvv, dmoo), None
+
+    (dmvv, dmoo), _ = jax.lax.scan(fn, (np.zeros((nvir, nvir)), np.zeros((nocc2,nocc2))),
+                                   (Lia.transpose(1,0,2), eia))
+    return dmvv, dmoo