Skip to content

Commit cd09930

Browse files
sunqmsunqm
authored
Improve PCM-TDDFT modules for PySCF integration (#558)
* Update TDDFT-PCM interface * Add the "from_cpu" interface * Add interface to pyscf * clean up attributes
1 parent 2e13995 commit cd09930

File tree

9 files changed

+139
-68
lines changed

9 files changed

+139
-68
lines changed

‎gpu4pyscf/grad/tdrhf.py‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
from gpu4pyscf import tdscf
2727

2828

29-
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
29+
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO,
30+
with_solvent=False):
3031
"""
3132
Electronic part of TDA, TDHF nuclear gradients
3233
@@ -36,6 +37,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
3637
x_y : a two-element list of numpy arrays
3738
TDDFT X and Y amplitudes. If Y is set to 0, this function computes
3839
TDA energy gradients.
40+
41+
Kwargs:
42+
with_solvent :
43+
Include the response of solvent in the gradients of the electronic
44+
energy.
3945
"""
4046
if singlet is None:
4147
singlet = True
@@ -62,7 +68,8 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
6268
dmxmy = reduce(cp.dot, (orbv, xmy, orbo.T)) # (X-Y) in ao basis
6369
dmzoo = reduce(cp.dot, (orbo, doo, orbo.T)) # T_{ij}*2 in ao basis
6470
dmzoo += reduce(cp.dot, (orbv, dvv, orbv.T)) # T_{ij}*2 + T_{ab}*2 in ao basis
65-
td_grad.dmxpy = dmxpy
71+
if with_solvent:
72+
td_grad._dmxpy = dmxpy
6673

6774
vj0, vk0 = mf.get_jk(mol, dmzoo, hermi=0)
6875
vj1, vk1 = mf.get_jk(mol, dmxpy + dmxpy.T, hermi=0)
@@ -82,13 +89,15 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
8289
vj = cp.stack((vj0, vj1, vj2))
8390
vk = cp.stack((vk0, vk1, vk2))
8491
veff0doo = vj[0] * 2 - vk[0] # 2 for alpha and beta
85-
veff0doo += td_grad.solvent_response(dmzoo)
92+
if with_solvent:
93+
veff0doo += td_grad.solvent_response(dmzoo)
8694
wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2
8795
if singlet:
8896
veff = vj[1] * 2 - vk[1]
8997
else:
9098
veff = -vk[1]
91-
veff += td_grad.solvent_response(dmxpy + dmxpy.T)
99+
if with_solvent:
100+
veff += td_grad.solvent_response(dmxpy + dmxpy.T)
92101
veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff))
93102
wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2 # 2 for dm + dm.T
94103
wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2
@@ -154,7 +163,8 @@ def fvind(x): # For singlet, closed shell ground state
154163
s1 = mf_grad.get_ovlp(mol)
155164

156165
dmz1doo = z1ao + dmzoo # P
157-
td_grad.dmz1doo = dmz1doo
166+
if with_solvent:
167+
td_grad._dmz1doo = dmz1doo
158168
oo0 = reduce(cp.dot, (orbo, orbo.T)) # D
159169

160170
if atmlst is None:
@@ -288,12 +298,9 @@ class Gradients(rhf_grad.GradientsBase):
288298
"cphf_conv_tol",
289299
"mol",
290300
"base",
291-
"chkfile",
292301
"state",
293302
"atmlst",
294303
"de",
295-
"dmz1doo",
296-
"dmxpy"
297304
}
298305

299306
def __init__(self, td):
@@ -302,12 +309,9 @@ def __init__(self, td):
302309
self.stdout = td.stdout
303310
self.mol = td.mol
304311
self.base = td
305-
self.chkfile = td.chkfile
306312
self.state = 1 # of which the gradients to be computed.
307313
self.atmlst = None
308314
self.de = None
309-
self.dmz1doo = None
310-
self.dmxpy = None
311315

312316
def dump_flags(self, verbose=None):
313317
log = logger.new_logger(self, verbose)
@@ -324,9 +328,7 @@ def dump_flags(self, verbose=None):
324328
log.info("\n")
325329
return self
326330

327-
@lib.with_doc(grad_elec.__doc__)
328-
def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO):
329-
return grad_elec(self, xy, singlet, atmlst, verbose)
331+
grad_elec = grad_elec
330332

