gpu4pyscf 1.4.2__tar.gz → 1.4.3__tar.gz

{gpu4pyscf-1.4.2 → gpu4pyscf-1.4.3}/CHANGELOG

@@ -1,3 +1,24 @@
+v1.4.3 (2025-08-20)
+-------------------
+* New Features
+  - Geometry optimization for excited states using TDDFT-ris methods.
+  - Non-adiabatic coupling vectors for TDDFT-ris methods.
+  - Analytical gradients for DFT+U with k-point sampling.
+  - Stress tensor calculations for semi-local DFT with k-point sampling and at the gamma-point.
+  - ASE interface to support crystal lattice optimization.
+  - Multigrid v2 algorithm for meta-GGA functionals.
+* Improvements
+  - GPU kernels for PBC overlap and kinetic integrals.
+  - Reduced GPU memory usage for TDDFT-ris by storing tensors in host memory.
+  - In PBC methods, scaled k-points (fractional coordinates) are stored to
+    simplify lattice optimization calculations.
+  - A preconditioned Krylov solver to accelerate convergence in TDDFT and
+    dynamic polarizability calculations.
+* Fixes
+  - Basis decontraction issue for d and f orbitals.
+  - A bug in Multigrid v2 algorithm related to non-orthogonal lattices.
+  - Incorrect virtual orbital energies when level_shift was enabled.
+
 v1.4.2 (2025-07-20)
 -------------------
 * New Features

{gpu4pyscf-1.4.2/gpu4pyscf.egg-info → gpu4pyscf-1.4.3}/PKG-INFO

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

{gpu4pyscf-1.4.2 → gpu4pyscf-1.4.3}/gpu4pyscf/__init__.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '1.4.2'
+__version__ = '1.4.3'
+
+from . import _patch_pyscf
 
 from . import lib, grad, hessian, solvent, scf, dft, tdscf, nac
 

gpu4pyscf-1.4.3/gpu4pyscf/_patch_pyscf.py

