Add geometry optimization for excited states using TDDFT-ris

examples/36-tddft-ris-grad-opt.py

@@ -0,0 +1,64 @@
+# 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.
+
+'''
+TDDFT-ris excited state gradient and geometry optimization
+'''
+
+import pyscf
+import gpu4pyscf.tdscf.ris as ris
+from gpu4pyscf.dft import rks
+from pyscf.geomopt.geometric_solver import optimize
+
+atom = """
+O       0.0000000000     0.0000000000     0.0000000000
+H       0.0000000000    -0.7570000000     0.5870000000
+H       0.0000000000     0.7570000000     0.5870000000
+"""
+
+bas0 = "ccpvdz"
+
+mol = pyscf.M(
+    atom=atom, basis=bas0, max_memory=32000)
+mf = rks.RKS(mol, xc='b3lyp').to_gpu()
+mf.kernel()
+td_ris = ris.TDDFT(mf=mf, nstates=5, spectra=False, single=False, GS=True)
+td_ris.conv_tol = 1.0E-4
+td_ris.Ktrunc = 0.0
+td_ris.kernel()
+
+"""
+TDDFT-ris excited state geometry optimization
+1st usage
+"""
+mol_gpu = optimize(td_ris)
+mff = rks.RKS(mol_gpu, xc='b3lyp').to_gpu()
+mff.kernel()
+tdf_ris = ris.TDDFT(mf=mff, nstates=5, spectra=False, single=False, GS=True)
+tdf_ris.conv_tol = 1.0E-4
+tdf_ris.Ktrunc = 0.0
+output = tdf_ris.kernel()
+
+"""
+TDDFT-ris excited state geometry optimization
+2nd usage
+"""
+excited_grad = td_ris.nuc_grad_method().as_scanner(state=1)
+mol_gpu = excited_grad.optimizer().kernel()
+
+"""
+TDDFT-ris excited state gradient
+"""
+excited_gradf_ris = tdf_ris.nuc_grad_method()
+excited_gradf_ris.kernel()

gpu4pyscf/grad/tdrks_ris.py

@@ -16,16 +16,14 @@
 from functools import reduce
 import cupy as cp
 import numpy as np
-from pyscf import lib
+from pyscf import lib, gto
 from gpu4pyscf.lib import logger
-from gpu4pyscf.lib.cupy_helper import contract, add_sparse, tag_array
+from gpu4pyscf.lib.cupy_helper import contract, tag_array
 from gpu4pyscf.df import int3c2e
 from gpu4pyscf.df.grad import tdrhf as tdrhf_df
-from gpu4pyscf.dft import numint,rks
-from pyscf.dft.numint import NumInt as numint_cpu
+from gpu4pyscf.dft import rks
 from gpu4pyscf.scf import cphf
 from gpu4pyscf.grad import rhf as rhf_grad
-from gpu4pyscf.grad import rks as rks_grad
 from gpu4pyscf.grad import tdrhf
 from gpu4pyscf.grad import tdrks
 from gpu4pyscf import tdscf
@@ -305,6 +303,7 @@ def fvind(x):
 def get_extra_force(atom_id, envs):
     return envs['dvhf'].aux[atom_id]
 
+
 def get_veff_ris(mf_J, mf_K, mol=None, dm=None, j_factor=1.0, k_factor=1.0, omega=0.0, hermi=0, verbose=None):
 
     if omega != 0.0:
@@ -341,10 +340,10 @@ def kernel(self, xy=None, state=None, singlet=None, atmlst=None):
                     "state=0 found in the input. Gradients of ground state is computed.",
                 )
                 return self.base._scf.nuc_grad_method().kernel(atmlst=atmlst)
-            if self.base.xy is not None:
+            if self.base.xy[1] is not None:
                 xy = (self.base.xy[0][state-1]*np.sqrt(0.5), self.base.xy[1][state-1]*np.sqrt(0.5))
             else:
-                xy = (self.base.X[state-1]*np.sqrt(0.5), self.base.X[state-1]*0.0)
+                xy = (self.base.xy[0][state-1]*np.sqrt(0.5), self.base.xy[0][state-1]*0.0)
 
         if singlet is None:
             singlet = self.base.singlet

gpu4pyscf/grad/tests/test_tddft_ris_grad.py

@@ -207,7 +207,7 @@ def test_grad_b3lyp_tda_singlet_ref(self):
         mf = dft.RKS(mol, xc='b3lyp').to_gpu()
         mf.kernel()
 