331333
def kernel(self, xy=None, state=None, singlet=None, atmlst=None):
332334
"""
@@ -413,5 +415,14 @@ def solvent_response(self, dm):
413415

414416
to_gpu = lib.to_gpu
415417

418+
@classmethod
419+
def from_cpu(cls, method):
420+
td = method.base.to_gpu()
421+
out = cls(td)
422+
out.cphf_max_cycle = method.cphf_max_cycle
423+
out.cphf_conv_tol = method.cphf_conv_tol
424+
out.state = method.state
425+
out.de = method.de
426+
return out
416427

417428
Grad = Gradients

‎gpu4pyscf/grad/tdrks.py‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@
3232
#
3333
# Given Y = 0, TDDFT gradients (XAX+XBY+YBX+YAY)^1 turn to TDA gradients (XAX)^1
3434
#
35-
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
35+
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO,
36+
with_solvent=False):
3637
"""
3738
Electronic part of TDA, TDDFT nuclear gradients
3839
@@ -42,6 +43,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
4243
x_y : a two-element list of cp arrays
4344
TDDFT X and Y amplitudes. If Y is set to 0, this function computes
4445
TDA energy gradients.
46+
47+
Kwargs:
48+
with_solvent :
49+
Include the response of solvent in the gradients of the electronic
50+
energy.
4551
"""
4652
if singlet is None:
4753
singlet = True
@@ -69,7 +75,8 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
6975
dmxmy = reduce(cp.dot, (orbv, xmy, orbo.T)) # (X-Y) in ao basis
7076
dmzoo = reduce(cp.dot, (orbo, doo, orbo.T)) # T_{ij}*2 in ao basis
7177
dmzoo += reduce(cp.dot, (orbv, dvv, orbv.T)) # T_{ij}*2 + T_{ab}*2 in ao basis
72-
td_grad.dmxpy = dmxpy
78+
if with_solvent:
79+
td_grad._dmxpy = dmxpy
7380

7481
ni = mf._numint
7582
ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True)
@@ -107,13 +114,15 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
107114
vk2 = cp.asarray(vk2)
108115
vk += cp.stack((vk0, vk1, vk2)) * (alpha - hyb)
109116
veff0doo = vj[0] * 2 - vk[0] + f1oo[0] + k1ao[0] * 2
110-
veff0doo += td_grad.solvent_response(dmzoo)
117+
if with_solvent:
118+
veff0doo += td_grad.solvent_response(dmzoo)
111119
wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2
112120
if singlet:
113121
veff = vj[1] * 2 - vk[1] + f1vo[0] * 2
114122
else:
115123
veff = f1vo[0] - vk[1]
116-
veff += td_grad.solvent_response(dmxpy + dmxpy.T)
124+
if with_solvent:
125+
veff += td_grad.solvent_response(dmxpy + dmxpy.T)
117126
veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff))
118127
wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2
119128
wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2
@@ -187,7 +196,8 @@ def fvind(x):
187196
s1 = mf_grad.get_ovlp(mol)
188197

189198
dmz1doo = z1ao + dmzoo
190-
td_grad.dmz1doo = dmz1doo
199+
if with_solvent:
200+
td_grad._dmz1doo = dmz1doo
191201
oo0 = reduce(cp.dot, (orbo, orbo.T))
192202

193203
if atmlst is None:
@@ -524,9 +534,7 @@ def _mgga_eval_mat_(mol, vmat, ao, wv, mask, shls_slice, ao_loc):
524534

525535

526536
class Gradients(tdrhf.Gradients):
527-
@lib.with_doc(grad_elec.__doc__)
528-
def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.info):
529-
return grad_elec(self, xy, singlet, atmlst, self.verbose)
537+
grad_elec = grad_elec
530538

531539

532540
Grad = Gradients

‎gpu4pyscf/grad/tduhf.py‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
from gpu4pyscf import tdscf
2727

2828

29-
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
29+
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO,
30+
with_solvent=False):
3031
"""
3132
Electronic part of TDA, TDHF nuclear gradients
3233
@@ -36,6 +37,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
3637
x_y : a two-element list of numpy arrays
3738
TDDFT X and Y amplitudes. If Y is set to 0, this function computes
3839
TDA energy gradients.
40+
41+
Kwargs:
42+
with_solvent :
43+
Include the response of solvent in the gradients of the electronic
44+
energy.
3945
"""
4046
if singlet is not True and singlet is not None:
4147
raise NotImplementedError("Only for spin-conserving TDHF")
@@ -85,7 +91,9 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
8591
dmzoob = reduce(cp.dot, (orbob, doob, orbob.T))
8692
dmzooa += reduce(cp.dot, (orbva, dvva, orbva.T))
8793
dmzoob += reduce(cp.dot, (orbvb, dvvb, orbvb.T))
88-
td_grad.dmxpy = (dmxpya + dmxpyb)*0.5
94+
dmxpy = (dmxpya + dmxpyb)*0.5
95+
if with_solvent:
96+
td_grad._dmxpy = dmxpy
8997

9098
vj0, vk0 = mf.get_jk(mol, cp.stack((dmzooa, dmzoob)), hermi=0)
9199
vj1, vk1 = mf.get_jk(mol, cp.stack((dmxpya + dmxpya.T, dmxpyb + dmxpyb.T)), hermi=0)
@@ -108,11 +116,13 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
108116
vk = vk.reshape(2, 3, nao, nao)
109117

110118
veff0doo = vj[0, 0] + vj[1, 0] - vk[:, 0]
111-
veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5)
119+
if with_solvent:
120+
veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5)
112121
wvoa = reduce(cp.dot, (orbva.T, veff0doo[0], orboa)) * 2
113122
wvob = reduce(cp.dot, (orbvb.T, veff0doo[1], orbob)) * 2
114123
veff = vj[0, 1] + vj[1, 1] - vk[:, 1]
115-
veff += td_grad.solvent_response((td_grad.dmxpy + td_grad.dmxpy.T))
124+
if with_solvent:
125+
veff += td_grad.solvent_response((dmxpy + dmxpy.T))
116126
veff0mopa = reduce(cp.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
117127
veff0mopb = reduce(cp.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
118128
wvoa -= contract("ki,ai->ak", veff0mopa[:nocca, :nocca], xpya) * 2
@@ -195,7 +205,8 @@ def fvind(x):
195205

196206
dmz1dooa = z1ao[0] + dmzooa
197207
dmz1doob = z1ao[1] + dmzoob
198-
td_grad.dmz1doo = (dmz1dooa + dmz1doob)*0.5
208+
if with_solvent:
209+
td_grad._dmz1doo = (dmz1dooa + dmz1doob)*0.5
199210
oo0a = reduce(cp.dot, (orboa, orboa.T))
200211
oo0b = reduce(cp.dot, (orbob, orbob.T))
201212

@@ -258,10 +269,7 @@ class Gradients(tdrhf.Gradients):
258269
to_cpu = utils.to_cpu
259270
to_gpu = utils.to_gpu
260271
device = utils.device
261-
262-
@lib.with_doc(grad_elec.__doc__)
263-
def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.info):
264-
return grad_elec(self, xy, singlet, atmlst, self.verbose)
272+
grad_elec = grad_elec
265273

266274

267275
Grad = Gradients

‎gpu4pyscf/grad/tduks.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@
3131
#
3232
# Given Y = 0, TDHF gradients (XAX+XBY+YBX+YAY)^1 turn to TDA gradients (XAX)^1
3333
#
34-
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
34+
def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO,
35+
with_solvent=False):
3536
"""
3637
Electronic part of TDA, TDDFT nuclear gradients
3738
@@ -41,6 +42,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
4142
x_y : a two-element list of numpy arrays
4243
TDDFT X and Y amplitudes. If Y is set to 0, this function computes
4344
TDA energy gradients.
45+
46+
Kwargs:
47+
with_solvent :
48+
Include the response of solvent in the gradients of the electronic
49+
energy.
4450
"""
4551
if singlet is not True and singlet is not None:
4652
raise NotImplementedError("Only for spin-conserving TDDFT")
@@ -91,7 +97,9 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
9197
dmzoob = reduce(cp.dot, (orbob, doob, orbob.T))
9298
dmzooa += reduce(cp.dot, (orbva, dvva, orbva.T))
9399
dmzoob += reduce(cp.dot, (orbvb, dvvb, orbvb.T))
94-
td_grad.dmxpy = (dmxpya + dmxpyb)*0.5
100+
dmxpy = (dmxpya + dmxpyb)*0.5
101+
if with_solvent:
102+
td_grad._dmxpy = dmxpy
95103

96104
ni = mf._numint
97105
ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True)
@@ -137,11 +145,13 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO):
137145
vk = vk.reshape(2, 3, nao, nao)
138146

139147
veff0doo = vj[0, 0] + vj[1, 0] - vk[:, 0] + f1oo[:, 0] + k1ao[:, 0] * 2
140-
veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5)
148+
if with_solvent:
149+
veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5)
141150
wvoa = reduce(cp.dot, (orbva.T, veff0doo[0], orboa)) * 2
142151
wvob = reduce(cp.dot, (orbvb.T, veff0doo[1], orbob)) * 2
143152
veff = vj[0, 1] + vj[1, 1] - vk[:, 1] + f1vo[:, 0] * 2
144-
veff += td_grad.solvent_response((td_grad.dmxpy + td_grad.dmxpy.T))
153+
if with_solvent:
154+
veff += td_grad.solvent_response((dmxpy + dmxpy.T))
145155
veff0mopa = reduce(cp.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
146156
veff0mopb = reduce(cp.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
147157
wvoa -= contract("ki,ai->ak", veff0mopa[:nocca, :nocca], xpya) * 2
@@ -245,7 +255,8 @@ def fvind(x):
245255

246256
dmz1dooa = z1ao[0] + dmzooa
247257
dmz1doob = z1ao[1] + dmzoob
248-
td_grad.dmz1doo = (dmz1dooa + dmz1doob)*0.5
258+
if with_solvent:
259+
td_grad._dmz1doo = (dmz1dooa + dmz1doob)*0.5
249260
oo0a = reduce(cp.dot, (orboa, orboa.T))
250261
oo0b = reduce(cp.dot, (orbob, orbob.T))
251262

@@ -505,9 +516,7 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True, with_k
505516

506517

507518
class Gradients(tdrhf.Gradients):
508-
@lib.with_doc(grad_elec.__doc__)
509-
def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.info):
510-
return grad_elec(self, xy, singlet, atmlst, self.verbose)
519+
grad_elec = grad_elec
511520

512521

513522
Grad = Gradients

0 commit comments

Comments
 (0)