1#!/usr/bin/env python 2# Copyright 2014-2019 The PySCF Developers. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# 16# Author: Qiming Sun <osirpt.sun@gmail.com> 17# 18# Ref: 19# J. Chem. Phys. 117, 7433 20# 21 22 23from functools import reduce 24import numpy 25from pyscf import lib 26from pyscf.lib import logger 27from pyscf.dft import rks 28from pyscf.dft import numint 29from pyscf.scf import cphf 30from pyscf.grad import rks as rks_grad 31from pyscf.grad import tdrhf 32 33 34# 35# Given Y = 0, TDDFT gradients (XAX+XBY+YBX+YAY)^1 turn to TDA gradients (XAX)^1 36# 37def grad_elec(td_grad, x_y, singlet=True, atmlst=None, 38 max_memory=2000, verbose=logger.INFO): 39 ''' 40 Electronic part of TDA, TDDFT nuclear gradients 41 42 Args: 43 td_grad : grad.tdrhf.Gradients or grad.tdrks.Gradients object. 44 45 x_y : a two-element list of numpy arrays 46 TDDFT X and Y amplitudes. If Y is set to 0, this function computes 47 TDA energy gradients. 48 ''' 49 log = logger.new_logger(td_grad, verbose) 50 time0 = logger.process_clock(), logger.perf_counter() 51 52 mol = td_grad.mol 53 mf = td_grad.base._scf 54 mo_coeff = mf.mo_coeff 55 mo_energy = mf.mo_energy 56 mo_occ = mf.mo_occ 57 nao, nmo = mo_coeff.shape 58 nocc = (mo_occ>0).sum() 59 nvir = nmo - nocc 60 x, y = x_y 61 xpy = (x+y).reshape(nocc,nvir).T 62 xmy = (x-y).reshape(nocc,nvir).T 63 orbv = mo_coeff[:,nocc:] 64 orbo = mo_coeff[:,:nocc] 65 66 dvv = numpy.einsum('ai,bi->ab', xpy, xpy) + numpy.einsum('ai,bi->ab', xmy, xmy) 67 doo =-numpy.einsum('ai,aj->ij', xpy, xpy) - numpy.einsum('ai,aj->ij', xmy, xmy) 68 dmxpy = reduce(numpy.dot, (orbv, xpy, orbo.T)) 69 dmxmy = reduce(numpy.dot, (orbv, xmy, orbo.T)) 70 dmzoo = reduce(numpy.dot, (orbo, doo, orbo.T)) 71 dmzoo+= reduce(numpy.dot, (orbv, dvv, orbv.T)) 72 73 mem_now = lib.current_memory()[0] 74 max_memory = max(2000, td_grad.max_memory*.9-mem_now) 75 76 ni = mf._numint 77 ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True) 78 omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) 79 # dm0 = mf.make_rdm1(mo_coeff, mo_occ), but it is not used when computing 80 # fxc since rho0 is passed to fxc function. 81 rho0, vxc, fxc = ni.cache_xc_kernel(mf.mol, mf.grids, mf.xc, 82 [mo_coeff]*2, [mo_occ*.5]*2, spin=1) 83 f1vo, f1oo, vxc1, k1ao = \ 84 _contract_xc_kernel(td_grad, mf.xc, dmxpy, 85 dmzoo, True, True, singlet, max_memory) 86 87 if abs(hyb) > 1e-10: 88 dm = (dmzoo, dmxpy+dmxpy.T, dmxmy-dmxmy.T) 89 vj, vk = mf.get_jk(mol, dm, hermi=0) 90 vk *= hyb 91 if abs(omega) > 1e-10: 92 vk += mf.get_k(mol, dm, hermi=0, omega=omega) * (alpha-hyb) 93 veff0doo = vj[0] * 2 - vk[0] + f1oo[0] + k1ao[0] * 2 94 wvo = reduce(numpy.dot, (orbv.T, veff0doo, orbo)) * 2 95 if singlet: 96 veff = vj[1] * 2 - vk[1] + f1vo[0] * 2 97 else: 98 veff = -vk[1] + f1vo[0] * 2 99 veff0mop = reduce(numpy.dot, (mo_coeff.T, veff, mo_coeff)) 100 wvo -= numpy.einsum('ki,ai->ak', veff0mop[:nocc,:nocc], xpy) * 2 101 wvo += numpy.einsum('ac,ai->ci', veff0mop[nocc:,nocc:], xpy) * 2 102 veff = -vk[2] 103 veff0mom = reduce(numpy.dot, (mo_coeff.T, veff, mo_coeff)) 104 wvo -= numpy.einsum('ki,ai->ak', veff0mom[:nocc,:nocc], xmy) * 2 105 wvo += numpy.einsum('ac,ai->ci', veff0mom[nocc:,nocc:], xmy) * 2 106 else: 107 vj = mf.get_j(mol, (dmzoo, dmxpy+dmxpy.T), hermi=1) 108 veff0doo = vj[0] * 2 + f1oo[0] + k1ao[0] * 2 109 wvo = reduce(numpy.dot, (orbv.T, veff0doo, orbo)) * 2 110 if singlet: 111 veff = vj[1] * 2 + f1vo[0] * 2 112 else: 113 veff = f1vo[0] * 2 114 veff0mop = reduce(numpy.dot, (mo_coeff.T, veff, mo_coeff)) 115 wvo -= numpy.einsum('ki,ai->ak', veff0mop[:nocc,:nocc], xpy) * 2 116 wvo += numpy.einsum('ac,ai->ci', veff0mop[nocc:,nocc:], xpy) * 2 117 veff0mom = numpy.zeros((nmo,nmo)) 118 119 # set singlet=None, generate function for CPHF type response kernel 120 vresp = mf.gen_response(singlet=None, hermi=1) 121 def fvind(x): 122 dm = reduce(numpy.dot, (orbv, x.reshape(nvir,nocc)*2, orbo.T)) 123 v1ao = vresp(dm+dm.T) 124 return reduce(numpy.dot, (orbv.T, v1ao, orbo)).ravel() 125 z1 = cphf.solve(fvind, mo_energy, mo_occ, wvo, 126 max_cycle=td_grad.cphf_max_cycle, 127 tol=td_grad.cphf_conv_tol)[0] 128 z1 = z1.reshape(nvir,nocc) 129 time1 = log.timer('Z-vector using CPHF solver', *time0) 130 131 z1ao = reduce(numpy.dot, (orbv, z1, orbo.T)) 132 veff = vresp(z1ao+z1ao.T) 133 134 im0 = numpy.zeros((nmo,nmo)) 135 im0[:nocc,:nocc] = reduce(numpy.dot, (orbo.T, veff0doo+veff, orbo)) 136 im0[:nocc,:nocc]+= numpy.einsum('ak,ai->ki', veff0mop[nocc:,:nocc], xpy) 137 im0[:nocc,:nocc]+= numpy.einsum('ak,ai->ki', veff0mom[nocc:,:nocc], xmy) 138 im0[nocc:,nocc:] = numpy.einsum('ci,ai->ac', veff0mop[nocc:,:nocc], xpy) 139 im0[nocc:,nocc:]+= numpy.einsum('ci,ai->ac', veff0mom[nocc:,:nocc], xmy) 140 im0[nocc:,:nocc] = numpy.einsum('ki,ai->ak', veff0mop[:nocc,:nocc], xpy)*2 141 im0[nocc:,:nocc]+= numpy.einsum('ki,ai->ak', veff0mom[:nocc,:nocc], xmy)*2 142 143 zeta = lib.direct_sum('i+j->ij', mo_energy, mo_energy) * .5 144 zeta[nocc:,:nocc] = mo_energy[:nocc] 145 zeta[:nocc,nocc:] = mo_energy[nocc:] 146 dm1 = numpy.zeros((nmo,nmo)) 147 dm1[:nocc,:nocc] = doo 148 dm1[nocc:,nocc:] = dvv 149 dm1[nocc:,:nocc] = z1 150 dm1[:nocc,:nocc] += numpy.eye(nocc)*2 # for ground state 151 im0 = reduce(numpy.dot, (mo_coeff, im0+zeta*dm1, mo_coeff.T)) 152 153 # Initialize hcore_deriv with the underlying SCF object because some 154 # extensions (e.g. QM/MM, solvent) modifies the SCF object only. 155 mf_grad = td_grad.base._scf.nuc_grad_method() 156 hcore_deriv = mf_grad.hcore_generator(mol) 157 s1 = mf_grad.get_ovlp(mol) 158 159 dmz1doo = z1ao + dmzoo 160 oo0 = reduce(numpy.dot, (orbo, orbo.T)) 161 if abs(hyb) > 1e-10: 162 dm = (oo0, dmz1doo+dmz1doo.T, dmxpy+dmxpy.T, dmxmy-dmxmy.T) 163 vj, vk = td_grad.get_jk(mol, dm) 164 vk *= hyb 165 if abs(omega) > 1e-10: 166 with mol.with_range_coulomb(omega): 167 vk += td_grad.get_k(mol, dm) * (alpha-hyb) 168 vj = vj.reshape(-1,3,nao,nao) 169 vk = vk.reshape(-1,3,nao,nao) 170 if singlet: 171 veff1 = vj * 2 - vk 172 else: 173 veff1 = numpy.vstack((vj[:2]*2-vk[:2], -vk[2:])) 174 else: 175 vj = td_grad.get_j(mol, (oo0, dmz1doo+dmz1doo.T, dmxpy+dmxpy.T)) 176 vj = vj.reshape(-1,3,nao,nao) 177 veff1 = numpy.zeros((4,3,nao,nao)) 178 if singlet: 179 veff1[:3] = vj * 2 180 else: 181 veff1[:2] = vj[:2] * 2 182 183 fxcz1 = _contract_xc_kernel(td_grad, mf.xc, z1ao, None, 184 False, False, True, max_memory)[0] 185 186 veff1[0] += vxc1[1:] 187 veff1[1] +=(f1oo[1:] + fxcz1[1:] + k1ao[1:]*2)*2 # *2 for dmz1doo+dmz1oo.T 188 veff1[2] += f1vo[1:] * 2 189 time1 = log.timer('2e AO integral derivatives', *time1) 190 191 if atmlst is None: 192 atmlst = range(mol.natm) 193 offsetdic = mol.offset_nr_by_atom() 194 de = numpy.zeros((len(atmlst),3)) 195 for k, ia in enumerate(atmlst): 196 shl0, shl1, p0, p1 = offsetdic[ia] 197 198 # Ground state gradients 199 h1ao = hcore_deriv(ia) 200 h1ao[:,p0:p1] += veff1[0,:,p0:p1] 201 h1ao[:,:,p0:p1] += veff1[0,:,p0:p1].transpose(0,2,1) 202 # oo0*2 for doubly occupied orbitals 203 e1 = numpy.einsum('xpq,pq->x', h1ao, oo0) * 2 204 205 e1 += numpy.einsum('xpq,pq->x', h1ao, dmz1doo) 206 e1 -= numpy.einsum('xpq,pq->x', s1[:,p0:p1], im0[p0:p1]) 207 e1 -= numpy.einsum('xqp,pq->x', s1[:,p0:p1], im0[:,p0:p1]) 208 209 e1 += numpy.einsum('xij,ij->x', veff1[1,:,p0:p1], oo0[p0:p1]) 210 e1 += numpy.einsum('xij,ij->x', veff1[2,:,p0:p1], dmxpy[p0:p1,:]) * 2 211 e1 += numpy.einsum('xij,ij->x', veff1[3,:,p0:p1], dmxmy[p0:p1,:]) * 2 212 e1 += numpy.einsum('xji,ij->x', veff1[2,:,p0:p1], dmxpy[:,p0:p1]) * 2 213 e1 -= numpy.einsum('xji,ij->x', veff1[3,:,p0:p1], dmxmy[:,p0:p1]) * 2 214 215 de[k] = e1 216 217 log.timer('TDDFT nuclear gradients', *time0) 218 return de 219 220 221# dmvo, dmoo in AO-representation 222# Note spin-trace is applied for fxc, kxc 223#TODO: to include the response of grids 224def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True, 225 with_kxc=True, singlet=True, max_memory=2000): 226 mol = td_grad.mol 227 mf = td_grad.base._scf 228 grids = mf.grids 229 230 ni = mf._numint 231 xctype = ni._xc_type(xc_code) 232 233 mo_coeff = mf.mo_coeff 234 mo_occ = mf.mo_occ 235 nao, nmo = mo_coeff.shape 236 shls_slice = (0, mol.nbas) 237 ao_loc = mol.ao_loc_nr() 238 239 # dmvo ~ reduce(numpy.dot, (orbv, Xai, orbo.T)) 240 dmvo = (dmvo + dmvo.T) * .5 # because K_{ia,jb} == K_{ia,jb} 241 242 f1vo = numpy.zeros((4,nao,nao)) # 0th-order, d/dx, d/dy, d/dz 243 deriv = 2 244 if dmoo is not None: 245 f1oo = numpy.zeros((4,nao,nao)) 246 else: 247 f1oo = None 248 if with_vxc: 249 v1ao = numpy.zeros((4,nao,nao)) 250 else: 251 v1ao = None 252 if with_kxc: 253 k1ao = numpy.zeros((4,nao,nao)) 254 deriv = 3 255 else: 256 k1ao = None 257 258 if xctype == 'LDA': 259 ao_deriv = 1 260 if singlet: 261 for ao, mask, weight, coords \ 262 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory): 263 rho = ni.eval_rho2(mol, ao[0], mo_coeff, mo_occ, mask, 'LDA') 264 vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 0, deriv=deriv)[1:] 265 266 wfxc = fxc[0] * weight * 2 # *2 for alpha+beta 267 rho1 = ni.eval_rho(mol, ao[0], dmvo, mask, 'LDA') 268 aow = numpy.einsum('pi,p,p->pi', ao[0], wfxc, rho1) 269 for k in range(4): 270 f1vo[k] += numint._dot_ao_ao(mol, ao[k], aow, mask, shls_slice, ao_loc) 271 if dmoo is not None: 272 rho2 = ni.eval_rho(mol, ao[0], dmoo, mask, 'LDA') 273 aow = numpy.einsum('pi,p,p->pi', ao[0], wfxc, rho2) 274 for k in range(4): 275 f1oo[k] += numint._dot_ao_ao(mol, ao[k], aow, mask, shls_slice, ao_loc) 276 if with_vxc: 277 aow = numpy.einsum('pi,p,p->pi', ao[0], vxc[0], weight) 278 for k in range(4): 279 v1ao[k] += numint._dot_ao_ao(mol, ao[k], aow, mask, shls_slice, ao_loc) 280 if with_kxc: 281 aow = numpy.einsum('pi,p,p,p->pi', ao[0], kxc[0], weight, rho1**2) 282 for k in range(4): 283 k1ao[k] += numint._dot_ao_ao(mol, ao[k], aow, mask, shls_slice, ao_loc) 284 vxc = fxc = kxc = aow = rho = rho1 = rho2 = None 285 if with_kxc: # for (rho1*2)^2, *2 for alpha+beta in singlet 286 k1ao *= 4 287 288 else: 289 raise NotImplementedError('LDA triplet') 290 291 elif xctype == 'GGA': 292 if singlet: 293 def gga_sum_(vmat, ao, wv, mask): 294 aow = numpy.einsum('pi,p->pi', ao[0], wv[0]) 295 aow += numpy.einsum('npi,np->pi', ao[1:4], wv[1:]) 296 tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc) 297 vmat[0] += tmp + tmp.T 298 rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv, mask, ao_loc) 299 ao_deriv = 2 300 for ao, mask, weight, coords \ 301 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory): 302 rho = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, 'GGA') 303 vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 0, deriv=deriv)[1:] 304 305 rho1 = ni.eval_rho(mol, ao, dmvo, mask, 'GGA') * 2 # *2 for alpha + beta 306 wv = numint._rks_gga_wv1(rho, rho1, vxc, fxc, weight) 307 gga_sum_(f1vo, ao, wv, mask) 308 309 if dmoo is not None: 310 rho2 = ni.eval_rho(mol, ao, dmoo, mask, 'GGA') * 2 311 wv = numint._rks_gga_wv1(rho, rho2, vxc, fxc, weight) 312 gga_sum_(f1oo, ao, wv, mask) 313 if with_vxc: 314 wv = numint._rks_gga_wv0(rho, vxc, weight) 315 gga_sum_(v1ao, ao, wv, mask) 316 if with_kxc: 317 wv = numint._rks_gga_wv2(rho, rho1, fxc, kxc, weight) 318 gga_sum_(k1ao, ao, wv, mask) 319 vxc = fxc = kxc = rho = rho1 = None 320 321 else: 322 raise NotImplementedError('GGA triplet') 323 324 else: 325 raise NotImplementedError('meta-GGA') 326 327 f1vo[1:] *= -1 328 if f1oo is not None: f1oo[1:] *= -1 329 if v1ao is not None: v1ao[1:] *= -1 330 if k1ao is not None: k1ao[1:] *= -1 331 return f1vo, f1oo, v1ao, k1ao 332 333 334class Gradients(tdrhf.Gradients): 335 @lib.with_doc(grad_elec.__doc__) 336 def grad_elec(self, xy, singlet, atmlst=None): 337 return grad_elec(self, xy, singlet, atmlst, self.max_memory, self.verbose) 338 339Grad = Gradients 340 341from pyscf import tdscf 342tdscf.rks.TDA.Gradients = tdscf.rks.TDDFT.Gradients = lib.class_as_method(Gradients) 343 344 345if __name__ == '__main__': 346 from pyscf import gto 347 from pyscf import dft 348 from pyscf import tddft 349 mol = gto.Mole() 350 mol.verbose = 0 351 mol.output = None 352 353 mol.atom = [ 354 ['H' , (0. , 0. , 1.804)], 355 ['F' , (0. , 0. , 0.)], ] 356 mol.unit = 'B' 357 mol.basis = '631g' 358 mol.build() 359 360 mf = dft.RKS(mol) 361 mf.xc = 'LDA,' 362 mf.grids.prune = False 363 mf.conv_tol = 1e-14 364# mf.grids.level = 6 365 mf.scf() 366 367 td = tddft.TDDFT(mf) 368 td.nstates = 3 369 e, z = td.kernel() 370 tdg = td.Gradients() 371 g1 = tdg.kernel(state=3) 372 print(g1) 373# [[ 0 0 -1.31315477e-01] 374# [ 0 0 1.31319442e-01]] 375 td_solver = td.as_scanner() 376 e1 = td_solver(mol.set_geom_('H 0 0 1.805; F 0 0 0', unit='B')) 377 e2 = td_solver(mol.set_geom_('H 0 0 1.803; F 0 0 0', unit='B')) 378 print(abs((e1[2]-e2[2])/.002 - g1[0,2]).max()) 379 380 mol.set_geom_('H 0 0 1.804; F 0 0 0', unit='B') 381 mf = dft.RKS(mol) 382 mf.xc = 'b3lyp' 383 mf._numint.libxc = dft.xcfun 384 mf.grids.prune = False 385 mf.conv_tol = 1e-14 386 mf.scf() 387 388 td = tddft.TDA(mf) 389 td.nstates = 3 390 e, z = td.kernel() 391 tdg = td.Gradients() 392 g1 = tdg.kernel(state=3) 393 print(g1) 394# [[ 0 0 -1.21504524e-01] 395# [ 0 0 1.21505341e-01]] 396 td_solver = td.as_scanner() 397 e1 = td_solver(mol.set_geom_('H 0 0 1.805; F 0 0 0', unit='B')) 398 e2 = td_solver(mol.set_geom_('H 0 0 1.803; F 0 0 0', unit='B')) 399 print(abs((e1[2]-e2[2])/.002 - g1[0,2]).max()) 400 401 mol.set_geom_('H 0 0 1.804; F 0 0 0', unit='B') 402 mf = dft.RKS(mol) 403 mf.xc = 'lda,' 404 mf.conv_tol = 1e-14 405 mf.kernel() 406 td = tddft.TDA(mf) 407 td.nstates = 3 408 td.singlet = False 409 e, z = td.kernel() 410 tdg = Gradients(td) 411 g1 = tdg.kernel(state=2) 412 print(g1) 413# [[ 0 0 -0.3633334] 414# [ 0 0 0.3633334]] 415 td_solver = td.as_scanner() 416 e1 = td_solver(mol.set_geom_('H 0 0 1.805; F 0 0 0', unit='B')) 417 e2 = td_solver(mol.set_geom_('H 0 0 1.803; F 0 0 0', unit='B')) 418 print(abs((e1[2]-e2[2])/.002 - g1[0,2]).max()) 419 420 mf = dft.RKS(mol) 421 mf.xc = 'b3lyp' 422 mf.conv_tol = 1e-14 423 mf.kernel() 424 td = tddft.TDA(mf) 425 td.nstates = 3 426 td.singlet = False 427 e, z = td.kernel() 428 tdg = td.Gradients() 429 g1 = tdg.kernel(state=2) 430 print(g1) 431# [[ 0 0 -0.3633334] 432# [ 0 0 0.3633334]] 433 td_solver = td.as_scanner() 434 e1 = td_solver(mol.set_geom_('H 0 0 1.805; F 0 0 0', unit='B')) 435 e2 = td_solver(mol.set_geom_('H 0 0 1.803; F 0 0 0', unit='B')) 436 print(abs((e1[2]-e2[2])/.002 - g1[0,2]).max()) 437