merge qmmm (#218)

* minor * minor * finish qmmm pbc * optimization in pbc qmmm * optimization in qmmm pbc * optimization in qmmm pbc * optimization in qmmm pbc * optimization in qmmm pbc * error estimation using octupole in qmmm pbc * add more checks in qmmm pbc * minor * allow parent class h1 run in background in qmmm * minor * allow zetas in qmmm.pbc functions * move qmmm_gas into qmmm.pbc.itrf * minor fix * minor * coding style * change cp.einsum to contract * add comments * keep changes minimal * add example and test * use get_j_int3c2e_pass2 in qmmm.pbc get_hcore * coding style * change cp.einsum to contract * minor * correct some docstrings * add test for error estimate * change some contract into broadcasting * avoid multiple gpus tasks run in background * change some contract into broadcasting * reduce duplication * reduce duplication * coding style * remove h1_on_cpu * remove async * revert unit test --------- Co-authored-by: Chenghan Li <chhli@login31.chn> Co-authored-by: xiaojie.wu <xiaojie.wu@bytedance.com> Co-authored-by: Xiaojie Wu <wxj6000@gmail.com>
pyscf · Oct 18, 2024 · c96f84c · c96f84c
1 parent 49f0f65
commit c96f84c
Show file tree

Hide file tree

Showing 10 changed files with 2,112 additions and 59 deletions.
diff --git a/examples/23-qmmm_pbc.py b/examples/23-qmmm_pbc.py
@@ -0,0 +1,52 @@
+# Copyright 2024 The GPU4PySCF Authors. All Rights Reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import pyscf
+from gpu4pyscf.scf import hf
+from gpu4pyscf.qmmm.pbc import itrf, mm_mole
+
+import numpy as np
+
+atom ='''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+
+mol = pyscf.M(atom=atom, basis='3-21g', max_memory=40000)
+
+mol.verbose = 4
+mf_GPU = hf.RHF(mol)
+
+# Add mm charges:
+# -0.8 charge at 1.0, 2.0,-1.0
+#  0.8 charge at 3.0, 4.0, 5.0
+# Place the QM and MM atoms in a box of 12 A
+# Real-space cutoff for Ewald set to 8 A
+# Exact MM potential computed for MM charges within 6 A of QM geometric center
+mf_GPU = itrf.add_mm_charges(mf_GPU, [[1,2,-1],[3,4,5]], np.eye(3)*12, [-0.8,0.8], [0.8,1.2], rcut_ewald=8, rcut_hcore=6)
+
+# Compute Energy of QM-QM and QM-MM but NOT MM-MM
+e_dft = mf_GPU.kernel()
+print(f"total energy = {e_dft}")
+
+# Compute Gradient
+g = mf_GPU.nuc_grad_method()
+g.max_memory = 40000
+g.auxbasis_response = True
+# energy gradient w.r.t. QM atom positions
+g_dft = g.kernel()
+# energy gradient w.r.t. MM atom positions
+g_mm = g.grad_nuc_mm() + g.grad_hcore_mm(mf_GPU.make_rdm1()) + g.de_ewald_mm
diff --git a/gpu4pyscf/df/df_jk.py b/gpu4pyscf/df/df_jk.py
@@ -24,7 +24,7 @@
 from pyscf.scf import dhf
 from pyscf.df import df_jk, addons
 from gpu4pyscf.lib import logger
-from gpu4pyscf.lib.cupy_helper import contract, take_last2d, transpose_sum, load_library, get_avail_mem
+from gpu4pyscf.lib.cupy_helper import contract, take_last2d, transpose_sum, load_library, get_avail_mem, empty_mapped
 from gpu4pyscf.dft import rks, uks, numint
 from gpu4pyscf.scf import hf, uhf
 from gpu4pyscf.df import df, int3c2e
@@ -52,11 +52,11 @@ def build_df():
             rsh_df.build(omega=omega)
         return
 
+    mf.h1e = cupy.asarray(mf.get_hcore(mf.mol))
+    mf.s1e = cupy.asarray(mf.get_ovlp(mf.mol))
     # pre-compute h1e and s1e and cderi for async workflow
     with lib.call_in_background(build_df) as build:
         build()
-        mf.s1e = cupy.asarray(mf.get_ovlp(mf.mol))
-        mf.h1e = cupy.asarray(mf.get_hcore(mf.mol))
         # for DFT object
         if hasattr(mf, '_numint'):
             ni = mf._numint

diff --git a/gpu4pyscf/grad/rhf.py b/gpu4pyscf/grad/rhf.py
@@ -280,37 +280,28 @@ def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
     dm0 = mf.make_rdm1(mo_coeff, mo_occ)
     dme0 = mf_grad.make_rdm1e(mo_energy, mo_coeff, mo_occ)
 
-    # CPU tasks are executed on background
-    def calculate_h1e(h1_gpu, s1_gpu):
-        # (\nabla i | hcore | j) - (\nabla i | j)
-        h1_cpu = mf_grad.get_hcore(mol)
-        s1_cpu = mf_grad.get_ovlp(mol)
-        h1_gpu[:] = cupy.asarray(h1_cpu)
-        s1_gpu[:] = cupy.asarray(s1_cpu)
-        return
+    # (\nabla i | hcore | j) - (\nabla i | j)
+    h1 = cupy.asarray(mf_grad.get_hcore(mol))
+    s1 = cupy.asarray(mf_grad.get_ovlp(mol))
 
-    h1 = cupy.empty([3, dm0.shape[0], dm0.shape[1]])
-    s1 = cupy.empty([3, dm0.shape[0], dm0.shape[1]])
-    with lib.call_in_background(calculate_h1e) as calculate_hs:
-        calculate_hs(h1, s1)
-        # (i | \nabla hcore | j)
-        t3 = log.init_timer()
-        dh1e = int3c2e.get_dh1e(mol, dm0)
+    # (i | \nabla hcore | j)
+    t3 = log.init_timer()
+    dh1e = int3c2e.get_dh1e(mol, dm0)
 
-        if mol.has_ecp():
-            dh1e += get_dh1e_ecp(mol, dm0)
-        t3 = log.timer_debug1('gradients of h1e', *t3)
+    if mol.has_ecp():
+        dh1e += get_dh1e_ecp(mol, dm0)
+    t3 = log.timer_debug1('gradients of h1e', *t3)
 
-        dvhf = mf_grad.get_veff(mol, dm0)
-        log.timer_debug1('gradients of veff', *t3)
-        log.debug('Computing Gradients of NR-HF Coulomb repulsion')
+    dvhf = mf_grad.get_veff(mol, dm0)
+    log.timer_debug1('gradients of veff', *t3)
+    log.debug('Computing Gradients of NR-HF Coulomb repulsion')
 
-        dm0 = tag_array(dm0, mo_coeff=mo_coeff, mo_occ=mo_occ)
-        extra_force = cupy.zeros((len(atmlst),3))
-        for k, ia in enumerate(atmlst):
-            extra_force[k] += mf_grad.extra_force(ia, locals())
+    dm0 = tag_array(dm0, mo_coeff=mo_coeff, mo_occ=mo_occ)
+    extra_force = cupy.zeros((len(atmlst),3))
+    for k, ia in enumerate(atmlst):
+        extra_force[k] += mf_grad.extra_force(ia, locals())
 
-        log.timer_debug1('gradients of 2e part', *t3)
+    log.timer_debug1('gradients of 2e part', *t3)
 
     dh = contract('xij,ij->xi', h1, dm0)
     ds = contract('xij,ij->xi', s1, dme0)

diff --git a/gpu4pyscf/grad/uhf.py b/gpu4pyscf/grad/uhf.py
@@ -22,7 +22,7 @@
 from pyscf.grad import uhf
 from pyscf.grad import rhf as rhf_grad_cpu
 from gpu4pyscf.lib.cupy_helper import load_library
-from gpu4pyscf.lib.cupy_helper import tag_array, contract
+from gpu4pyscf.lib.cupy_helper import tag_array, contract, empty_mapped
 from gpu4pyscf.df import int3c2e      #TODO: move int3c2e to out of df
 from gpu4pyscf.lib import logger
 from gpu4pyscf.scf.int4c2e import _VHFOpt
@@ -278,34 +278,26 @@ def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
         atmlst = range(mol.natm)
     aoslices = mol.aoslice_by_atom()
     de = cupy.zeros((len(atmlst),3))
-
-    def calculate_h1e(h1_gpu, s1_gpu):
-        # (\nabla i | hcore | j) - (\nabla i | j)
-        h1_cpu = mf_grad.get_hcore(mol)
-        s1_cpu = mf_grad.get_ovlp(mol)
-        h1_gpu[:] = cupy.asarray(h1_cpu)
-        s1_gpu[:] = cupy.asarray(s1_cpu)
-        return
-
-    h1 = cupy.empty([3, dm0.shape[1], dm0.shape[2]])
-    s1 = cupy.empty([3, dm0.shape[1], dm0.shape[2]])
-    with lib.call_in_background(calculate_h1e) as calculate_hs:
-        calculate_hs(h1, s1)
-        # (i | \nabla hcore | j)
-        t3 = log.init_timer()
-        dh1e = int3c2e.get_dh1e(mol, dm0_sf)
-
-        log.timer_debug1("get_dh1e", *t3)
-        if mol.has_ecp():
-            dh1e += rhf_grad.get_dh1e_ecp(mol, dm0_sf)
-        t1 = log.timer_debug1('gradients of h1e', *t0)
-        log.debug('Computing Gradients of NR-HF Coulomb repulsion')
-        dvhf = mf_grad.get_veff(mol, dm0)
-
-        extra_force = cupy.zeros((len(atmlst),3))
-        for k, ia in enumerate(atmlst):
-            extra_force[k] += mf_grad.extra_force(ia, locals())
-        log.timer_debug1('gradients of 2e part', *t1)
+
+    # (\nabla i | hcore | j) - (\nabla i | j)
+    h1 = cupy.asarray(mf_grad.get_hcore(mol))
+    s1 = cupy.asarray(mf_grad.get_ovlp(mol))
+
+    # (i | \nabla hcore | j)
+    t3 = log.init_timer()
+    dh1e = int3c2e.get_dh1e(mol, dm0_sf)
+
+    log.timer_debug1("get_dh1e", *t3)
+    if mol.has_ecp():
+        dh1e += rhf_grad.get_dh1e_ecp(mol, dm0_sf)
+    t1 = log.timer_debug1('gradients of h1e', *t0)
+    log.debug('Computing Gradients of NR-HF Coulomb repulsion')
+    dvhf = mf_grad.get_veff(mol, dm0)
+
+    extra_force = cupy.zeros((len(atmlst),3))
+    for k, ia in enumerate(atmlst):
+        extra_force[k] += mf_grad.extra_force(ia, locals())
+    log.timer_debug1('gradients of 2e part', *t1)
 
     dh = contract('xij,ij->xi', h1, dm0_sf)
     ds = contract('xij,ij->xi', s1, dme0_sf)

diff --git a/gpu4pyscf/lib/tests/test_to_gpu.py b/gpu4pyscf/lib/tests/test_to_gpu.py
@@ -129,4 +129,5 @@ def test_df_RKS(self):
 
 if __name__ == "__main__":
     print("Full tests for to_gpu module")
-    unittest.main()
+    unittest.main()
+
diff --git a/gpu4pyscf/qmmm/pbc/__init__.py b/gpu4pyscf/qmmm/pbc/__init__.py
@@ -0,0 +1,16 @@
+# Copyright 2024 The GPU4PySCF Authors. All Rights Reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from gpu4pyscf.qmmm.pbc import mm_mole, itrf, tools