-        td = ris.TDDFT(mf=mf, nstates=5, spectra=True, single=False)
+        td = ris.TDDFT(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
         td.conv_tol = 1.0E-4
         td.Ktrunc = 0.0
         td.kernel()

gpu4pyscf/grad/tests/test_tddft_ris_opt.py

@@ -0,0 +1,93 @@
+# 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.
+
+import pyscf
+import numpy as np
+import unittest
+import pytest
+from pyscf import dft
+from pyscf.geomopt.geometric_solver import optimize
+import gpu4pyscf.tdscf.ris as ris
+
+atom = """
+H     1.2953527433   -0.4895463266    0.8457608681
+C     0.6689912970   -0.0128659340    0.0499408027
+H     1.3504336752    0.5361460613   -0.6478375784
+C    -0.6690192526   -0.0870427249   -0.0501820705
+H    -1.4008634673    0.6483035475    0.3700152345
+H    -1.2449949956   -0.8949946232   -0.5680972562
+"""
+
+bas0 = "def2tzvp"
+
+def setUpModule():
+    global mol
+    mol = pyscf.M(
+        atom=atom, basis=bas0, max_memory=32000, output="/dev/null", verbose=1)
+
+
+def tearDownModule():
+    global mol
+    mol.stdout.close()
+    del mol
+
+class KnownValues(unittest.TestCase):
+    def test_opt_rks_tda_1(self):
+        mf = dft.RKS(mol, xc='pbe0').to_gpu()
+        mf.kernel()
+        assert mf.converged
+        td_ris = ris.TDA(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
+        td_ris.conv_tol = 1.0E-5
+        td_ris.Ktrunc = 0.0
+        td_ris.kernel()
+        mol_gpu = optimize(td_ris)
+
+        mff = dft.RKS(mol_gpu, xc='pbe0').to_gpu()
+        mff.kernel()
+        assert mff.converged
+        tdf_ris = ris.TDA(mf=mff, nstates=5, spectra=False, single=False, gram_schmidt=True)
+        tdf_ris.conv_tol = 1.0E-5
+        tdf_ris.Ktrunc = 0.0
+        tdf_ris.kernel()
+        excited_gradf_ris = tdf_ris.nuc_grad_method()
+        excited_gradf_ris.kernel()
+        assert np.linalg.norm(excited_gradf_ris.de) < 3.0e-4
+
+    def test_opt_rks_tda_2(self):
+        mf = dft.RKS(mol, xc='pbe0').to_gpu()
+        mf.kernel()
+        assert mf.converged
+        td_ris = ris.TDA(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
+        td_ris.conv_tol = 1.0E-5
+        td_ris.Ktrunc = 0.0
+        td_ris.kernel()
+
+        excited_grad = td_ris.nuc_grad_method().as_scanner(state=1)
+        mol_gpu = excited_grad.optimizer().kernel()
+
+        mff = dft.RKS(mol_gpu, xc='pbe0').to_gpu()
+        mff.kernel()
+        assert mff.converged
+        tdf_ris = ris.TDA(mf=mff, nstates=5, spectra=False, single=False, gram_schmidt=True)
+        tdf_ris.conv_tol = 1.0E-5
+        tdf_ris.Ktrunc = 0.0
+        tdf_ris.kernel()
+        excited_gradf_ris = tdf_ris.nuc_grad_method()
+        excited_gradf_ris.kernel()
+        assert np.linalg.norm(excited_gradf_ris.de) < 3.0e-4
+
+
+if __name__ == "__main__":
+    print("Full Tests for geomtry optimization for excited states using TDDFT-ris.")
+    unittest.main()

gpu4pyscf/tdscf/_krylov_tools.py

@@ -478,6 +478,7 @@ def krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
     if ii == max_iter - 1 and max_norm >= conv_tol:
         log.warn(f'=== Warning: {problem_type.capitalize()} solver not converged below {conv_tol:.2e} ===')
         log.warn(f'Current residual norms: {r_norms.tolist()}')
+    converged = r_norms < conv_tol
     log.info(f'Finished in {ii+1} steps')
     log.info(f'Maximum residual norm = {max_norm:.2e}')
     log.info(f'Final subspace size = {sub_A.shape[0]}')
@@ -497,9 +498,9 @@ def krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
     log.info(f'========== {problem_type.capitalize()} Solver Done ==========')
 
     if problem_type == 'eigenvalue':
-        return omega, full_X
+        return converged, omega, full_X
     elif problem_type in ['linear', 'shifted_linear']:
-        return full_X
+        return converged, full_X
 
 def nested_krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
         rhs=None, omega_shift=None, n_states=20, conv_tol=1e-5, 
@@ -626,15 +627,15 @@ def matrix_vector_product(x):
 
     hdiag = cp.diag(A)
 
-    eigenvalues, eigenvecters = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
+    _, eigenvalues, eigenvecters = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
                             problem_type='eigenvalue', n_states=5,
                             conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
 
-    solution_vectors = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
+    _, solution_vectors = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
                             problem_type='linear', rhs=rhs,
                             conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
 
-    solution_vectors_shifted = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
+    _, solution_vectors_shifted = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
                             problem_type='shifted_linear', rhs=rhs, omega_shift=omega_shift,
                             conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
 

gpu4pyscf/tdscf/_lr_eig.py

@@ -1274,6 +1274,7 @@ def Davidson(matrix_vector_product,
         residual = AV - omega.reshape(-1, 1) * full_X
 
         r_norms = cp.linalg.norm(residual, axis=1)
+        conv = r_norms[:N_states] <= conv_tol
         max_norm = cp.max(r_norms)
         log.info(f'iter: {ii+1:<3d}   max|R|: {max_norm:<12.2e}  subspace: {sub_A.shape[0]:<8d}')
         if max_norm < conv_tol or ii == (max_iter-1):
@@ -1312,7 +1313,7 @@ def Davidson(matrix_vector_product,
 
 
     log.info('========== Davidson Diagonalization Done ==========')
-    return omega, full_X
+    return conv, omega, full_X
 
 # TODO: merge with real_eig, write a Class of krylov method for Casida problem, allowing ris initial guess/preconditioner
 def Davidson_Casida(matrix_vector_product,
@@ -1509,7 +1510,7 @@ def Davidson_Casida(matrix_vector_product,
         r_norms = cp.linalg.norm(residual, axis=1)
 
         max_norm = cp.max(r_norms)
-
+        conv = r_norms[:N_states] <= conv_tol
         log.info(f'iter: {ii+1:<3d}, max|R|: {max_norm:<10.2e} subspace_size = {sub_A.shape[0]}')
 
         if max_norm < conv_tol or ii == (max_iter -1):
@@ -1560,5 +1561,5 @@ def Davidson_Casida(matrix_vector_product,
 
     log.info('======= TDDFT Eigen Solver Done =======' )
 
-    return omega, X_full, Y_full
+    return conv, omega, X_full, Y_full