@@ -0,0 +1,164 @@
+# Copyright 2025 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import cupy
+import pyscf
+
+if int(pyscf.__version__.split('.')[1]) <= 10:
+    def _fftdf_to_gpu(self):
+        from gpu4pyscf.pbc.df.fft import FFTDF
+        return FFTDF(self.cell, self.kpts)
+    from pyscf.pbc.df.fft import FFTDF
+    FFTDF.to_gpu = _fftdf_to_gpu
+
+    def _aftdf_to_gpu(self):
+        from gpu4pyscf.pbc.df.aft import AFTDF
+        return AFTDF(self.cell, self.kpts)
+    from pyscf.pbc.df.aft import AFTDF
+    AFTDF.to_gpu = _aftdf_to_gpu
+
+    def _gdf_to_gpu(self):
+        from gpu4pyscf.pbc.df.df import GDF
+        return GDF(self.cell, self.kpts)
+    from pyscf.pbc.df.df import GDF
+    GDF.to_gpu = _gdf_to_gpu
+
+    # patch PySCF Cell class, updating lattice parameters is not avail in pyscf 2.10
+    from pyscf import lib
+    from pyscf.lib import logger
+    from pyscf.gto import mole
+    from pyscf.pbc.gto.cell import Cell
+    def set_geom_(self, atoms_or_coords=None, unit=None, symmetry=None,
+                  a=None, inplace=True):
+        '''Update geometry and lattice parameters
+
+        Kwargs:
+            atoms_or_coords : list, str, or numpy.ndarray
+                When specified in list or str, it is processed as the Mole.atom
+                attribute. If inputing a (N, 3) numpy array, this array
+                represents the coordinates of the atoms in the molecule.
+            a : list, str, or numpy.ndarray
+                If specified, it is assigned to the cell.a attribute. Its data
+                format should be the same to cell.a
+            unit : str
+                The unit for the input `atoms_or_coords` and `a`. If specified,
+                cell.unit will be updated to this value. If not provided, the
+                current cell.unit will be used for the two inputs.
+            symmetry : bool
+                Whether to enable space_group_symmetry. It is a reserved input
+                argument. This functionality is not supported yet.
+            inplace : bool
+                Whether to overwrite the existing Mole object.
+        '''
+        if inplace:
+            cell = self
+        else:
+            cell = self.copy(deep=False)
+            cell._env = cell._env.copy()
+
+        if unit is not None and cell.unit != unit:
+            if isinstance(unit, str):
+                if mole.is_au(unit):
+                    _unit = 1.
+                else:
+                    _unit = lib.param.BOHR
+            else:
+                _unit = unit
+            if a is None:
+                a = self.lattice_vectors() * _unit
+            if atoms_or_coords is None:
+                atoms_or_coords = self.atom_coords() * _unit
+
+        if a is not None:
+            logger.info(cell, 'Set new lattice vectors')
+            logger.info(cell, '%s', a)
+            cell.a = a
+            if cell._mesh_from_build:
+                cell.mesh = None
+            if cell._rcut_from_build:
+                cell.rcut = None
+            cell._built = False
+        cell.enuc = None
+
+        if atoms_or_coords is not None:
+            cell = mole.MoleBase.set_geom_(cell, atoms_or_coords, unit, symmetry)
+        if not cell._built:
+            cell.build(False, False)
+        return cell
+    Cell.set_geom_ = set_geom_
+
+    def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None, discard=True):
+        '''This version employs more strict criteria when discarding images in lattice sum.
+        It can be replaced by the built-in version available in PySCF 2.10.
+        '''
+        if dimension is None:
+            # For atoms near the boundary of the cell, it is necessary (even in low-
+            # dimensional systems) to include lattice translations in all 3 dimensions.
+            if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
+                dimension = cell.dimension
+            else:
+                dimension = 3
+        if rcut is None:
+            rcut = cell.rcut
+
+        if dimension == 0 or rcut <= 0 or cell.natm == 0:
+            return np.zeros((1, 3))
+
+        a = cell.lattice_vectors()
+
+        scaled_atom_coords = cell.get_scaled_atom_coords()
+        atom_boundary_max = scaled_atom_coords[:,:dimension].max(axis=0)
+        atom_boundary_min = scaled_atom_coords[:,:dimension].min(axis=0)
+        if (np.any(atom_boundary_max > 1) or np.any(atom_boundary_min < -1)):
+            atom_boundary_max[atom_boundary_max > 1] = 1
+            atom_boundary_min[atom_boundary_min <-1] = -1
+        ovlp_penalty = atom_boundary_max - atom_boundary_min
+        dR = ovlp_penalty.dot(a[:dimension])
+        dR_basis = np.diag(dR)
+
+        # Search the minimal x,y,z requiring |x*a[0]+y*a[1]+z*a[2]+dR|^2 > rcut^2
+        # Ls boundary should be derived by decomposing (a, Rij) for each atom-pair.
+        # For reasons unclear, the so-obtained Ls boundary seems not large enough.
+        # The upper-bound of the Ls boundary is generated by find_boundary function.
+        def find_boundary(a):
+            aR = np.vstack([a, dR_basis])
+            r = np.linalg.qr(aR.T)[1]
+            ub = (rcut + abs(r[2,3:]).sum()) / abs(r[2,2])
+            return ub
+
+        xb = find_boundary(a[[1,2,0]])
+        if dimension > 1:
+            yb = find_boundary(a[[2,0,1]])
+        else:
+            yb = 0
+        if dimension > 2:
+            zb = find_boundary(a)
+        else:
+            zb = 0
+        bounds = np.ceil([xb, yb, zb]).astype(int)
+        Ts = lib.cartesian_prod((np.arange(-bounds[0], bounds[0]+1),
+                                 np.arange(-bounds[1], bounds[1]+1),
+                                 np.arange(-bounds[2], bounds[2]+1)))
+        Ls = np.dot(Ts[:,:dimension], a[:dimension])
+
+        if discard and len(Ls) > 1:
+            r = cell.atom_coords()
+            rr = r[:,None] - r
+            dist_max = np.linalg.norm(rr, axis=2).max()
+            Ls_mask = np.linalg.norm(Ls, axis=1) < rcut + dist_max
+            Ls = Ls[Ls_mask]
+        return np.asarray(Ls, order='C')
+    # Patch the get_lattice_Ls for pyscf-2.9 or older
+    Cell.get_lattice_Ls = get_lattice_Ls

