Add SF dispersion gradients and improve geometry scanners

.github/workflows/ci.yml

@@ -31,5 +31,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - run: python -m pip install -e ".[dev]"
+      - name: Check dispersion and optimization dependencies
+        run: python -c "from pyscf.dispersion import dftd3, dftd4; import geometric"
       - run: python -m ruff check .
       - run: python -m pytest src -q

examples/grad/02_sftda_dispersion.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+# Copyright 2026 The NEST 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.
+
+"""
+SF-TDA total energies and gradients with D3/D4 dispersion.
+Required dependencies: pyscf-dispersion
+"""
+
+from pyscf import gto
+from nest import sftda  # Register SFTDA/SFTDDFT on PySCF reference objects.
+
+mol = gto.M(
+    atom="O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587",
+    basis="sto-3g",
+    spin=2,
+    verbose=3,
+)
+mf = mol.UKS(xc="PBE")
+# Set dispersion BEFORE SCF. Parameters are selected using mf.xc.
+# 'd3bj' or 'd3bj2b': D3(BJ), two-body only
+# 'd3bjatm':          D3(BJ), including the ATM three-body term
+# 'd4':              D4, including the ATM three-body term
+mf.disp = "d4"
+mf.kernel()
+
+# The same gradient interface is available for mf.SFTDDFT().
+td = mf.SFTDA().set(nstates=2, collinear="col").run()
+print("TD roots relative to the reference (Hartree):", td.e)
+print("Total state energies including dispersion (Hartree):", td.e_tot)
+
+# state=1 means the first TD root, which can lie below the SCF reference.
+# state=0 requests the SCF reference gradient instead.
+grad = td.Gradients().kernel(state=1)
+print("First-root total gradient including dispersion (Hartree/Bohr):")
+print(grad)
+
+# For geometry optimization, the scanner returns consistent total energy and
+# gradient and recomputes dispersion when the geometry changes.
+# It reuses the previous SCF density and TD amplitudes as initial guesses.
+# The root index does not track state character.
+scanner = td.Gradients().as_scanner(state=1)
+energy, gradient = scanner(mol)
+print("Scanner total energy (Hartree):", energy)
+# mf.e_tot, td.e_tot, and the total gradients already include dispersion.

examples/grad/03_sftda_geomopt.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+# Copyright 2026 The NEST 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.
+
+"""
+Optimize the first SF-TDA root with geomeTRIC and D4 dispersion.
+Required dependencies: geometric, pyscf-dispersion
+"""
+
+from pyscf import gto
+from pyscf.geomopt import geometric_solver
+from nest import sftda  # Register SF-TDA on PySCF reference objects.
+
+mol = gto.M(
+    atom="O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587",
+    basis="6-31g",
+    spin=2,
+)
+mf = mol.UKS(xc="B3LYP").set(disp="d4")
+
+# The TD constructor runs SCF if needed; the optimizer triggers the TD calculation.
+td = mf.SFTDA().set(extype=1, collinear='mcol', nstates=3)
+
+# It reuses the SCF density and projects old TD amplitudes onto the new MO basis.
+# A fixed root index does not track electronic character through root crossings.
+scanner = td.Gradients().as_scanner(state=1)
+optimizer = geometric_solver.GeometryOptimizer(scanner)
+optimizer.max_cycle = 50
+mol_eq = optimizer.kernel()
+
+print("Geometry optimization converged:", optimizer.converged)
+print("Optimized geometry (Angstrom):")
+print(mol_eq.tostring(format="xyz"))

examples/sftda/02_sftddft_roks.py

@@ -16,8 +16,9 @@
 '''
 Spin-flip TDDFT/TDA examples with ROKS reference.
 
-nest.sftda.SFTDDFT/SFTDA can also be applied to ROKS objects,
-or equivalently to mf.to_uks() objects with methods.
+ROKS and UKS share the mf.SFTDDFT()/mf.SFTDA() interface.
+The nest.sftda.SFTDDFT(mf)/SFTDA(mf) function interface is also available.
+Spin-flip TDDFT gradients for RO references are NOT implemented.
 '''
 
 import numpy as np
@@ -43,7 +44,7 @@ def print_header(title):
 
 # Spin-flip down TDDFT
 print_header("CALCULATION 1: Spin-Flip-Down TDDFT")
