gpu4pyscf 1.3.2__tar.gz → 1.4.0__tar.gz

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/CHANGELOG

@@ -1,8 +1,28 @@
+v1.4.0 (2025-03-27)
+-------------------
+* New Features
+  - RKS and UKS TDDFT Gradients for density fitting and direct-SCF methods.
+  - ECP integrals and its first and second derivatives accelerated on GPU.
+  - Multigrid algorithm for Coulomb matrix and LDA, GGA, MGGA functionals computation.
+  - PBC Gaussian density fitting integrals.
+  - ASE interface for molecular systems.
+* Improvements
+  - Reduce memory footprint in SCF driver.
+  - Reduce memory requirements for PCM energy and gradients.
+  - Reduce memory requirements for DFT gradients.
+  - Utilize the sparsity in cart2sph coefficients in the cart2sph transformation in scf.jk kernel
+  - Molecular 3c2e integrals generated using the block-divergent alogrithm.
+  - Support I orbitals in DFT.
+* Fixes
+  - LRU cached cart2sph under the multiple GPU environment.
+  - A maxDynamicSharedMemorySize setting bug in gradient and hessian calculation under the multiple GPU environment.
+  - Remove the limits of 6000 GTO shells in DFT numerical integration module.
+
 v1.3.2 (2025-03-10)
 -------------------
 * Improvements
   - Dump xc info and grids into to log file
-  - Optimize 4-center integral evalulation CUDA kernels using warp divergent algorithm
+  - Optimize 4-center integral evaluation CUDA kernels using warp divergent algorithm
   - Support up to I orbitals in DFT
   - Fix out-of-bound issue in DFT hessian for heavy atoms (>=19)
 * Deprecation