{gpu4pyscf-1.4.2 → gpu4pyscf-1.4.3}/gpu4pyscf/df/df.py

@@ -26,6 +26,7 @@ from gpu4pyscf.lib.cupy_helper import (
 from gpu4pyscf.df import int3c2e, df_jk
 from gpu4pyscf.df import int3c2e_bdiv
 from gpu4pyscf.lib import logger
+from gpu4pyscf.lib import utils
 from gpu4pyscf import __config__
 from gpu4pyscf.__config__ import _streams, num_devices
 
@@ -38,7 +39,6 @@ LINEAR_DEP_THR = incore.LINEAR_DEP_THR
 GROUP_SIZE = 256
 
 class DF(lib.StreamObject):
-    from gpu4pyscf.lib.utils import to_gpu, device
 
     _keys = {'intopt', 'nao', 'naux', 'cd_low', 'mol', 'auxmol', 'use_gpu_memory'}
 
@@ -71,10 +71,12 @@ class DF(lib.StreamObject):
             self.reset()
             self._auxbasis = x
 
+    to_gpu = utils.to_gpu
+    device = utils.device
+
     def to_cpu(self):
-        from gpu4pyscf.lib.utils import to_cpu
-        obj = to_cpu(self)
-        return obj.reset()
+        from pyscf.df.df import DF
+        return utils.to_cpu(self, out=DF(self.mol, auxbasis=self.auxbasis))
 
     def build(self, direct_scf_tol=1e-14, omega=None):
         mol = self.mol

{gpu4pyscf-1.4.2 → gpu4pyscf-1.4.3}/gpu4pyscf/df/df_jk.py

@@ -25,7 +25,7 @@ from gpu4pyscf.lib import logger
 from gpu4pyscf.lib.cupy_helper import (
     contract, transpose_sum, reduce_to_device, tag_array)
 from gpu4pyscf.dft import rks, uks, numint
-from gpu4pyscf.scf import hf, uhf
+from gpu4pyscf.scf import hf, uhf, rohf
 from gpu4pyscf.df import df, int3c2e
 from gpu4pyscf.__config__ import _streams, num_devices
 
@@ -192,13 +192,24 @@ class _DFHF:
         if dm is None: dm = self.make_rdm1()
         assert not self.direct_scf
 
+        if isinstance(self, rohf.ROHF):
+            if getattr(dm, 'mo_coeff', None) is not None:
+                mo_coeff = cupy.repeat(dm.mo_coeff[None], 2, axis=0)
+                mo_occ = cupy.asarray([dm.mo_occ>0, dm.mo_occ==2],
+                                      dtype=numpy.double)
+                if dm.ndim == 2:  # RHF DM
+                    dm = cupy.repeat(dm[None]*.5, 2, axis=0)
+                dm = tag_array(dm, mo_coeff=mo_coeff, mo_occ=mo_occ)
+            elif dm.ndim == 2:  # RHF DM
+                dm = cupy.repeat(dm[None]*.5, 2, axis=0)
+
         # for DFT
         if isinstance(self, rks.KohnShamDFT):
             t0 = logger.init_timer(self)
             rks.initialize_grids(self, mol, dm)
             ni = self._numint
-            if dm.ndim == 2: # RKS
-                n, exc, vxc = ni.nr_rks(mol, self.grids, self.xc, dm)
+            if isinstance(self, (uhf.UHF, rohf.ROHF)): # UKS
+                n, exc, vxc = ni.nr_uks(mol, self.grids, self.xc, dm)
                 logger.debug(self, 'nelec by numeric integration = %s', n)
                 if self.do_nlc():
                     if ni.libxc.is_nlc(self.xc):
@@ -206,30 +217,31 @@ class _DFHF:
                     else:
                         assert ni.libxc.is_nlc(self.nlc)
                         xc = self.nlc
-                    n, enlc, vnlc = ni.nr_nlc_vxc(mol, self.nlcgrids, xc, dm)
+                    n, enlc, vnlc = ni.nr_nlc_vxc(mol, self.nlcgrids, xc, dm[0]+dm[1])
                     exc += enlc
                     vxc += vnlc
                     logger.debug(self, 'nelec with nlc grids = %s', n)
-                t0 = logger.timer_debug1(self, 'vxc tot', *t0)
+                t0 = logger.timer(self, 'vxc', *t0)
 
                 if not ni.libxc.is_hybrid_xc(self.xc):
-                    vj = self.get_j(mol, dm, hermi)
+                    vj = self.get_j(mol, dm[0]+dm[1], hermi)
                     vxc += vj
                 else:
                     omega, alpha, hyb = ni.rsh_and_hybrid_coeff(self.xc, spin=mol.spin)
                     vj, vk = self.get_jk(mol, dm, hermi)
+                    vj = vj[0] + vj[1]
                     vxc += vj
                     vk *= hyb
-                    if omega != 0:
-                        vklr = self.get_k(mol, dm, hermi, omega=abs(omega))
+                    if abs(omega) > 1e-10:
+                        vklr = self.get_k(mol, dm, hermi, omega=omega)
                         vklr *= (alpha - hyb)
                         vk += vklr
-                    vxc -= vk * .5
-                    exc -= cupy.einsum('ij,ji', dm, vk).real * .25
-                ecoul = cupy.einsum('ij,ji', dm, vj).real * .5
+                    vxc -= vk
+                    exc -= cupy.einsum('sij,sji->', dm, vk).real * .5
+                ecoul = cupy.einsum('sij,ji->', dm, vj).real * .5
 
-            elif dm.ndim == 3: # UKS
-                n, exc, vxc = ni.nr_uks(mol, self.grids, self.xc, dm)
+            elif isinstance(self, hf.RHF):
+                n, exc, vxc = ni.nr_rks(mol, self.grids, self.xc, dm)
                 logger.debug(self, 'nelec by numeric integration = %s', n)
                 if self.do_nlc():
                     if ni.libxc.is_nlc(self.xc):
@@ -237,39 +249,41 @@ class _DFHF:
                     else:
                         assert ni.libxc.is_nlc(self.nlc)
                         xc = self.nlc
-                    n, enlc, vnlc = ni.nr_nlc_vxc(mol, self.nlcgrids, xc, dm[0]+dm[1])
+                    n, enlc, vnlc = ni.nr_nlc_vxc(mol, self.nlcgrids, xc, dm)
                     exc += enlc
                     vxc += vnlc
                     logger.debug(self, 'nelec with nlc grids = %s', n)
-                t0 = logger.timer(self, 'vxc', *t0)
+                t0 = logger.timer_debug1(self, 'vxc tot', *t0)
 
                 if not ni.libxc.is_hybrid_xc(self.xc):
-                    vj = self.get_j(mol, dm[0]+dm[1], hermi)
+                    vj = self.get_j(mol, dm, hermi)
                     vxc += vj
                 else:
                     omega, alpha, hyb = ni.rsh_and_hybrid_coeff(self.xc, spin=mol.spin)
                     vj, vk = self.get_jk(mol, dm, hermi)
-                    vj = vj[0] + vj[1]
                     vxc += vj
                     vk *= hyb
-                    if abs(omega) > 1e-10:
-                        vklr = self.get_k(mol, dm, hermi, omega=omega)
+                    if omega != 0:
+                        vklr = self.get_k(mol, dm, hermi, omega=abs(omega))
                         vklr *= (alpha - hyb)
                         vk += vklr
-                    vxc -= vk
-                    exc -= cupy.einsum('sij,sji->', dm, vk).real * .5
-                ecoul = cupy.einsum('sij,ji->', dm, vj).real * .5
+                    vxc -= vk * .5
+                    exc -= cupy.einsum('ij,ji', dm, vk).real * .25
+                ecoul = cupy.einsum('ij,ji', dm, vj).real * .5
+
+            else:
+                raise NotImplementedError("DF only supports R/U/RO KS.")
             t0 = logger.timer_debug1(self, 'jk total', *t0)
             return tag_array(vxc, ecoul=ecoul, exc=exc, vj=None, vk=None)
 
-        if dm.ndim == 2:
-            vj, vk = self.get_jk(mol, dm, hermi=hermi)
-            return vj - vk * .5
-        elif dm.ndim == 3:
+        if isinstance(self, (uhf.UHF, rohf.ROHF)):
             vj, vk = self.get_jk(mol, dm, hermi=hermi)
             return vj[0] + vj[1] - vk
+        elif isinstance(self, hf.RHF):
+            vj, vk = self.get_jk(mol, dm, hermi=hermi)
+            return vj - vk * .5
         else:
-            raise NotImplementedError("Please check the dimension of the density matrix, it should not reach here.")
+            raise NotImplementedError("DF only supports R/U/RO HF.")
 
     def to_cpu(self):
         obj = self.undo_df().to_cpu().density_fit()

gpu4pyscf-1.4.3/gpu4pyscf/df/grad/tdrks_ris.py

@@ -0,0 +1,39 @@
+# Copyright 2021-2025 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gpu4pyscf.df import int3c2e, df
+from gpu4pyscf.df.grad import tdrhf as tdrhf_grad_df
+from gpu4pyscf.tdscf import ris
+from gpu4pyscf.grad import tdrks_ris as tdrks_ris_grad
+from gpu4pyscf import __config__
+
+class Gradients(tdrks_ris_grad.Gradients):
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'with_df', 'auxbasis_response'}
+    def __init__(self, td):
+        # Whether to include the response of DF auxiliary basis when computing
+        # nuclear gradients of J/K matrices
+        tdrks_ris_grad.Gradients.__init__(self, td)
+
+    auxbasis_response = True
+    get_jk = tdrhf_grad_df.get_jk
+
+    def check_sanity(self):
+        assert isinstance(self.base._scf, df.df_jk._DFHF)
+        assert isinstance(self.base, ris.TDDFT) or isinstance(self.base, ris.TDA)
+
+    get_veff = tdrhf_grad_df.get_veff
+
+Grad = Gradients