-sfd_tddft = sftda.SFTDDFT(mf)  # same as mf.to_uks().SFTDDFT()
+sfd_tddft = mf.SFTDDFT()
 sfd_tddft.nstates = 5
 sfd_tddft.extype = 1  # 1 for spin-flip-down excitations
 sfd_tddft.collinear_samples = 20
@@ -89,7 +90,7 @@ def norm_xy(z):
 
 # Spin-flip down TDA
 print_header("CALCULATION 2: Spin-Flip-Down TDA")
-sfd_tda = sftda.SFTDA(mf)  # same as mf.to_uks().SFTDA()
+sfd_tda = mf.SFTDA()
 sfd_tda.extype = 1
 sfd_tda.nstates = 5
 sfd_tda.collinear = 'col'  # Use collinear functional

pyproject.toml

@@ -13,6 +13,8 @@ license-files = ["LICENSE"]
 dependencies = [
     "mcfun>=0.2.6",
     "pyscf>=2.13.0",
+    "pyscf-dispersion",
+    "geometric",
     "numpy",
     "scipy",
 ]

src/nest/grad/tduks_sf.py

@@ -357,11 +357,78 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True, with_k
     return f1vo, f1oo, v1ao, k1ao
 
 
+class TDSCF_GradScanner(tdrhf_grad.TDSCF_GradScanner):
+    @property
+    def converged(self):
+        return self.base._scf.converged and self.base.converged[self.state - 1]
+
+
 class Gradients(tdrhf_grad.Gradients):
     cphf_max_cycle = tdrhf_grad.Gradients.cphf_max_cycle + 20
 
+    def as_scanner(self, state=1):
+        if isinstance(self, lib.GradScanner):
+            return self
+        if state == 0:
+            return self.base._scf.nuc_grad_method().as_scanner()
+        name = self.__class__.__name__ + TDSCF_GradScanner.__name_mixin__
+        return lib.set_class(TDSCF_GradScanner(self, state),
+                             (TDSCF_GradScanner, self.__class__), name)
+
     @lib.with_doc(grad_elec.__doc__)
-    def grad_elec(self, xy, singlet=None, atmlst=None):
+    def grad_elec(self, xy, atmlst=None):
         return grad_elec(self, xy, atmlst, self.max_memory, self.verbose)
 
+    def dump_flags(self, verbose=None):
+        super().dump_flags(verbose)
+        log = logger.new_logger(self, verbose)
+        td = self.base
+        log.info('extype = %s', td.extype)
+        log.info('collinear = %s', td.collinear)
+        if td.collinear == 'mcol':
+            log.info('collinear_samples = %s', td.collinear_samples)
+        mf = td._scf
+        log.info('dispersion = %s', (mf.disp or mf.xc) if mf.do_disp() else False)
+        return self
+
+    def kernel(self, xy=None, state=None, atmlst=None):
+        """Rewrite the `kernel` method to include the reference's D3/D4 correction."""
+        if xy is None:
+            if state is None:
+                state = self.state
+            else:
+                self.state = state
+            if state == 0:
+                logger.warn(self, 'state=0: computing the SCF reference gradient.')
+                # The SCF gradient kernel already includes dispersion.
+                return self.base._scf.nuc_grad_method().kernel(atmlst=atmlst)
+
+        if xy is None:
+            if self.base.xy is None:
+                self.base.run()
+            xy = self.base.xy[state - 1]
+
+        if atmlst is None:
+            atmlst = self.atmlst
+        else:
+            self.atmlst = atmlst
+
+        if self.verbose >= logger.WARN:
+            self.check_sanity()
+        if self.verbose >= logger.INFO:
+            self.dump_flags()
+
+        de = self.grad_elec(xy, atmlst)
+        self.de = de + self.grad_nuc(atmlst=atmlst)
+        if self.mol.symmetry:
+            self.de = self.symmetrize(self.de, atmlst)
+        mf = self.base._scf
+        if mf.do_disp():
+            dispersion = mf.nuc_grad_method().get_dispersion()
+            if atmlst is not None:
+                dispersion = dispersion[np.asarray(atmlst, dtype=int)]
+            self.de += dispersion
+        self._finalize()
+        return self.de
+
 Grad = Gradients

src/nest/grad/tests/test_tduks_sf_grad.py

@@ -18,6 +18,11 @@
 from pyscf import gto
 from nest import sftda
 
