fixed a bug in nlc

gpu4pyscf/df/grad/rhf.py

@@ -34,7 +34,7 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
     if mol is None: mol = mf_grad.mol
     #TODO: dm has to be the SCF density matrix in this version.  dm should be
     # extended to any 1-particle density matrix
-    
+
     if(dm0 is None): dm0 = mf_grad.base.make_rdm1()
     mf = mf_grad.base
     if omega is None:

gpu4pyscf/df/tests/test_df_grad.py

@@ -59,8 +59,10 @@ def _check_grad(grid_response=False, xc=xc0, disp=disp0, tol=1e-6):
     mf = rks.RKS(mol, xc=xc, disp=disp).density_fit(auxbasis=auxbasis0)
     mf.grids.level = grids_level
     mf.nlcgrids.level = nlcgrids_level
-    mf.conv_tol = 1e-12
+    mf.conv_tol = 1e-10
+    mf.verbose = 1
     e_tot = mf.kernel()
+
     g = mf.nuc_grad_method()
     g.auxbasis_response = True
     g.grid_response = grid_response

gpu4pyscf/dft/numint.py

@@ -382,14 +382,15 @@ def nr_rks(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
     dm_shape = dms.shape
     dms = [coeff @ dm @ coeff.T for dm in dms.reshape(-1,nao0,nao0)]
     nset = len(dms)
-    ao_loc = mol.ao_loc_nr()
 
     if mo_coeff is not None:
         mo_coeff = coeff @ mo_coeff
 
     nelec = cupy.zeros(nset)
     excsum = cupy.zeros(nset)
     vmat = cupy.zeros((nset, nao, nao))
+    '''
+    ao_loc = mol.ao_loc_nr()
     if USE_SPARSITY == 1:
         nbins = NBINS * 2 - int(NBINS * np.log(ni.cutoff) / np.log(grids.cutoff))
         pair2shls, pairs_locs = _make_pairs2shls_idx(ni.pair_mask, opt.l_bas_offsets, hermi)
@@ -398,7 +399,7 @@ def nr_rks(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
                                                                opt.l_bas_offsets)
         else:
             pair2shls_full, pairs_locs_full = pair2shls, pairs_locs
-
+    '''
     release_gpu_stack()
     if xctype == 'LDA':
         ao_deriv = 0
@@ -882,7 +883,7 @@ def nr_nlc_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
         wv[0] *= .5
         aow = _scale_ao(ao, wv)
         #vmat += ao[0].dot(aow.T)
-        add_sparse(vmat, ao.dot(aow.T), mask)
+        add_sparse(vmat, ao[0].dot(aow.T), mask)
 
     vmat = vmat + vmat.T
     vmat = contract('pi,pq->iq', coeff, vmat)

gpu4pyscf/dft/rks.py

@@ -79,7 +79,6 @@ def initialize_grids(ks, mol=None, dm=None):
             # Filter grids the first time setup grids
             ks.grids = prune_small_rho_grids_(ks, ks.mol, dm, ks.grids)
         t0 = logger.timer_debug1(ks, 'setting up grids', *t0)
-
         is_nlc = ks.nlc or ks._numint.libxc.is_nlc(ks.xc)
         if is_nlc and ks.nlcgrids.coords is None:
             if ks.nlcgrids.coords is None:
@@ -235,7 +234,14 @@ def __init__(self, mol, xc='LDA,VWN', disp=None):
         self._numint = numint.NumInt(xc=xc)
         self.disp = disp
         self.screen_tol = 1e-14
+
+        grids_level = self.grids.level
         self.grids = gen_grid.Grids(mol)
+        self.grids.level = grids_level
+
+        nlcgrids_level = self.nlcgrids.level
+        self.nlcgrids = gen_grid.Grids(mol)
+        self.nlcgrids.level = nlcgrids_level
 
     def get_dispersion(self):
         if self.disp is None:

gpu4pyscf/grad/rks.py

@@ -108,22 +108,22 @@ def _get_veff(ks_grad, mol=None, dm=None):
 
 def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
             max_memory=2000, verbose=None):
-    log = logger.new_logger(mol, verbose)
     xctype = ni._xc_type(xc_code)
     opt = getattr(ni, 'gdftopt', None)
     if opt is None:
         ni.build(mol, grids.coords)
         opt = ni.gdftopt
     mo_occ = cupy.asarray(dms.mo_occ)
     mo_coeff = cupy.asarray(dms.mo_coeff)
-
     coeff = cupy.asarray(opt.coeff)
     nao, nao0 = coeff.shape
     dms = cupy.asarray(dms)
     dms = [cupy.einsum('pi,ij,qj->pq', coeff, dm, coeff)
            for dm in dms.reshape(-1,nao0,nao0)]
     mo_coeff = coeff @ mo_coeff
+
     nset = len(dms)
+    assert nset == 1
 
     if xctype == 'LDA':
         ao_deriv = 1
@@ -139,9 +139,6 @@ def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
                 rho = numint.eval_rho2(opt.mol, ao_mask[0], mo_coeff_mask, mo_occ, None, xctype)
                 vxc = ni.eval_xc_eff(xc_code, rho, 1, xctype=xctype)[1]
                 wv = weight * vxc[0]
-                #mask = cupy.any(cupy.abs(ao) > AO_THRESHOLD, axis=[0,2])
-                #idx = cupy.argwhere(mask).astype(numpy.int32)[:,0]
-                #ao_mask = ao[:,idx,:]
                 aow = numint._scale_ao(ao_mask[0], wv)
                 vtmp = _d1_dot_(ao_mask[1:4], aow.T)
                 #idx = cupy.ix_(mask, mask)
@@ -160,9 +157,6 @@ def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
                 vxc = ni.eval_xc_eff(xc_code, rho, 1, xctype=xctype)[1]
                 wv = weight * vxc
                 wv[0] *= .5
-                #mask = cupy.any(cupy.abs(ao) > AO_THRESHOLD, axis=[0,2])
-                #idx = cupy.argwhere(mask).astype(numpy.int32)[:,0]
-                #ao_mask = ao[:,idx,:]
                 vtmp = _gga_grad_sum_(ao_mask, wv)
                 #idx = cupy.ix_(mask, mask)
                 #vmat[idm][0][idx] += vtmp[0]
@@ -184,9 +178,6 @@ def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
                 wv = weight * vxc
                 wv[0] *= .5
                 wv[4] *= .5  # for the factor 1/2 in tau
-                #mask = cupy.any(cupy.abs(ao) > AO_THRESHOLD, axis=[0,2])
-                #idx = cupy.argwhere(mask).astype(numpy.int32)[:,0]
-                #ao_mask = ao[:,idx,:]
                 vtmp = _gga_grad_sum_(ao_mask, wv)
                 vtmp += _tau_grad_dot_(ao_mask, wv[4])
                 #idx = cupy.ix_(mask, mask)