gpu4pyscf-1.4.3/gpu4pyscf/df/nac/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import tdrhf, tdrks, tdrks_ris

gpu4pyscf-1.4.3/gpu4pyscf/df/nac/tdrhf.py

@@ -0,0 +1,38 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gpu4pyscf.nac import tdrhf as tdrhf_nac
+from gpu4pyscf.df.grad.tdrhf import get_jk, get_veff
+from gpu4pyscf.tdscf import rhf as tdrhf
+from gpu4pyscf.lib import logger
+from gpu4pyscf.df import df
+from gpu4pyscf.lib.cupy_helper import tag_array
+
+
+class NAC(tdrhf_nac.NAC):
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'with_df', 'auxbasis_response'}
+    def __init__(self, td):
+        # Whether to include the response of DF auxiliary basis when computing
+        # nuclear gradients of J/K matrices
+        tdrhf_nac.NAC.__init__(self, td)
+
+    auxbasis_response = True
+    get_jk = get_jk
+    get_veff = get_veff
+
+    def check_sanity(self):
+        assert isinstance(self.base._scf, df.df_jk._DFHF)
+        assert isinstance(self.base, tdrhf.TDHF) or isinstance(self.base, tdrhf.TDA)

gpu4pyscf-1.4.3/gpu4pyscf/df/nac/tdrks.py

