accelerated DFUHF hessian && smoke test

.github/workflows/nightly_build.yml

@@ -0,0 +1,36 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: nightly build on self-hosted machine
+
+on:
+  schedule:
+    - cron: "0 3 * * *"
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: self-hosted
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Install dependencies
+      run: |
+        python3 -m pip install --upgrade pip
+        pip3 install flake8 pytest coverage gpu4pyscf-libxc-cuda11x pytest-cov dftd3 dftd4
+        pip3 install pyscf --upgrade
+        git config --global core.compression 9
+    - name: Build GPU4PySCF
+      run: |
+        export CUDA_HOME=/usr/local/cuda
+        export PATH=${CUDA_HOME}/bin:${PATH}
+        export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+        sh build.sh
+    - name: Smoke Test
+      run: |
+        echo $GITHUB_WORKSPACE
+        export PYTHONPATH="${PYTHONPATH}:$(pwd)"
+        pytest -m smoke --durations=0

.github/workflows/unittest.yml

@@ -35,4 +35,4 @@ jobs:
       run: |
         echo $GITHUB_WORKSPACE
         export PYTHONPATH="${PYTHONPATH}:$(pwd)"
-        pytest --cov=$GITHUB_WORKSPACE
+        pytest -m "not smoke" --cov=$GITHUB_WORKSPACE

gpu4pyscf/df/df_jk.py

