Skip to content

Commit 5965dc4

Browse files
authored
Add geometry optimization for excited states using TDDFT-ris (#443)
* add ris optimization * change some comments * fix some typos * change the initial geometry * fix a bug in generating auxmol * change the codes as review comments : * fix a bug * merge the master ris * fix some bugs * fix a typo
1 parent bb5cd0f commit 5965dc4

8 files changed

Lines changed: 265 additions & 49 deletions

File tree

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# Copyright 2021-2025 The PySCF Developers. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
'''
16+
TDDFT-ris excited state gradient and geometry optimization
17+
'''
18+
19+
import pyscf
20+
import gpu4pyscf.tdscf.ris as ris
21+
from gpu4pyscf.dft import rks
22+
from pyscf.geomopt.geometric_solver import optimize
23+
24+
atom = """
25+
O 0.0000000000 0.0000000000 0.0000000000
26+
H 0.0000000000 -0.7570000000 0.5870000000
27+
H 0.0000000000 0.7570000000 0.5870000000
28+
"""
29+
30+
bas0 = "ccpvdz"
31+
32+
mol = pyscf.M(
33+
atom=atom, basis=bas0, max_memory=32000)
34+
mf = rks.RKS(mol, xc='b3lyp').to_gpu()
35+
mf.kernel()
36+
td_ris = ris.TDDFT(mf=mf, nstates=5, spectra=False, single=False, GS=True)
37+
td_ris.conv_tol = 1.0E-4
38+
td_ris.Ktrunc = 0.0
39+
td_ris.kernel()
40+
41+
"""
42+
TDDFT-ris excited state geometry optimization
43+
1st usage
44+
"""
45+
mol_gpu = optimize(td_ris)
46+
mff = rks.RKS(mol_gpu, xc='b3lyp').to_gpu()
47+
mff.kernel()
48+
tdf_ris = ris.TDDFT(mf=mff, nstates=5, spectra=False, single=False, GS=True)
49+
tdf_ris.conv_tol = 1.0E-4
50+
tdf_ris.Ktrunc = 0.0
51+
output = tdf_ris.kernel()
52+
53+
"""
54+
TDDFT-ris excited state geometry optimization
55+
2nd usage
56+
"""
57+
excited_grad = td_ris.nuc_grad_method().as_scanner(state=1)
58+
mol_gpu = excited_grad.optimizer().kernel()
59+
60+
"""
61+
TDDFT-ris excited state gradient
62+
"""
63+
excited_gradf_ris = tdf_ris.nuc_grad_method()
64+
excited_gradf_ris.kernel()

gpu4pyscf/grad/tdrks_ris.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,14 @@
1616
from functools import reduce
1717
import cupy as cp
1818
import numpy as np
19-
from pyscf import lib
19+
from pyscf import lib, gto
2020
from gpu4pyscf.lib import logger
21-
from gpu4pyscf.lib.cupy_helper import contract, add_sparse, tag_array
21+
from gpu4pyscf.lib.cupy_helper import contract, tag_array
2222
from gpu4pyscf.df import int3c2e
2323
from gpu4pyscf.df.grad import tdrhf as tdrhf_df
24-
from gpu4pyscf.dft import numint,rks
25-
from pyscf.dft.numint import NumInt as numint_cpu
24+
from gpu4pyscf.dft import rks
2625
from gpu4pyscf.scf import cphf
2726
from gpu4pyscf.grad import rhf as rhf_grad
28-
from gpu4pyscf.grad import rks as rks_grad
2927
from gpu4pyscf.grad import tdrhf
3028
from gpu4pyscf.grad import tdrks
3129
from gpu4pyscf import tdscf
@@ -305,6 +303,7 @@ def fvind(x):
305303
def get_extra_force(atom_id, envs):
306304
return envs['dvhf'].aux[atom_id]
307305

306+
308307
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):
309308

310309
if omega != 0.0:
@@ -341,10 +340,10 @@ def kernel(self, xy=None, state=None, singlet=None, atmlst=None):
341340
"state=0 found in the input. Gradients of ground state is computed.",
342341
)
343342
return self.base._scf.nuc_grad_method().kernel(atmlst=atmlst)
344-
if self.base.xy is not None:
343+
if self.base.xy[1] is not None:
345344
xy = (self.base.xy[0][state-1]*np.sqrt(0.5), self.base.xy[1][state-1]*np.sqrt(0.5))
346345
else:
347-
xy = (self.base.X[state-1]*np.sqrt(0.5), self.base.X[state-1]*0.0)
346+
xy = (self.base.xy[0][state-1]*np.sqrt(0.5), self.base.xy[0][state-1]*0.0)
348347