@@ -0,0 +1,36 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gpu4pyscf.nac import tdrks as tdrks_nac
+from gpu4pyscf.df.nac import tdrhf as tdrhf_nac_df
+from gpu4pyscf.df import df
+from gpu4pyscf.tdscf import rks as tdrks
+
+class NAC(tdrks_nac.NAC):
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'with_df', 'auxbasis_response'}
+    def __init__(self, td):
+        # Whether to include the response of DF auxiliary basis when computing
+        # nuclear gradients of J/K matrices
+        tdrks_nac.NAC.__init__(self, td)
+
+    auxbasis_response = True
+    get_jk = tdrhf_nac_df.get_jk
+
+    def check_sanity(self):
+        assert isinstance(self.base._scf, df.df_jk._DFHF)
+        assert isinstance(self.base, tdrks.TDDFT) or isinstance(self.base, tdrks.TDA)
+
+    get_veff = tdrhf_nac_df.get_veff

gpu4pyscf-1.4.3/gpu4pyscf/df/nac/tdrks_ris.py

@@ -0,0 +1,35 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gpu4pyscf.nac import tdrks_ris as tdrks_ris_nac
+from gpu4pyscf.df.nac import tdrhf as tdrhf_nac_df
+from gpu4pyscf.df import df
+from gpu4pyscf.tdscf import ris
+
+class NAC(tdrks_ris_nac.NAC):
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'with_df', 'auxbasis_response'}
+    def __init__(self, td):
+        # Whether to include the response of DF auxiliary basis when computing
+        # nuclear gradients of J/K matrices
+        tdrks_ris_nac.NAC.__init__(self, td)
+
+    auxbasis_response = True
+    get_veff = tdrhf_nac_df.get_veff
+    get_jk = tdrhf_nac_df.get_jk
+
+    def check_sanity(self):
+        assert isinstance(self.base._scf, df.df_jk._DFHF)
+        assert isinstance(self.base, ris.TDDFT) or isinstance(self.base, ris.TDA)

