fixed a bug in screen_index (#68)

Author: wxj6000

examples/00-h2o.py

@@ -43,7 +43,7 @@
 
 # Compute Energy
 e_dft = mf_GPU.kernel()
-print(f"total energy = {e_dft}")
+print(f"total energy = {e_dft}") # -76.26736519501688
 
 # Compute Gradient
 g = mf_GPU.nuc_grad_method()

examples/dft_driver.py

@@ -55,7 +55,6 @@
 e_tot = mf_df.kernel()
 scf_time = time.time() - start_time
 print(f'compute time for energy: {scf_time:.3f} s')
-exit()
 
 start_time = time.time()
 g = mf_df.nuc_grad_method()
@@ -70,6 +69,3 @@
 h_dft = h.kernel()
 hess_time = time.time() - start_time
 print(f'compute time for hessian: {hess_time:.3f} s')
-
-import numpy
-numpy.savez('gpu4pyscf_out.npz', e_tot=e_tot, f=f, h_dft=h_dft)

gpu4pyscf/__init__.py

@@ -1,5 +1,5 @@
 from . import lib, grad, hessian, solvent, scf, dft
-__version__ = '0.6.10'
+__version__ = '0.6.11'
 
 # monkey patch libxc reference due to a bug in nvcc
 from pyscf.dft import libxc

gpu4pyscf/dft/rks.py

@@ -226,6 +226,8 @@ def energy_elec(ks, dm=None, h1e=None, vhf=None):
 class RKS(scf.hf.RHF, rks.RKS):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
+    _keys = {'disp'}
+
     def __init__(self, mol, xc='LDA,VWN', disp=None):
         super().__init__(mol, xc)
         self._numint = numint.NumInt(xc=xc)

gpu4pyscf/grad/rhf.py

@@ -560,6 +560,8 @@ class Gradients(rhf.Gradients):
     get_veff = get_veff
     get_jk = _get_jk
 
+    _keys = {'auxbasis_response'}
+
     def get_j(self, mol=None, dm=None, hermi=0, omega=None):
         vj, _ = self.get_jk(mol, dm, with_k=False, omega=omega)
         return vj

gpu4pyscf/lib/gdft/nr_eval_gto.cu

@@ -76,7 +76,7 @@ void _screen_index(int *non0shl_idx, double cutoff, int l, int ish, int nprim, d
     double min_exp = 1e9;
     for (int ip = 0; ip < nprim; ++ip) {
         min_exp = MIN(min_exp, exps[ip]);
-        maxc = MAX(maxc, coeffs[ip]);
+        maxc = MAX(maxc, fabs(coeffs[ip]));
     }
     double gto_sup = -min_exp * rr + .5 * log(rr) * l + log(maxc);
     int is_large = gto_sup > log(cutoff);