+try:
+    from pyscf.dispersion import dftd3, dftd4
+except ImportError:
+    dftd3 = dftd4 = None
+
 
 def cal_exact_sf_tda_gradient(mf, extype=1, collinear='mcol', collinear_samples=20, state=1):
     a, b = sftda.uhf_sf.get_ab_sf(mf, collinear=collinear, collinear_samples=collinear_samples)
@@ -234,6 +239,67 @@ def test_mcol_tpss(self):
         self.assertAlmostEqual(abs(grad_iter - ref).max(), 0, delta=1e-5)
 
 
+    @unittest.skipIf(dftd3 is None, 'pyscf-dispersion is not installed')
+    def test_col_d3bj(self):
+        mf = self.mol.UKS(xc='PBE').run()
+        td = mf.SFTDA().set(nstates=2, collinear='col').run()
+        roots = td.e.copy()
+        energies = td.e_tot.copy()
+        grad = td.Gradients().kernel(state=1)
+
+        mf.disp = 'd3bj'
+        mf.e_tot = mf.energy_tot()
+        td.run()
+        disp_energy = mf.get_dispersion()
+        disp_grad = mf.nuc_grad_method().get_dispersion()
+        self.assertAlmostEqual(abs(td.e - roots).max(), 0, delta=1e-6)
+        self.assertAlmostEqual(abs(td.e_tot - energies - disp_energy).max(), 0, delta=1e-6)
+        corrected = td.Gradients().kernel(state=1)
+        self.assertAlmostEqual(abs(corrected - grad - disp_grad).max(), 0, delta=1e-6)
+
+    @unittest.skipIf(dftd4 is None, 'pyscf-dispersion is not installed')
+    def test_mcol_d4(self):
+        mf = self.mol.UKS(xc='PBE').run()
+        td = mf.SFTDA().set(nstates=2, collinear='mcol').run()
+        roots = td.e.copy()
+        energies = td.e_tot.copy()
+        grad = td.Gradients().kernel(state=1)
+
+        mf.disp = 'd4'
+        mf.e_tot = mf.energy_tot()
+        td.run()
+        disp_energy = mf.get_dispersion()
+        disp_grad = mf.nuc_grad_method().get_dispersion()
+        self.assertAlmostEqual(abs(td.e - roots).max(), 0, delta=1e-6)
+        self.assertAlmostEqual(abs(td.e_tot - energies - disp_energy).max(), 0, delta=1e-6)
+        corrected = td.Gradients().kernel(state=1)
+        self.assertAlmostEqual(abs(corrected - grad - disp_grad).max(), 0, delta=1e-6)
+
+    @unittest.skipIf(dftd4 is None, 'pyscf-dispersion is not installed')
+    def test_d4_gradient_scanner(self):
+        # Finite differences require tightly converged energies at both geometries.
+        mf = self.mol.UKS(xc='PBE').set(disp='d4', conv_tol=1e-14, conv_tol_grad=1e-9)
+        mf.grids.level = 4
+        mf.run()
+        td = mf.SFTDA().set(nstates=1, collinear='col', conv_tol=1e-9).run()
+        grad = td.Gradients().kernel(state=1)
+        scanner = td.Gradients().as_scanner(state=1)
+        energy, scanner_grad = scanner(self.mol)
+        self.assertTrue(scanner.converged)
+        self.assertAlmostEqual(energy, td.e_tot[0], delta=1e-6)
+        self.assertAlmostEqual(abs(scanner_grad - grad).max(), 0, delta=1e-6)
+
+        # Differentiate the corrected total energy along one nuclear coordinate.
+        step = 1e-3  # Bohr
+        coords = self.mol.atom_coords()
+        coords[1, 1] += step
+        e_plus = scanner(self.mol.set_geom_(coords, unit='Bohr', inplace=False))[0]
+        coords[1, 1] -= 2 * step
+        e_minus = scanner(self.mol.set_geom_(coords, unit='Bohr', inplace=False))[0]
+        derivative = (e_plus - e_minus) / (2 * step)
+        # Allow finite-difference and numerical integration errors, in Hartree/Bohr.
+        self.assertAlmostEqual(derivative, grad[1, 1], delta=1e-5)
+
 if __name__ == '__main__':
     print('Full tests for SF-TDA and SF-TDDFT analytic gradients')
     unittest.main()