{gpu4pyscf-1.4.2 → gpu4pyscf-1.4.3}/gpu4pyscf/dft/numint.py

@@ -1528,18 +1528,24 @@ def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=0,
 def batch_square(a):
     return a[0]**2 + a[1]**2 + a[2]**2
 
-def eval_xc_eff(ni, xc_code, rho, deriv=1, omega=None, xctype=None, verbose=None):
+def eval_xc_eff(ni, xc_code, rho, deriv=1, omega=None, xctype=None,
+                verbose=None, spin=None):
     '''
     Different from PySCF, this function employ cuda version libxc
     '''
     if omega is None: omega = ni.omega
     if xctype is None: xctype = ni._xc_type(xc_code)
 
-    spin_polarized = rho.ndim >= 2 and rho.shape[0] == 2
-    xcfuns = ni._init_xcfuns(xc_code, spin_polarized)
+    if spin is None:
+        spin_polarized = rho.ndim >= 2 and rho.shape[0] == 2
+        if spin_polarized:
+            spin = 1
+        else:
+            spin = 0
+    xcfuns = ni._init_xcfuns(xc_code, spin)
 
     inp = {}
-    if not spin_polarized:
+    if spin == 0:
         assert rho.dtype == np.float64
         if xctype == 'LDA':
             inp['rho'] = rho.ravel()
@@ -1599,23 +1605,16 @@ def eval_xc_eff(ni, xc_code, rho, deriv=1, omega=None, xctype=None, verbose=None
     kxc = None
 
     exc = ret_full["zk"]
-    if not spin_polarized:
-        vxc = [ret_full[label] for label in vxc_labels if label in ret_full]
-        if do_fxc:
-            fxc = [ret_full[label] for label in fxc_labels if label in ret_full]
-        if do_kxc:
-            kxc = [ret_full[label] for label in kxc_labels if label in ret_full]
-    else:
-        vxc = [ret_full[label] for label in vxc_labels if label in ret_full]
-        if do_fxc:
-            fxc = [ret_full[label] for label in fxc_labels if label in ret_full]
-        if do_kxc:
-            kxc = [ret_full[label] for label in kxc_labels if label in ret_full]
+    vxc = [ret_full[label] for label in vxc_labels if label in ret_full]
+    if do_fxc:
+        fxc = [ret_full[label] for label in fxc_labels if label in ret_full]
+    if do_kxc:
+        kxc = [ret_full[label] for label in kxc_labels if label in ret_full]
     if do_kxc:
-        kxc = xc_deriv.transform_kxc(rho, fxc, kxc, xctype, spin_polarized)
+        kxc = xc_deriv.transform_kxc(rho, fxc, kxc, xctype, spin)
     if do_fxc:
-        fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin_polarized)
-    vxc = xc_deriv.transform_vxc(rho, vxc, xctype, spin_polarized)
+        fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin)
+    vxc = xc_deriv.transform_vxc(rho, vxc, xctype, spin)
     return exc, vxc, fxc, kxc
 
 def _init_xcfuns(xc_code, spin):