{gpu4pyscf-1.3.2/gpu4pyscf.egg-info → gpu4pyscf-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 1.3.2
+Version: 1.4.0
 Summary: GPU extensions for PySCF
 Author: PySCF developers
 License: Apache-2.0

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/README.md

@@ -7,7 +7,7 @@ Installation
 --------
 
 > [!NOTE]
-> The compiled binary packages support compute capability 6.0 and later (Pascal and later, such as Tesla P100, RTX 10 series and later).
+> The compiled binary packages support compute capability 7.0 and later (Volta and later, such as Tesla V100, RTX 20 series and later).
 
 Run ```nvcc --version``` in your terminal to check the installed CUDA toolkit version. Then, choose the proper package based on your CUDA toolkit version.
 
@@ -27,7 +27,7 @@ cmake --build build/temp.gpu4pyscf -j 4
 CURRENT_PATH=`pwd`
 export PYTHONPATH="${PYTHONPATH}:${CURRENT_PATH}"
 ```
-Then install cutensor and cupy for acceleration (please switch the versions according to your nvcc version!)
+Then install cutensor and cupy for acceleration (please switch the versions according to your runtime CUDA environment!)
 ```sh
 pip3 install cutensor-cu12 cupy-cuda12x
 ```
@@ -45,17 +45,21 @@ Features
 - LDA, GGA, mGGA, hybrid, and range-separated functionals via [libXC](https://gitlab.com/libxc/libxc/-/tree/master/);
 - Spin-conserved and spin-flip TDA and TDDFT for excitated states
 - Geometry optimization and transition state search via [geomeTRIC](https://geometric.readthedocs.io/en/latest/);
+- Atomic Simulation Environment ([ASE](https://gitlab.com/ase/ase)) interface;
 - Dispersion corrections via [DFTD3](https://github.com/dftd3/simple-dftd3) and [DFTD4](https://github.com/dftd4/dftd4);
 - Nonlocal functional correction (vv10) for SCF and gradient;
-- ECP is supported and calculated on CPU;
-- PCM models, SMD model, their analytical gradients, and semi-analytical Hessian matrix;
+- ECP is supported and calculated on GPU;
+- PCM models, their analytical gradients, and analytical Hessian matrix;
+- SMD solvent model;
 - Unrestricted Hartree-Fock and unrestricted DFT, gradient, and Hessian;
-- MP2/DF-MP2 and CCSD (experimental);
-- Polarizability, IR, and NMR shielding (experimental);
-- QM/MM with PBC;
 - CHELPG, ESP, and RESP atomic charge;
-- Multi-GPU for both direct SCF and density fitting (experimental)
-- SCF and DFT with periodic boundary condition (experimental)
+
+The following features are still in the experimental stage
+- MP2/DF-MP2 and CCSD;
+- Polarizability, IR, and NMR shielding;
+- QM/MM with PBC;
+- Multi-GPU for both direct SCF and density fitting
+- SCF and DFT with periodic boundary condition
 
 Limitations
 --------
@@ -65,6 +69,7 @@ Limitations
 - Density fitting scheme up to ~168 atoms with def2-tzvpd basis, bounded by CPU memory;
 - meta-GGA without density laplacian;
 - Double hybrid functionals are not supported;
+- Hessian of TDDFT is not supported;
 
 Examples
 --------
@@ -113,7 +118,7 @@ Find more examples in [gpu4pyscf/examples](https://github.com/pyscf/gpu4pyscf/tr
 
 Benchmarks
 --------
-Speedup with GPU4PySCF v0.6.0 on A100-80G over Q-Chem 6.1 on 32-cores CPU (Desity fitting, SCF, def2-tzvpp, def2-universal-jkfit, B3LYP, (99,590))
+Speedup with GPU4PySCF v0.6.0 on A100-80G over Q-Chem 6.1 on 32-cores CPU (density fitting, SCF, def2-tzvpp, def2-universal-jkfit, B3LYP, (99,590))
 
 | mol               |   natm |    LDA |    PBE |   B3LYP |    M06 |   wB97m-v |
 |:------------------|-------:|-------:|-------:|--------:|-------:|----------:|

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/gpu4pyscf/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '1.3.2'
+__version__ = '1.4.0'
 
-from . import lib, grad, hessian, solvent, scf, dft
+from . import lib, grad, hessian, solvent, scf, dft, tdscf

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/gpu4pyscf/df/grad/__init__.py

@@ -12,9 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import rhf, rks, uhf, uks
+from . import rhf, rks, uhf, uks, tdrhf, tdrks, tduhf, tduks
 
 RHF = rhf.Gradients
 RKS = rks.Gradients
 UHF = uhf.Gradients
 UKS = uks.Gradients
+TDRHF = tdrhf.Gradients
+TDRKS = tdrks.Gradients
+TDUHF = tduhf.Gradients
+TDUKS = tduks.Gradients

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/gpu4pyscf/df/grad/jk.py

@@ -86,24 +86,23 @@ def get_rhojk(with_df, dm, orbo, with_j=True, with_k=True):
 def _jk_ip_task(intopt, rhoj_cart, dm_cart, rhok_cart, orbo_cart, task_list,
                 with_j=True, with_k=True, device_id=0, omega=None):
     mol = intopt.mol
+    natm = mol.natm
     with cupy.cuda.Device(device_id), _streams[device_id]:
         log = logger.new_logger(mol, mol.verbose)
         t0 = (logger.process_clock(), logger.perf_counter())
 
         orbo_cart = cupy.asarray(orbo_cart)
         cart_aux_loc = intopt.cart_aux_loc
-        nao_cart = dm_cart.shape[0]
-        naux_cart = intopt._sorted_auxmol.nao
-        vj = vk = vjaux = vkaux = None
+        ej = ek = ejaux = ekaux = None
         if with_j:
             rhoj_cart = cupy.asarray(rhoj_cart)
             dm_cart = cupy.asarray(dm_cart)
-            vj = cupy.zeros((3,nao_cart), order='C')
-            vjaux = cupy.zeros((3,naux_cart))
+            ej = cupy.zeros((natm,3), order='C')
+            ejaux = cupy.zeros((natm,3))
         if with_k:
             rhok_cart = cupy.asarray(rhok_cart)
-            vk = cupy.zeros((3,nao_cart), order='C')
-            vkaux = cupy.zeros((3,naux_cart))
+            ek = cupy.zeros((natm,3), order='C')
+            ekaux = cupy.zeros((natm,3))
         
         for cp_kl_id in task_list:
             k0, k1 = cart_aux_loc[cp_kl_id], cart_aux_loc[cp_kl_id+1]
@@ -135,16 +134,16 @@ def _jk_ip_task(intopt, rhoj_cart, dm_cart, rhok_cart, orbo_cart, task_list,
             vjaux_outcore = contract('xp,p->xp', rhoj_tmp0, rhoj_cart[k0:k1])
             vkaux_outcore = contract('xpji,pji->xp', int3c_ip, rhok_tmp)
             '''
-            vj_tmp, vk_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip1', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
-            if with_j: vj += vj_tmp
-            if with_k: vk += vk_tmp
-            vj_tmp, vk_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip2', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
-            if with_j: vjaux[:, k0:k1] = vj_tmp
-            if with_k: vkaux[:, k0:k1] = vk_tmp
-
-            rhoj_tmp = rhok_tmp = vj_tmp = vk_tmp = None
+            ej_tmp, ek_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip1', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
+            if with_j: ej += ej_tmp
+            if with_k: ek += ek_tmp
+            ej_tmp, ek_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip2', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
+            if with_j: ejaux += ej_tmp
+            if with_k: ekaux += ek_tmp
+
+            rhoj_tmp = rhok_tmp = ej_tmp = ek_tmp = None
             t0 = log.timer_debug1(f'calculate {cp_kl_id:3d} / {len(intopt.aux_log_qs):3d}, {k1-k0:3d} slices', *t0)
-    return vj, vk, vjaux, vkaux
+    return ej, ek, ejaux, ekaux
 
 def get_grad_vjk(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbo_cart, 
                  with_j=True, with_k=True, omega=None):
@@ -191,3 +190,159 @@ def get_grad_vjk(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbo_cart,
         rhok = reduce_to_device(rhok_total)
         vkaux = reduce_to_device(vkaux_total)
     return rhoj, rhok, vjaux, vkaux
+
+
+def _jk_task_td(with_df, dm, orbol, orbor, with_j=True, with_k=True, device_id=0):
+    '''  
+    (L|ij) -> rhoj: (L), rhok: (L|lr), for dm0 from scf, rhok is (L|oo) 
+    '''
+    rhoj = rhok = None
+    with cupy.cuda.Device(device_id), _streams[device_id]:
+        log = logger.new_logger(with_df.mol, with_df.verbose)
+        assert isinstance(with_df.verbose, int)
+        t0 = log.init_timer()
+        dm = cupy.asarray(dm)
+        naux_slice = with_df._cderi[device_id].shape[0]
+        nl = orbol.shape[-1]
+        nr = orbor.shape[-1]
+
+        rows = with_df.intopt.cderi_row
+        cols = with_df.intopt.cderi_col
+        dm_sparse = dm[rows, cols]
+        dm_sparse[with_df.intopt.cderi_diag] *= .5
+
+        blksize = with_df.get_blksize()
+        if with_j:
+            rhoj = cupy.empty([naux_slice])
+        if with_k:
+            rhok = cupy.empty([naux_slice, nl, nr], order='C')
+        p0 = p1 = 0
+
+        for cderi, cderi_sparse in with_df.loop(blksize=blksize):
+            p1 = p0 + cderi.shape[0]
+            if with_j:
+                rhoj[p0:p1] = 2.0*dm_sparse.dot(cderi_sparse)
+            if with_k:
+                rhok[p0:p1] = 0
+                tmp = contract('Lij,jk->Lki', cderi, orbol)
+                contract('Lki,il->Lkl', tmp, orbor, out=rhok[p0:p1])
+            p0 = p1
+            cupy.cuda.get_current_stream().synchronize()
+        t0 = log.timer_debug1(f'rhoj and rhok on Device {device_id}', *t0)
+    return rhoj, rhok
+
+
+def get_rhojk_td(with_df, dm, orbol, orbor, with_j=True, with_k=True):
+    ''' Calculate rhoj and rhok on Multi-GPU system
+    '''
+    futures = []
+    cupy.cuda.get_current_stream().synchronize()
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
+            future = executor.submit(
+                _jk_task_td, with_df, dm, orbol, orbor,
+                with_j=with_j, with_k=with_k, device_id=device_id)
+            futures.append(future)
+
+    rhoj_total = []
+    rhok_total = []
+    for future in futures:
+        rhoj, rhok = future.result()
+        rhoj_total.append(rhoj)
+        rhok_total.append(rhok)
+
+    rhoj = rhok = None
+    if with_j:
+        rhoj = concatenate(rhoj_total)
+    if with_k:
+        rhok = concatenate(rhok_total)
+
+    return rhoj, rhok
+
+
+def _jk_ip_task_td(intopt, rhoj_cart, dm_cart, rhok_cart, orbol_cart, orbor_cart, task_list,
+                with_j=True, with_k=True, device_id=0, omega=None):
+    mol = intopt.mol
+    natm = mol.natm
+    with cupy.cuda.Device(device_id), _streams[device_id]:
+        log = logger.new_logger(mol, mol.verbose)
+        t0 = (logger.process_clock(), logger.perf_counter())
+
+        cart_aux_loc = intopt.cart_aux_loc
+        ej = ek = ejaux = ekaux = None
+        if with_j:
+            rhoj_cart = cupy.asarray(rhoj_cart)
+            dm_cart = cupy.asarray(dm_cart)
+            ej = cupy.zeros((natm,3), order='C')
+            ejaux = cupy.zeros((natm,3))
+        if with_k:
+            rhok_cart = cupy.asarray(rhok_cart)
+            ek = cupy.zeros((natm,3), order='C')
+            ekaux = cupy.zeros((natm,3))
+        
+        for cp_kl_id in task_list:
+            k0, k1 = cart_aux_loc[cp_kl_id], cart_aux_loc[cp_kl_id+1]
+            rhoj_tmp = rhok_tmp = None
+            if with_j:
+                rhoj_tmp = rhoj_cart[k0:k1]
+            if with_k:
+                rhok_tmp = contract('por,ir->pio', rhok_cart[k0:k1], orbol_cart)
+                rhok_tmp = contract('pio,jo->pji', rhok_tmp, orbor_cart)
+
+            ej_tmp, ek_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip1', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
+            if with_j: ej += ej_tmp
+            if with_k: ek += ek_tmp
+            ej_tmp, ek_tmp = get_int3c2e_ip_jk(intopt, cp_kl_id, 'ip2', rhoj_tmp, rhok_tmp, dm_cart, omega=omega)
+            if with_j: ejaux += ej_tmp
+            if with_k: ekaux += ek_tmp
+
+            rhoj_tmp = rhok_tmp = ej_tmp = ek_tmp = None
+            t0 = log.timer_debug1(f'calculate {cp_kl_id:3d} / {len(intopt.aux_log_qs):3d}, {k1-k0:3d} slices', *t0)
+    return ej, ek, ejaux, ekaux
+
+
+def get_grad_vjk_td(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbol_cart, orbor_cart, 
+                 with_j=True, with_k=True, omega=None):
+    '''
+    Calculate vj    = (i'j|L)(L|kl)(ij)(kl), vk    = (i'j|L)(L|kl)(ik)(jl)
+              vjaux = (ij|L')(L|kl)(ij)(kl), vkaux = (ij|L')(L|kl)(ik)(jl)
+    '''
+    nao_cart = dm_cart.shape[0]
+    block_size = with_df.get_blksize(nao=nao_cart)
+
+    intopt = VHFOpt(mol, auxmol, 'int2e')
+    intopt.build(1e-14, diag_block_with_triu=True, aosym=False,
+                 group_size_aux=block_size, verbose=0)#, group_size=block_size)
+
+    aux_ao_loc = np.array(intopt.aux_ao_loc)
+    loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
+    task_list = _split_tasks(loads, num_devices)
+
+    futures = []
+    cupy.cuda.get_current_stream().synchronize()
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
+            future = executor.submit(
+                _jk_ip_task_td, intopt, rhoj_cart, dm_cart, rhok_cart, orbol_cart, orbor_cart, task_list[device_id],
+                with_j=with_j, with_k=with_k, device_id=device_id, omega=omega)
+            futures.append(future)
+
+    rhoj_total = []
+    rhok_total = []
+    ejaux_total = []
+    ekaux_total = []
+    for future in futures:
+        rhoj, rhok, ejaux, ekaux = future.result()
+        rhoj_total.append(rhoj)
+        rhok_total.append(rhok)
+        ejaux_total.append(ejaux)
+        ekaux_total.append(ekaux)
+
+    rhoj = rhok = ejaux = ekaux = None
+    if with_j:
+        rhoj = reduce_to_device(rhoj_total)
+        ejaux = reduce_to_device(ejaux_total)
+    if with_k:
+        rhok = reduce_to_device(rhok_total)
+        ekaux = reduce_to_device(ekaux_total)
+    return rhoj, rhok, ejaux, ekaux

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/gpu4pyscf/df/grad/rhf.py

@@ -59,7 +59,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
+    # extended to any 1-particle density matrix. The get_jk in tddft supports this function.
 
     if(dm0 is None): dm0 = mf_grad.base.make_rdm1()
     if omega is None:
@@ -106,6 +106,8 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
     aux_cart2sph = intopt.aux_cart2sph
     low = with_df.cd_low
     low_t = low.T.copy()
+
+    ejaux = ekaux = None
     if with_j:
         if low.tag == 'eig':
             rhoj = cupy.dot(low_t.T, rhoj)
@@ -120,7 +122,7 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
         rhoj = intopt.unsort_orbitals(rhoj, aux_axis=[0])
         tmp = contract('xpq,q->xp', int2c_e1, rhoj)
         vjaux = -contract('xp,p->xp', tmp, rhoj)
-        vjaux_2c = cupy.array([-vjaux[:,p0:p1].sum(axis=1) for p0, p1 in auxslices[:,2:]])
+        ejaux = cupy.array([-vjaux[:,p0:p1].sum(axis=1) for p0, p1 in auxslices[:,2:]])
         rhoj = vjaux = tmp = None
     if with_k:
         nocc = orbo.shape[-1]
@@ -133,7 +135,7 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
         tmp = contract('pij,qij->pq', rhok, rhok)
         tmp = intopt.unsort_orbitals(tmp, aux_axis=[0,1])
         vkaux = -contract('xpq,pq->xp', int2c_e1, tmp)
-        vkaux_2c = cupy.array([-vkaux[:,p0:p1].sum(axis=1) for p0, p1 in auxslices[:,2:]])
+        ekaux = cupy.array([-vkaux[:,p0:p1].sum(axis=1) for p0, p1 in auxslices[:,2:]])
         vkaux = tmp = None
         if not auxmol.cart:
             rhok_cart = contract('pq,qkl->pkl', aux_cart2sph, rhok)
@@ -153,39 +155,16 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
         dm_cart = cart2sph @ dm @ cart2sph.T
         
     with_df._cderi = None # release GPU memory
-    vj, vk, vjaux, vkaux = get_grad_vjk(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbo_cart,
+    ej, ek, ejaux_3c, ekaux_3c = get_grad_vjk(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbo_cart,
                                         with_j=with_j, with_k=with_k, omega=omega)
-    # NOTE: vj and vk are still in cartesian
-    _sorted_mol = intopt._sorted_mol
-    natm = _sorted_mol.natm
-    nao_cart = _sorted_mol.nao
-    ao2atom = numpy.zeros([nao_cart, natm])
-    ao_loc = _sorted_mol.ao_loc
-    for ibas, iatm in enumerate(_sorted_mol._bas[:,gto.ATOM_OF]):
-        ao2atom[ao_loc[ibas]:ao_loc[ibas+1],iatm] = 1
-    ao2atom = cupy.asarray(ao2atom)
     if with_j:
-        vj = -ao2atom.T @ vj.T
+        ej = -ej
+        ejaux -= ejaux_3c
     if with_k:
-        vk = -ao2atom.T @ vk.T
+        ek = -ek
+        ekaux -= ekaux_3c
     t0 = log.timer_debug1('(di,j|P) and (i,j|dP)', *t0)
-
-    _sorted_auxmol = intopt._sorted_auxmol
-    natm = _sorted_auxmol.natm
-    naux_cart = _sorted_auxmol.nao
-    aux2atom = numpy.zeros([naux_cart, natm])
-    ao_loc = _sorted_auxmol.ao_loc
-    for ibas, iatm in enumerate(_sorted_auxmol._bas[:,gto.ATOM_OF]):
-        aux2atom[ao_loc[ibas]:ao_loc[ibas+1],iatm] = 1
-    aux2atom = cupy.asarray(aux2atom)
-    if with_j:
-        vjaux_3c = aux2atom.T @ vjaux.T
-        vjaux = vjaux_2c - vjaux_3c
-
-    if with_k:
-        vkaux_3c = aux2atom.T @ vkaux.T
-        vkaux = vkaux_2c - vkaux_3c
-    return vj, vk, vjaux, vkaux
+    return ej, ek, ejaux, ekaux
 
 
 class Gradients(rhf_grad.Gradients):
@@ -228,4 +207,4 @@ class Gradients(rhf_grad.Gradients):
         else:
             return 0
 
-Grad = Gradients
+Grad = Gradients

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.0}/gpu4pyscf/df/grad/rks.py

@@ -55,14 +55,14 @@ def get_veff(ks_grad, mol=None, dm=None, verbose=None):
     mem_now = lib.current_memory()[0]
     max_memory = max(2000, ks_grad.max_memory*.9-mem_now)
     if ks_grad.grid_response:
-        exc, vxc = rks_grad.get_vxc_full_response(
+        exc, exc1 = rks_grad.get_exc_full_response(
                 ni, mol, grids, mf.xc, dm,
                 max_memory=max_memory, verbose=ks_grad.verbose)
         #logger.debug1(ks_grad, 'sum(grids response) %s', exc.sum(axis=0))
         if mf.do_nlc():
             raise NotImplementedError
     else:
-        exc, vxc = rks_grad.get_vxc(
+        exc, exc1 = rks_grad.get_exc(
                 ni, mol, grids, mf.xc, dm,
                 max_memory=max_memory, verbose=ks_grad.verbose)
         if mf.do_nlc():
@@ -70,47 +70,42 @@ def get_veff(ks_grad, mol=None, dm=None, verbose=None):
                 xc = mf.xc
             else:
                 xc = mf.nlc
-            enlc, vnlc = rks_grad.get_nlc_vxc(
+            enlc, exc1_nlc = rks_grad.get_nlc_exc(
                 ni, mol, nlcgrids, xc, dm,
                 max_memory=max_memory, verbose=ks_grad.verbose)
-            vxc += vnlc
+            exc1 += exc1_nlc
     t0 = logger.timer(ks_grad, 'vxc total', *t0)
 
-    # this can be moved into vxc calculations
-    occ_coeff = cupy.asarray(mf.mo_coeff[:, mf.mo_occ>0.5], order='C')
-    tmp = contract('nij,jk->nik', vxc, occ_coeff)
-    vxc = 2.0*contract('nik,ik->ni', tmp, occ_coeff)
-
     aoslices = mol.aoslice_by_atom()
-    vxc = [vxc[:,p0:p1].sum(axis=1) for p0, p1 in aoslices[:,2:]]
-    vxc = cupy.asarray(vxc)
+    exc1 = [exc1[:,p0:p1].sum(axis=1) for p0, p1 in aoslices[:,2:]]
+    exc1 = cupy.asarray(exc1)
     if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
-        vj, vjaux = ks_grad.get_j(mol, dm)
-        vxc += vj
+        ej, ejaux = ks_grad.get_j(mol, dm)
+        exc1 += ej
         if ks_grad.auxbasis_response:
-            e1_aux = vjaux
+            e1_aux = ejaux
     else:
-        vj, vk, vjaux, vkaux = ks_grad.get_jk(mol, dm)
+        ej, ek, ejaux, ekaux = ks_grad.get_jk(mol, dm)
 
         if ks_grad.auxbasis_response:
-            vk_aux = vkaux * hyb
-        vk *= hyb
+            ek_aux = ekaux * hyb
+        ek *= hyb
         if abs(omega) > 1e-10:  # For range separated Coulomb operator
-            vk_lr, vkaux_lr = ks_grad.get_k(mol, dm, omega=omega)
-            vk += vk_lr * (alpha - hyb)
+            ek_lr, ekaux_lr = ks_grad.get_k(mol, dm, omega=omega)
+            ek += ek_lr * (alpha - hyb)
             if ks_grad.auxbasis_response:
-                vk_aux += vkaux_lr * (alpha - hyb)
+                ek_aux += ekaux_lr * (alpha - hyb)
 
-        vxc += vj - vk * .5
+        exc1 += ej - ek * .5
         if ks_grad.auxbasis_response:
-            e1_aux = vjaux - vk_aux * .5
+            e1_aux = ejaux - ek_aux * .5
 
     if ks_grad.auxbasis_response:
         logger.debug1(ks_grad, 'sum(auxbasis response) %s', e1_aux.sum(axis=0))
     else:
         e1_aux = None
-    vxc = tag_array(vxc, aux=e1_aux, exc1_grid=exc)
-    return vxc
+    exc1 = tag_array(exc1, aux=e1_aux, exc1_grid=exc)
+    return exc1
 
 class Gradients(rks_grad.Gradients):
     from gpu4pyscf.lib.utils import to_gpu, device