gpu4pyscf 1.3.2__tar.gz → 1.4.1__tar.gz

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.1}/CHANGELOG

@@ -1,8 +1,47 @@
+v1.4.1 (2025-05-20)
+-------------------
+* New Features
+  - Analytical hessian for VV10 functionals
+  - DFT polarizability with VV10 functionals
+  - TDDFT for VV10 functionals
+  - Non-adiabatic coupling constants for TDDFT states
+  - TDDFT gradients and geometry optimization solver for excited states
+  - LR-PCM for TDDFT and TDDFT gradients
+  - TDDFT-ris method
+* Improvements
+  - Optimization CUDA kernel and integral screening for MD J-engine. The MD
+    J-engine is utilized by default for large system HF and DFT computation.
+  - Optimization for PBC gaussian density fitting at gamma point.
+  - ECP gradients CUDA kernel
+  - Reduced atomicAdd overhead in Rys JK kernel
+* Fixes
+  - MINAO initial guess for ghost atoms
+
+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.1}/PKG-INFO

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

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

@@ -7,18 +7,34 @@ 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.
+To check your installed CUDA Toolkit version, run
+```sh
+nvcc --version
+```
+Then, install the appropriate package based on your CUDA version:
 
 | Platform      | Command                               | cutensor (**highly recommended**)|
 ----------------| --------------------------------------|----------------------------------|
 | **CUDA 11.x** |  ```pip3 install gpu4pyscf-cuda11x``` | ```pip3 install cutensor-cu11``` |
 | **CUDA 12.x** |  ```pip3 install gpu4pyscf-cuda12x``` | ```pip3 install cutensor-cu12``` |
 
+The versions of CuPy and cuTENSOR are strongly interdependent and should not be combined arbitrarily.
+The recommended combinations include:
+1. CuPy 13.3.0 + cuTENSOR 2.0.2
+2. CuPy 13.4.1 + cuTENSOR 2.2.0
+
+Using other versions or combinations may lead to failures in functionality. 
+We **recommend** creating a dedicated environment using:
+```sh
+pip3 install --no-cache-dir -r requirements.txt
+```
+This ensures compatibility and reproducibility, especially since this configuration is used in our nightly benchmarks.
+
 Compilation
 --------
-One can compile the package with
+To compile the package, run the following commands:
 ```sh
 git clone https://github.com/pyscf/gpu4pyscf.git
 cd gpu4pyscf
@@ -27,7 +43,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
 ```