349348
if singlet is None:
350349
singlet = self.base.singlet

gpu4pyscf/grad/tests/test_tddft_ris_grad.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ def test_grad_b3lyp_tda_singlet_ref(self):
207207
mf = dft.RKS(mol, xc='b3lyp').to_gpu()
208208
mf.kernel()
209209

210-
td = ris.TDDFT(mf=mf, nstates=5, spectra=True, single=False)
210+
td = ris.TDDFT(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
211211
td.conv_tol = 1.0E-4
212212
td.Ktrunc = 0.0
213213
td.kernel()
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# Copyright 2021-2025 The PySCF Developers. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pyscf
16+
import numpy as np
17+
import unittest
18+
import pytest
19+
from pyscf import dft
20+
from pyscf.geomopt.geometric_solver import optimize
21+
import gpu4pyscf.tdscf.ris as ris
22+
23+
atom = """
24+
H 1.2953527433 -0.4895463266 0.8457608681
25+
C 0.6689912970 -0.0128659340 0.0499408027
26+
H 1.3504336752 0.5361460613 -0.6478375784
27+
C -0.6690192526 -0.0870427249 -0.0501820705
28+
H -1.4008634673 0.6483035475 0.3700152345
29+
H -1.2449949956 -0.8949946232 -0.5680972562
30+
"""
31+
32+
bas0 = "def2tzvp"
33+
34+
def setUpModule():
35+
global mol
36+
mol = pyscf.M(
37+
atom=atom, basis=bas0, max_memory=32000, output="/dev/null", verbose=1)
38+
39+
40+
def tearDownModule():
41+
global mol
42+
mol.stdout.close()
43+
del mol
44+
45+
class KnownValues(unittest.TestCase):
46+
def test_opt_rks_tda_1(self):
47+
mf = dft.RKS(mol, xc='pbe0').to_gpu()
48+
mf.kernel()
49+
assert mf.converged
50+
td_ris = ris.TDA(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
51+
td_ris.conv_tol = 1.0E-5
52+
td_ris.Ktrunc = 0.0
53+
td_ris.kernel()
54+
mol_gpu = optimize(td_ris)
55+
56+
mff = dft.RKS(mol_gpu, xc='pbe0').to_gpu()
57+
mff.kernel()
58+
assert mff.converged
59+
tdf_ris = ris.TDA(mf=mff, nstates=5, spectra=False, single=False, gram_schmidt=True)
60+
tdf_ris.conv_tol = 1.0E-5
61+
tdf_ris.Ktrunc = 0.0
62+
tdf_ris.kernel()
63+
excited_gradf_ris = tdf_ris.nuc_grad_method()
64+
excited_gradf_ris.kernel()
65+
assert np.linalg.norm(excited_gradf_ris.de) < 3.0e-4
66+
67+
def test_opt_rks_tda_2(self):
68+
mf = dft.RKS(mol, xc='pbe0').to_gpu()
69+
mf.kernel()
70+
assert mf.converged
71+
td_ris = ris.TDA(mf=mf, nstates=5, spectra=False, single=False, gram_schmidt=True)
72+
td_ris.conv_tol = 1.0E-5
73+
td_ris.Ktrunc = 0.0
74+
td_ris.kernel()
75+
76+
excited_grad = td_ris.nuc_grad_method().as_scanner(state=1)
77+
mol_gpu = excited_grad.optimizer().kernel()
78+
79+
mff = dft.RKS(mol_gpu, xc='pbe0').to_gpu()
80+
mff.kernel()
81+
assert mff.converged
82+
tdf_ris = ris.TDA(mf=mff, nstates=5, spectra=False, single=False, gram_schmidt=True)
83+
tdf_ris.conv_tol = 1.0E-5
84+
tdf_ris.Ktrunc = 0.0
85+
tdf_ris.kernel()
86+
excited_gradf_ris = tdf_ris.nuc_grad_method()
87+
excited_gradf_ris.kernel()
88+
assert np.linalg.norm(excited_gradf_ris.de) < 3.0e-4
89+
90+
91+
if __name__ == "__main__":
92+
print("Full Tests for geomtry optimization for excited states using TDDFT-ris.")
93+
unittest.main()

gpu4pyscf/tdscf/_krylov_tools.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -478,6 +478,7 @@ def krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
478478
if ii == max_iter - 1 and max_norm >= conv_tol:
479479
log.warn(f'=== Warning: {problem_type.capitalize()} solver not converged below {conv_tol:.2e} ===')
480480
log.warn(f'Current residual norms: {r_norms.tolist()}')
481+
converged = r_norms < conv_tol
481482
log.info(f'Finished in {ii+1} steps')
482483
log.info(f'Maximum residual norm = {max_norm:.2e}')
483484
log.info(f'Final subspace size = {sub_A.shape[0]}')
@@ -497,9 +498,9 @@ def krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
497498
log.info(f'========== {problem_type.capitalize()} Solver Done ==========')
498499

499500
if problem_type == 'eigenvalue':
500-
return omega, full_X
501+
return converged, omega, full_X
501502
elif problem_type in ['linear', 'shifted_linear']:
502-
return full_X
503+
return converged, full_X
503504

504505
def nested_krylov_solver(matrix_vector_product, hdiag, problem_type='eigenvalue',
505506
rhs=None, omega_shift=None, n_states=20, conv_tol=1e-5,
@@ -626,15 +627,15 @@ def matrix_vector_product(x):
626627

627628
hdiag = cp.diag(A)
628629

629-
eigenvalues, eigenvecters = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
630+
_, eigenvalues, eigenvecters = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
630631
problem_type='eigenvalue', n_states=5,
631632
conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
632633

633-
solution_vectors = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
634+
_, solution_vectors = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
634635
problem_type='linear', rhs=rhs,
635636
conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
636637

637-
solution_vectors_shifted = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
638+
_, solution_vectors_shifted = krylov_solver(matrix_vector_product=matrix_vector_product, hdiag=hdiag,
638639
problem_type='shifted_linear', rhs=rhs, omega_shift=omega_shift,
639640
conv_tol=1e-5, max_iter=35,gram_schmidt=True, verbose=5, single=False)
640641

gpu4pyscf/tdscf/_lr_eig.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1274,6 +1274,7 @@ def Davidson(matrix_vector_product,
12741274
residual = AV - omega.reshape(-1, 1) * full_X
12751275

12761276
r_norms = cp.linalg.norm(residual, axis=1)
1277+
conv = r_norms[:N_states] <= conv_tol
12771278
max_norm = cp.max(r_norms)
12781279
log.info(f'iter: {ii+1:<3d} max|R|: {max_norm:<12.2e} subspace: {sub_A.shape[0]:<8d}')
12791280
if max_norm < conv_tol or ii == (max_iter-1):
@@ -1312,7 +1313,7 @@ def Davidson(matrix_vector_product,
13121313

13131314

13141315
log.info('========== Davidson Diagonalization Done ==========')
1315-
return omega, full_X
1316+
return conv, omega, full_X
13161317

13171318
# TODO: merge with real_eig, write a Class of krylov method for Casida problem, allowing ris initial guess/preconditioner
13181319
def Davidson_Casida(matrix_vector_product,
@@ -1509,7 +1510,7 @@ def Davidson_Casida(matrix_vector_product,
15091510
r_norms = cp.linalg.norm(residual, axis=1)
15101511

15111512
max_norm = cp.max(r_norms)
1512-
1513+
conv = r_norms[:N_states] <= conv_tol
15131514
log.info(f'iter: {ii+1:<3d}, max|R|: {max_norm:<10.2e} subspace_size = {sub_A.shape[0]}')
15141515

15151516
if max_norm < conv_tol or ii == (max_iter -1):
@@ -1560,5 +1561,5 @@ def Davidson_Casida(matrix_vector_product,
15601561

15611562
log.info('======= TDDFT Eigen Solver Done =======' )
15621563

1563-
return omega, X_full, Y_full
1564+
return conv, omega, X_full, Y_full
15641565

0 commit comments

Comments
 (0)