src/nest/sftda/__init__.py

@@ -14,33 +14,50 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pyscf import scf, dft
+from pyscf import scf
 from nest.sftda import uhf_sf
 
 
-def TDA_SF(mf):
+def TDA_SF(mf, extype=1, collinear="mcol", collinear_samples=20):
+    """Construct SF-TDA from a UHF/UKS or ROHF/ROKS reference."""
     mf = mf.remove_soscf()
-    if isinstance(mf, scf.rohf.ROHF) or isinstance(mf, scf.hf_symm.SymAdaptedROHF):
-        if isinstance(mf, dft.roks.ROKS) or isinstance(mf, dft.rks_symm.SymAdaptedROKS):
+    restricted = isinstance(mf, scf.rohf.ROHF)
+    if restricted:
+        if mf.mo_coeff is None:
+            mf.run()
+        if isinstance(mf, scf.hf.KohnShamDFT):
             mf = mf.to_uks()
         else:
             mf = mf.to_uhf()
-    return mf.TDA_SF()
+    td = uhf_sf.TDA_SF(mf, extype, collinear, collinear_samples)
+    td._ro_reference = restricted
+    return td
 
 
-def TDDFT_SF(mf):
+def TDDFT_SF(mf, extype=1, collinear="mcol", collinear_samples=20):
+    """Construct SF-TDDFT from a UHF/UKS or ROHF/ROKS reference."""
     mf = mf.remove_soscf()
-    if isinstance(mf, scf.rohf.ROHF) or isinstance(mf, scf.hf_symm.SymAdaptedROHF):
-        if isinstance(mf, dft.roks.ROKS) or isinstance(mf, dft.rks_symm.SymAdaptedROKS):
+    restricted = isinstance(mf, scf.rohf.ROHF)
+    if restricted:
+        if mf.mo_coeff is None:
+            mf.run()
+        if isinstance(mf, scf.hf.KohnShamDFT):
             mf = mf.to_uks()
         else:
             mf = mf.to_uhf()
-    return mf.TDDFT_SF()
+    td = uhf_sf.TDDFT_SF(mf, extype, collinear, collinear_samples)
+    td._ro_reference = restricted
+    return td
 
 
 SFTDA = TDA_SF
 SFTDDFT = TDDFT_SF
 
+scf.uhf.UHF.TDA_SF = scf.uhf.UHF.SFTDA = TDA_SF
+scf.uhf.UHF.TDDFT_SF = scf.uhf.UHF.SFTDDFT = TDDFT_SF
+scf.rohf.ROHF.TDA_SF = scf.rohf.ROHF.SFTDA = TDA_SF
+scf.rohf.ROHF.TDDFT_SF = scf.rohf.ROHF.SFTDDFT = TDDFT_SF
+
 __all__ = [
     "SFTDA",
     "SFTDDFT",

src/nest/sftda/tests/test_sftda.py

@@ -128,9 +128,9 @@ def test_col_cam_tda(self):
         self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6)
 
     def test_hf_tda_roks(self):
-        mf = self.mol.ROKS(xc='HF').run()
+        mf = self.mol.ROKS(xc='HF')
         ref = np.array([-0.2204522712, -0.0023966488])
-        td = sftda.TDA_SF(mf).set(extype=1, collinear_samples=50, nstates=2).run()
+        td = mf.SFTDA().set(extype=1, collinear_samples=50, nstates=2).run()
         self.assertTrue(np.all(td.converged))
         self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6)
         e = diagonalize_tda(mf, extype=1, collinear_samples=50, nstates=2)
@@ -190,6 +190,16 @@ def test_col_cam_tda_roks(self):
         e = diagonalize_tda(mf, extype=1, collinear="col", collinear_samples=50, nstates=2)
         self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6)
 
+    def test_tda_scanner(self):
+        mf = self.mol.UKS(xc='B3LYP').run()
+        for extype in (0, 1):
+            td = mf.SFTDA().set(extype=extype, nstates=3)
+            ref = td.kernel()[0].copy()
+            td_scan = td.as_scanner()
+            td_scan.max_cycle = 1
+            td_scan(self.mol)
+            self.assertAlmostEqual(abs(td_scan.e - ref).max(), 0, delta=1e-6)
+
 if __name__ == "__main__":
     print("Full Tests for spin-flip-TDA with UKS and ROKS references")
     unittest.main()