@@ -36,26 +52,28 @@ There shouldn't be cupy or cutensor compilation during pip install process. If y
 <repo_path>/gpu4pyscf/lib/cutensor.py:<line_number>: UserWarning: using cupy as the tensor contraction engine.
 ```
 
-The package also provides multiple dockerfiles in ```dockerfiles```. One can use them as references to create the compilation envrionment.
-
 Features
 --------
 - Density fitting scheme and direct SCF scheme;
 - SCF, analytical gradient, and analytical Hessian calculations for Hartree-Fock and DFT;
 - 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
+- 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;
+- Analytical gradient and analytical Hessian for nonlocal functional correction (vv10);
+- GPU accelerated ECP;
+- 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 +83,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 +132,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.1}/gpu4pyscf/__config__.py

@@ -24,10 +24,10 @@ for device_id in range(num_devices):
 
 props = cupy.cuda.runtime.getDeviceProperties(0)
 GB = 1024*1024*1024
-min_ao_blksize = 128
-min_grid_blksize = 128*128
-ao_aligned = 32
-grid_aligned = 256
+min_ao_blksize = 256        # maxisum batch size of AOs
+min_grid_blksize = 128*128  # maximum batch size of grids for DFT
+ao_aligned = 32             # global AO alignment for slicing
+grid_aligned = 256          # 256 alignment for grids globally
 
 # Use smaller blksize for old gaming GPUs
 if props['totalGlobalMem'] < 16 * GB:

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.1}/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.1'
 
-from . import lib, grad, hessian, solvent, scf, dft
+from . import lib, grad, hessian, solvent, scf, dft, tdscf, nac

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.1}/gpu4pyscf/df/df.py

@@ -20,8 +20,9 @@ import numpy as np
 from cupyx.scipy.linalg import solve_triangular
 from pyscf import lib
 from pyscf.df import df, addons, incore
-from gpu4pyscf.lib.cupy_helper import (cholesky, tag_array, get_avail_mem, 
-                                       cart2sph, p2p_transfer, copy_array)
+from gpu4pyscf.lib.cupy_helper import (
+    cholesky, tag_array, get_avail_mem, cart2sph, p2p_transfer, copy_array,
+    asarray)
 from gpu4pyscf.df import int3c2e, df_jk
 from gpu4pyscf.df import int3c2e_bdiv
 from gpu4pyscf.lib import logger
@@ -125,7 +126,7 @@ class DF(lib.StreamObject):
                 j2c_cpu = auxmol.intor('int2c2e', hermi=1)
         else:
             j2c_cpu = auxmol.intor('int2c2e', hermi=1)
-        j2c = cupy.asarray(j2c_cpu, order='C')
+        j2c = asarray(j2c_cpu)
         t0 = log.timer_debug1('2c2e', *t0)
         intopt = int3c2e.VHFOpt(mol, auxmol, 'int2e')
         intopt.build(direct_scf_tol, diag_block_with_triu=False, aosym=True,
@@ -205,7 +206,7 @@ class DF(lib.StreamObject):
             if isinstance(cderi_sparse, np.ndarray):
                 # first block
                 if buf_prefetch is None:
-                    buf = cupy.asarray(cderi_sparse[p0:p1,:])
+                    buf = asarray(cderi_sparse[p0:p1,:])
                 buf_prefetch = cupy.empty([p2-p1,cderi_sparse.shape[1]])
             if isinstance(cderi_sparse, np.ndarray) and p1 < p2:
                 buf_prefetch.set(cderi_sparse[p1:p2,:])
@@ -322,7 +323,7 @@ def _cderi_task(intopt, cd_low, task_list, _cderi, aux_blksize,
         log = logger.new_logger(mol, mol.verbose)
         t1 = log.init_timer()
         cd_low_tag = cd_low.tag
-        cd_low = cupy.asarray(cd_low)
+        cd_low = asarray(cd_low)
 
         cart_ao_loc = intopt.cart_ao_loc
         aux_ao_loc = intopt.aux_ao_loc
@@ -411,7 +412,7 @@ def _cholesky_eri_bdiv(intopt, omega=None):
     if intopt.mol.cart:
         eri3c = intopt.orbital_pair_cart2sph(eri3c)
     auxmol = intopt.auxmol
-    j2c = cupy.asarray(auxmol.intor('int2c2e', hermi=1), order='C')
+    j2c = asarray(auxmol.intor('int2c2e', hermi=1), order='C')
     cd_low = cholesky(j2c)
     aux_coeff = cupy.array(intopt.aux_coeff, copy=True)
     cd_low = solve_triangular(cd_low, aux_coeff.T, lower=True, overwrite_b=True)

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.1}/gpu4pyscf/df/df_jk.py

@@ -34,35 +34,6 @@ def _pin_memory(array):
     ret[...] = array
     return ret
 
-def init_workflow(mf, dm0=None):
-    # build CDERI for omega = 0 and omega ! = 0
-    def build_df():
-        mf.with_df.build()
-        if hasattr(mf, '_numint'):
-            omega, _, _ = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mf.mol.spin)
-            if abs(omega) <= 1e-10: return
-            key = '%.6f' % omega
-            if key in mf.with_df._rsh_df:
-                rsh_df = mf.with_df._rsh_df[key]
-            else:
-                rsh_df = mf.with_df._rsh_df[key] = mf.with_df.copy().reset()
-            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()
-        # for DFT object
-        if hasattr(mf, '_numint'):
-            ni = mf._numint
-            rks.initialize_grids(mf, mf.mol, dm0)
-            ni.build(mf.mol, mf.grids.coords)
-            mf._numint.xcfuns = numint._init_xcfuns(mf.xc, dm0.ndim==3)
-    dm0 = cupy.asarray(dm0)
-    return
-
 def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
     '''For the given SCF object, update the J, K matrix constructor with
     corresponding density fitting integrals.
@@ -143,8 +114,6 @@ class _DFHF:
         self.with_df.reset(mol)
         return super().reset(mol)
 
-    init_workflow = init_workflow
-
     def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
                omega=None):
         if dm is None: dm = self.make_rdm1()
@@ -517,23 +486,6 @@ def get_jk(dfobj, dms_tag, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-
             vk = vk.get()
         return vj, vk
 
-def _get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True,
-            direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
-            omega=None):
-    if omega is None:
-        return get_jk(dfobj, dm, hermi, with_j, with_k, direct_scf_tol)
-
-    # A temporary treatment for RSH-DF integrals
-    key = '%.6f' % omega
-    if key in dfobj._rsh_df:
-        rsh_df = dfobj._rsh_df[key]
-    else:
-        rsh_df = dfobj._rsh_df[key] = dfobj.copy().reset()
-        logger.info(dfobj, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
-
-    with rsh_df.mol.with_range_coulomb(omega):
-        return get_jk(rsh_df, dm, hermi, with_j, with_k, direct_scf_tol)
-
 def get_j(dfobj, dm, hermi=1, direct_scf_tol=1e-13):
     intopt = getattr(dfobj, 'intopt', None)
     if intopt is None:

{gpu4pyscf-1.3.2 → gpu4pyscf-1.4.1}/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.1}/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.1}/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