@@ -285,7 +285,6 @@ def get_jk(dfobj, dms_tag, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-
 
     # SCF K matrix with occ
     if getattr(dms_tag, 'mo_coeff', None) is not None:
-        #occ_coeff = cupy.asarray(dms_tag.occ_coeff[ao_idx, :], order='C')
         mo_occ = dms_tag.mo_occ
         mo_coeff = dms_tag.mo_coeff
         nmo = mo_occ.shape[-1]
@@ -307,7 +306,6 @@ def get_jk(dfobj, dms_tag, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-
             for i in range(nset):
                 if with_k:
                     rhok = contract('Lij,jk->Lki', cderi, occ_coeff[i])
-                    #vk[0] += 2.0 * contract('Lki,Lkj->ij', rhok, rhok)
                     cublas.syrk('T', rhok.reshape([-1,nao]), out=vk[i], alpha=1.0, beta=1.0, lower=True)
         if with_j:
             vj[:,rows,cols] = vj_packed
@@ -319,30 +317,40 @@ def get_jk(dfobj, dms_tag, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-
             transpose_sum(vk)
     # CP-HF K matrix
     elif hasattr(dms_tag, 'mo1'):
+        occ_coeffs = dms_tag.occ_coeff
+        mo1s = dms_tag.mo1
+        mo_occ = dms_tag.mo_occ
+        if not isinstance(occ_coeffs, list):
+            occ_coeffs = [occ_coeffs * 2.0] # For restricted
+        if not isinstance(mo1s, list):
+            mo1s = [mo1s]
+
+        occ_coeffs = [occ_coeff[ao_idx] for occ_coeff in occ_coeffs]
+        mo1s = [mo1[:,ao_idx] for mo1 in mo1s]
+
         if with_j:
             vj_sparse = cupy.zeros_like(dm_sparse)
-        mo1 = dms_tag.mo1[:,ao_idx,:]
-        nocc = mo1.shape[2]
-        # 2.0 due to rhok and rhok1, put it here for symmetry
-        occ_coeff = dms_tag.occ_coeff[ao_idx,:] * 2.0
+
+        nocc = max([mo1.shape[2] for mo1 in mo1s])
+
         blksize = dfobj.get_blksize(extra=2*nao*nocc)
         for cderi, cderi_sparse in dfobj.loop(blksize=blksize, unpack=with_k):
             if with_j:
-                #vj += get_j(cderi_sparse)
                 rhoj = 2.0*dm_sparse.dot(cderi_sparse)
                 vj_sparse += cupy.dot(rhoj, cderi_sparse.T)
             if with_k:
-                rhok = contract('Lij,jk->Lki', cderi, occ_coeff)
-                for i in range(mo1.shape[0]):
-                    rhok1 = contract('Lij,jk->Lki', cderi, mo1[i])
-                    #vk[i] += contract('Lki,Lkj->ij', rhok, rhok1)
-                    contract('Lki,Lkj->ij', rhok, rhok1, alpha=1.0, beta=1.0, out=vk[i])
+                iset = 0
+                for occ_coeff, mo1 in zip(occ_coeffs, mo1s):
+                    rhok = contract('Lij,jk->Lki', cderi, occ_coeff)
+                    for i in range(mo1.shape[0]):
+                        rhok1 = contract('Lij,jk->Lki', cderi, mo1[i])
+                        contract('Lki,Lkj->ij', rhok, rhok1, alpha=1.0, beta=1.0, out=vk[iset])
+                        iset += 1
         occ_coeff = rhok1 = rhok = mo1 = None
         if with_j:
             vj[:,rows,cols] = vj_sparse
             vj[:,cols,rows] = vj_sparse
         if with_k:
-            #vk = vk + vk.transpose(0,2,1)
             transpose_sum(vk)
     # general K matrix with density matrix
     else:

gpu4pyscf/dft/numint.py

@@ -1093,10 +1093,12 @@ def nr_uks_fxc(ni, mol, grids, xc_code, dm0=None, dms=None, relativity=0, hermi=
     # AO basis -> gdftopt AO basis
     with_mocc = hasattr(dms, 'mo1')
     if with_mocc:
-        mo1a = contract('nio,pi->npo', dma.mo1, coeff) * 2.0**0.5
-        mo1b = contract('nio,pi->npo', dmb.mo1, coeff) * 2.0**0.5
-        occ_coeff_a = contract('io,pi->po', dma.occ_coeff, coeff) * 2.0**0.5
-        occ_coeff_b = contract('io,pi->po', dmb.occ_coeff, coeff) * 2.0**0.5
+        mo1a, mo1b = dms.mo1
+        occ_coeffa, occ_coeffb = dms.occ_coeff
+        mo1a = contract('nio,pi->npo', mo1a, coeff)
+        mo1b = contract('nio,pi->npo', mo1b, coeff)
+        occ_coeff_a = contract('io,pi->po', occ_coeffa, coeff)
+        occ_coeff_b = contract('io,pi->po', occ_coeffb, coeff)
 
     dma = cupy.asarray(dma).reshape(-1,nao0,nao0)
     dmb = cupy.asarray(dmb).reshape(-1,nao0,nao0)

gpu4pyscf/hessian/uhf.py

@@ -417,21 +417,26 @@ def fx(mo1):
         nset = len(mo1)
 
         x = mo1[:,:nmoa*nocca].reshape(nset,nmoa,nocca)
-        mo1_mo = contract('npo,ip->nio', x, mo_coeff[0])
-        dma = contract('nio,jo->nij', mo1_mo, mocca)
+        mo1_moa = contract('npo,ip->nio', x, mo_coeff[0])
+        dma = contract('nio,jo->nij', mo1_moa, mocca)
 
         x = mo1[:,nmoa*nocca:].reshape(nset,nmob,noccb)
-        mo1_mo = contract('npo,ip->nio', x, mo_coeff[1])
-        dmb = contract('nio,jo->nij', mo1_mo, moccb)
+        mo1_mob = contract('npo,ip->nio', x, mo_coeff[1])
+        dmb = contract('nio,jo->nij', mo1_mob, moccb)
 
         dm1 = cupy.empty([2,nset,nao,nao])
         dm1[0] = dma + dma.transpose(0,2,1)
         dm1[1] = dmb + dmb.transpose(0,2,1)
 
+        #v1_old = vresp(dm1)
         # TODO: improve the efficiency with occ_coeff
-        #dm1 = tag_array(dm1, mo1=mo1_mo, occ_coeff=mocc, mo_occ=mo_occ)
+        dm1 = tag_array(dm1, mo1=[mo1_moa,mo1_mob], occ_coeff=[mocca,moccb], mo_occ=mo_occ)
+        #print(dm1.shape)
         v1 = vresp(dm1)
-
+        #print(cupy.linalg.norm(v1 - v1_old))
+        #print(cupy.linalg.norm(v1))
+        #print(v1.shape)
+        #exit()
         v1vo = cupy.empty_like(mo1)
         tmp = contract('nij,jo->nio', v1[0], mocca)
         v1vo[:,:nmoa*nocca] = contract('nio,ip->npo', tmp, mo_coeff[0]).reshape(nset,-1)

gpu4pyscf/solvent/grad/pcm.py

@@ -22,9 +22,8 @@
 import cupy
 from cupyx import scipy
 from pyscf import lib
-from pyscf import gto, df
+from pyscf import gto
 from pyscf.grad import rhf as rhf_grad
-from pyscf.solvent import ddcosmo_grad
 
 from gpu4pyscf.solvent.pcm import PI, switch_h
 from gpu4pyscf.df import int3c2e

gpu4pyscf/solvent/grad/smd.py

@@ -22,17 +22,14 @@
 import cupy
 #from cupyx import scipy, jit
 from pyscf import lib
-from pyscf import gto, df
 from pyscf.grad import rhf as rhf_grad
 from pyscf.data import radii
-from pyscf.solvent import ddcosmo_grad
 
 from gpu4pyscf.solvent import pcm, smd
 from gpu4pyscf.solvent.grad import pcm as pcm_grad
 from gpu4pyscf.solvent.smd import (
     sigma_water, sigma_n, sigma_alpha, sigma_beta, r_zz, swtich_function,
     hartree2kcal)
-from gpu4pyscf.df import int3c2e
 from gpu4pyscf.lib.cupy_helper import contract
 from gpu4pyscf.lib import logger
 

gpu4pyscf/solvent/hessian/pcm.py

@@ -20,9 +20,7 @@
 
 import numpy
 import cupy
-from cupyx import scipy
-from pyscf import lib, gto, df
-from pyscf.grad import rhf as rhf_grad
+from pyscf import lib, gto
 from gpu4pyscf import scf
 from gpu4pyscf.solvent.pcm import PI, switch_h
 from gpu4pyscf.solvent.grad.pcm import grad_switch_h, get_dF_dA, get_dD_dS, grad_qv, grad_solver, grad_nuc

gpu4pyscf/solvent/hessian/smd.py

@@ -20,20 +20,16 @@
 
 import numpy as np
 import cupy
-from pyscf import gto, df, scf, lib
-from pyscf.grad import rhf as rhf_grad
-from pyscf.data import radii
-from pyscf.solvent import ddcosmo_grad
+from pyscf import lib
 
-from gpu4pyscf.solvent import pcm, smd
-from gpu4pyscf.solvent.grad import pcm as pcm_grad
+from gpu4pyscf.solvent import smd
 from gpu4pyscf.solvent.grad import smd as smd_grad
 from gpu4pyscf.solvent.hessian import pcm as pcm_hess
 
 from gpu4pyscf.solvent.smd import (
     sigma_water, sigma_n, sigma_alpha, sigma_beta, r_zz, swtich_function,
     hartree2kcal)
-from gpu4pyscf.df import int3c2e
+from gpu4pyscf import scf
 from gpu4pyscf.lib.cupy_helper import contract
 from gpu4pyscf.lib import logger
 

gpu4pyscf/tests/test_dft.py

@@ -0,0 +1,149 @@
+# 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 numpy as np
+import pyscf
+import pytest
+import cupy
+
+atom = '''
+C                 -0.07551087    1.68127663   -0.10745193
+O                  1.33621755    1.87147409   -0.39326987
+C                  1.67074668    2.95729545    0.49387976
+C                  0.41740763    3.77281969    0.78495878
+C                 -0.60481480    3.07572636    0.28906224
+H                 -0.19316298    1.01922455    0.72486113
+O                  0.35092043    5.03413298    1.45545728
+H                  0.42961487    5.74279041    0.81264173
+O                 -1.95331750    3.53349874    0.15912025
+H                 -2.55333895    2.78846397    0.23972698
+O                  2.81976302    3.20110148    0.94542226
+C                 -0.81772499    1.09230218   -1.32146482
+H                 -0.70955636    1.74951833   -2.15888136
+C                 -2.31163857    0.93420736   -0.98260166
+H                 -2.72575463    1.89080093   -0.74107186
+H                 -2.41980721    0.27699120   -0.14518512
+O                 -0.26428017   -0.18613595   -1.64425697
+H                 -0.72695910   -0.55328886   -2.40104423
+O                 -3.00083741    0.38730252   -2.10989934
+H                 -3.93210821    0.28874990   -1.89865997
+'''
+
+mol = pyscf.M(atom=atom, basis='def2-tzvpp', max_memory=32000, cart=0)
+mol.output = '/dev/null'
+mol.build()
+mol.verbose = 1
+
+@pytest.mark.smoke
+def test_DFRKS():
+    print('-------- DFRKS with D3(BJ) -------')
+    from gpu4pyscf.dft import rks
+    mf = rks.RKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-10
+    mf.conv_tol_cpscf = 1e-8
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0326965348272) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.17498362161082373) < 1e-6
+
+    h = mf.Hessian().kernel()
+    assert np.abs(cupy.linalg.norm(h) - 3.7684319231335377) < 1e-4
+
+@pytest.mark.smoke
+def test_DFUKS():
+    print('------- DFUKS with D3(BJ) -------')
+    from gpu4pyscf.dft import uks
+    mf = uks.UKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-10
+    mf.conv_tol_cpscf = 1e-8
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0326965349493) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.17498264516108836) < 1e-6
+
+    h = mf.Hessian().kernel()
+    assert np.abs(cupy.linalg.norm(h) - 3.768429871470736) < 1e-4
+
+@pytest.mark.smoke
+def test_RKS():
+    print('-------- RKS with D3(BJ) -------')
+    from gpu4pyscf.dft import rks
+    mf = rks.RKS(mol, xc='b3lyp')
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-12
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0325611822375) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.1750368231223345) < 1e-6
+
+@pytest.mark.smoke
+def test_UKS():
+    print('--------- UKS with D3(BJ) -------')
+    from gpu4pyscf.dft import uks
+    mf = uks.UKS(mol, xc='b3lyp')
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-12
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0325611823603) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.17503584692057772) < 1e-6
+
+@pytest.mark.smoke
+def test_DFRKS_with_SMD():
+    print('----- DFRKS with SMD -----')
+    from gpu4pyscf.dft import rks
+    mf = rks.RKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+    mf = mf.SMD()
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-10
+    mf.conv_tol_cpscf = 1e-8
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0579697676872) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.16832696143145143) < 1e-6
+
+    h = mf.Hessian().kernel()
+    assert np.abs(cupy.linalg.norm(h) - 3.750511574920911) < 1e-4
+
+@pytest.mark.smoke
+def test_DFUKS_with_SMD():
+    print('------- DFUKS with SMD ---------')
+    from gpu4pyscf.dft import uks
+    mf = uks.UKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+    mf = mf.SMD()
+    mf.grids.atom_grid = (99,590)
+    mf.conv_tol = 1e-10
+    mf.conv_tol_cpscf = 1e-8
+    mf.disp = 'd3bj'
+    e_dft = mf.kernel()
+    assert np.abs(e_dft - -685.0579697677729) < 1e-8
+
+    g = mf.nuc_grad_method().kernel()
+    assert np.abs(cupy.linalg.norm(g) - 0.16832715562574982) < 1e-6
+
+    h = mf.Hessian().kernel()
+    assert np.abs(cupy.linalg.norm(h) - 3.7505074541182823) < 1e-4