def test_mgga_deriv2(self):
ng = 7
xctype = 'MGGA'
np.random.seed(8)
rho = np.random.rand(2, 5, ng)
rho1 = np.random.rand(2, 5, ng)
weight = 1
exc, vxc, fxc, kxc = eval_xc(f'R-{xctype}', ng)
ref = v6to5(
numint._rks_mgga_wv1(v5to6(rho[0]), v5to6(rho1[0]), vxc, fxc,
weight))
ref[0] *= 2
ref[4] *= 4
v1 = xc_deriv.transform_fxc(rho[0], vxc, fxc, xctype, spin=0)
v1 = np.einsum('xg,xyg->yg', rho1[0], v1)
self.assertAlmostEqual(abs(v1 - ref).max(), 0, 12)
exc, vxc, fxc, kxc = eval_xc(f'U-{xctype}', ng)
ref = v6to5(
np.array(
numint._uks_mgga_wv1(v5to6(rho), v5to6(rho1), vxc, fxc,
weight)))
ref[:, 0] *= 2
ref[:, 4] *= 4
v1 = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=1)
v1 = np.einsum('axg,axbyg->byg', rho1, v1)
self.assertAlmostEqual(abs(v1 - ref).max(), 0, 12)
def _contract_xc_kernel(td_grad,
xc_code,
dmvo,
dmoo=None,
with_vxc=True,
with_kxc=True,
singlet=True,
max_memory=2000):
mol = td_grad.mol
mf = td_grad.base._scf
grids = mf.grids
ni = mf._numint
xctype = ni._xc_type(xc_code)
mo_coeff = mf.mo_coeff
mo_occ = mf.mo_occ
nao, nmo = mo_coeff.shape
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_nr()
# dmvo ~ reduce(numpy.dot, (orbv, Xai, orbo.T))
dmvo = (dmvo + dmvo.T) * .5 # because K_{ia,jb} == K_{ia,jb}
f1vo = numpy.zeros((4, nao, nao)) # 0th-order, d/dx, d/dy, d/dz
deriv = 2
if dmoo is not None:
f1oo = numpy.zeros((4, nao, nao))
else:
f1oo = None
if with_vxc:
v1ao = numpy.zeros((4, nao, nao))
else:
v1ao = None
if with_kxc:
k1ao = numpy.zeros((4, nao, nao))
deriv = 3
else:
k1ao = None
if xctype == 'LDA':
ao_deriv = 1
if singlet:
def lda_sum_(vmat, ao, wv, mask):
aow = numint._scale_ao(ao[0], wv)
for k in range(4):
vmat[k] += numint._dot_ao_ao(mol, ao[k], aow, mask,
shls_slice, ao_loc)
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[0], mo_coeff, mo_occ, mask, xctype)
vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 0, deriv=deriv)[1:]
wfxc = fxc[0] * weight * 2 # *2 for alpha+beta
rho1 = ni.eval_rho(mol, ao[0], dmvo, mask, xctype)
lda_sum_(f1vo, ao, wfxc * rho1, mask)
if dmoo is not None:
rho2 = ni.eval_rho(mol, ao[0], dmoo, mask, xctype)
lda_sum_(f1oo, ao, wfxc * rho2, mask)
if with_vxc:
lda_sum_(v1ao, ao, vxc[0] * weight, mask)
if with_kxc:
lda_sum_(k1ao, ao, kxc[0] * weight * rho1**2, mask)
if with_kxc: # for (rho1*2)^2, *2 for alpha+beta in singlet
k1ao *= 4
else:
raise NotImplementedError('LDA triplet')
elif xctype == 'GGA':
if singlet:
def gga_sum_(vmat, ao, wv, mask):
aow = numint._scale_ao(ao[:4], wv[:4])
tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice,
ao_loc)
vmat[0] += tmp + tmp.T
rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv, mask, ao_loc)
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, xctype)
vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 0, deriv=deriv)[1:]
rho1 = ni.eval_rho(mol, ao, dmvo, mask,
xctype) * 2 # *2 for alpha + beta
wv = numint._rks_gga_wv1(rho, rho1, vxc, fxc, weight)
gga_sum_(f1vo, ao, wv, mask)
if dmoo is not None:
rho2 = ni.eval_rho(mol, ao, dmoo, mask, xctype) * 2
wv = numint._rks_gga_wv1(rho, rho2, vxc, fxc, weight)
gga_sum_(f1oo, ao, wv, mask)
if with_vxc:
wv = numint._rks_gga_wv0(rho, vxc, weight)
gga_sum_(v1ao, ao, wv, mask)
if with_kxc:
wv = numint._rks_gga_wv2(rho, rho1, fxc, kxc, weight)
gga_sum_(k1ao, ao, wv, mask)
vxc = fxc = kxc = rho = rho1 = None
else:
raise NotImplementedError('GGA triplet')
elif xctype == 'MGGA':
if singlet:
def mgga_sum_(vmat, ao, wv, mask):
aow = numint._scale_ao(ao[:4], wv[:4])
tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice,
ao_loc)
aow = numint._scale_ao(ao[1], wv[5], aow)
tmp += numint._dot_ao_ao(mol, ao[1], aow, mask, shls_slice,
ao_loc)
aow = numint._scale_ao(ao[2], wv[5], aow)
tmp += numint._dot_ao_ao(mol, ao[2], aow, mask, shls_slice,
ao_loc)
aow = numint._scale_ao(ao[3], wv[5], aow)
tmp += numint._dot_ao_ao(mol, ao[3], aow, mask, shls_slice,
ao_loc)
vmat[0] += tmp + tmp.T
rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv[:4], mask,
ao_loc)
rks_grad._tau_grad_dot_(vmat[1:], mol, ao, wv[5] * 2, mask,
ao_loc, True)
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, xctype)
vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 0, deriv=deriv)[1:]
rho1 = ni.eval_rho(mol, ao, dmvo, mask,
xctype) * 2 # *2 for alpha + beta
wv = numint._rks_mgga_wv1(rho, rho1, vxc, fxc, weight)
mgga_sum_(f1vo, ao, wv, mask)
if dmoo is not None:
rho2 = ni.eval_rho(mol, ao, dmoo, mask, xctype) * 2
wv = numint._rks_mgga_wv1(rho, rho2, vxc, fxc, weight)
mgga_sum_(f1oo, ao, wv, mask)
if with_vxc:
wv = numint._rks_mgga_wv0(rho, vxc, weight)
mgga_sum_(v1ao, ao, wv, mask)
if with_kxc:
wv = numint._rks_mgga_wv2(rho, rho1, fxc, kxc, weight)
mgga_sum_(k1ao, ao, wv, mask)
vxc = fxc = kxc = rho = rho1 = None
else:
raise NotImplementedError('MGGA triplet')
elif xctype == 'HF':
pass
else:
raise NotImplementedError(f'td-rks for functional {xc_code}')
f1vo[1:] *= -1
if f1oo is not None: f1oo[1:] *= -1
if v1ao is not None: v1ao[1:] *= -1
if k1ao is not None: k1ao[1:] *= -1
return f1vo, f1oo, v1ao, k1ao
def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
mol = hessobj.mol
mf = hessobj.base
if hessobj.grids is not None:
grids = hessobj.grids
else:
grids = mf.grids
if grids.coords is None:
grids.build(with_non0tab=True)
nao, nmo = mo_coeff.shape
ni = mf._numint
xctype = ni._xc_type(mf.xc)
aoslices = mol.aoslice_by_atom()
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_nr()
dm0 = mf.make_rdm1(mo_coeff, mo_occ)
v_ip = numpy.zeros((3, nao, nao))
vmat = numpy.zeros((mol.natm, 3, nao, nao))
max_memory = max(2000, max_memory - vmat.size * 8 / 1e6)
if xctype == 'LDA':
ao_deriv = 1
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[0], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
aow = numint._scale_ao(ao[0], weight * vrho)
rks_grad._d1_dot_(v_ip, mol, ao[1:4], aow, mask, ao_loc, True)
ao_dm0 = numint._dot_ao_dm(mol, ao[0], dm0, mask, shls_slice,
ao_loc)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
# First order density = rho1 * 2. *2 is not applied because + c.c. in the end
rho1 = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
ao_dm0[:, p0:p1])
wv = weight * frr * rho1
aow = [numint._scale_ao(ao[0], wv[i]) for i in range(3)]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
True)
ao_dm0 = aow = None
elif xctype == 'GGA':
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[:4], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
wv = numint._rks_gga_wv0(rho, vxc, weight)
rks_grad._gga_grad_sum_(v_ip, mol, ao, wv, mask, ao_loc)
ao_dm0 = [
numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices)
wv[0] = numint._rks_gga_wv1(rho, dR_rho1[0], vxc, fxc, weight)
wv[1] = numint._rks_gga_wv1(rho, dR_rho1[1], vxc, fxc, weight)
wv[2] = numint._rks_gga_wv1(rho, dR_rho1[2], vxc, fxc, weight)
aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
True)
ao_dm0 = aow = None
elif xctype == 'MGGA':
if grids.level < 5:
logger.warn(mol, 'MGGA Hessian is sensitive to dft grids.')
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[:10], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
wv = numint._rks_mgga_wv0(rho, vxc, weight)
rks_grad._gga_grad_sum_(v_ip, mol, ao, wv, mask, ao_loc)
# *2 because wv[5] is scaled by 0.5 in _rks_mgga_wv0
rks_grad._tau_grad_dot_(v_ip, mol, ao, wv[5] * 2, mask, ao_loc,
True)
ao_dm0 = [
numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices, xctype)
wv[0] = numint._rks_mgga_wv1(rho, dR_rho1[0], vxc, fxc, weight)
wv[1] = numint._rks_mgga_wv1(rho, dR_rho1[1], vxc, fxc, weight)
wv[2] = numint._rks_mgga_wv1(rho, dR_rho1[2], vxc, fxc, weight)
aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
True)
for j in range(1, 4):
aow = [numint._scale_ao(ao[j], wv[i, 5]) for i in range(3)]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[j], mask, ao_loc,
True)
ao_dm0 = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmat[ia, :, p0:p1] += v_ip[:, p0:p1]
vmat[ia] = -vmat[ia] - vmat[ia].transpose(0, 2, 1)
return vmat
def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
mol = hessobj.mol
mf = hessobj.base
if hessobj.grids is not None:
grids = hessobj.grids
else:
grids = mf.grids
if grids.coords is None:
grids.build(with_non0tab=True)
nao, nmo = mo_coeff.shape
ni = mf._numint
xctype = ni._xc_type(mf.xc)
aoslices = mol.aoslice_by_atom()
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_nr()
dm0 = mf.make_rdm1(mo_coeff, mo_occ)
vmat = numpy.zeros((mol.natm, 3, 3, nao, nao))
ipip = numpy.zeros((3, 3, nao, nao))
if xctype == 'LDA':
ao_deriv = 1
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[0], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
wv = weight * vrho
aow = [numint._scale_ao(ao[i], wv) for i in range(1, 4)]
_d1d2_dot_(ipip, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0 = numint._dot_ao_dm(mol, ao[0], dm0, mask, shls_slice,
ao_loc)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
# *2 for \nabla|ket> in rho1
rho1 = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
ao_dm0[:, p0:p1]) * 2
# aow ~ rho1 ~ d/dR1
wv = weight * frr * rho1
aow = [numint._scale_ao(ao[0], wv[i]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0 = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmat[ia, :, :, :, p0:p1] += ipip[:, :, :, p0:p1]
elif xctype == 'GGA':
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[:4], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
wv = numint._rks_gga_wv0(rho, vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wv)
_d1d2_dot_(ipip, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0 = [
numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices)
for i in range(3):
wv[i] = numint._rks_gga_wv1(rho, dR_rho1[i], vxc, fxc,
weight)
aow = rks_grad._make_dR_dao_w(ao, wv[i])
rks_grad._d1_dot_(vmat[ia, i], mol, aow, ao[0], mask,
ao_loc, True)
aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0 = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmat[ia, :, :, :, p0:p1] += ipip[:, :, :, p0:p1]
vmat[ia, :, :, :, p0:p1] += ipip[:, :, p0:p1].transpose(1, 0, 3, 2)
elif xctype == 'MGGA':
XX, XY, XZ = 4, 5, 6
YX, YY, YZ = 5, 7, 8
ZX, ZY, ZZ = 6, 8, 9
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rho = ni.eval_rho2(mol, ao[:10], mo_coeff, mo_occ, mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
wv = numint._rks_mgga_wv0(rho, vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wv)
_d1d2_dot_(ipip, mol, aow, ao[1:4], mask, ao_loc, False)
aow = [numint._scale_ao(ao[i], wv[5]) for i in range(4, 10)]
_d1d2_dot_(ipip, mol, [aow[0], aow[1], aow[2]],
[ao[XX], ao[XY], ao[XZ]], mask, ao_loc, False)
_d1d2_dot_(ipip, mol, [aow[1], aow[3], aow[4]],
[ao[YX], ao[YY], ao[YZ]], mask, ao_loc, False)
_d1d2_dot_(ipip, mol, [aow[2], aow[4], aow[5]],
[ao[ZX], ao[ZY], ao[ZZ]], mask, ao_loc, False)
ao_dm0 = [
numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices, xctype)
for i in range(3):
wv[i] = numint._rks_mgga_wv1(rho, dR_rho1[i], vxc, fxc,
weight)
aow = rks_grad._make_dR_dao_w(ao, wv[i])
rks_grad._d1_dot_(vmat[ia, i], mol, aow, ao[0], mask,
ao_loc, True)
aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, ao[1:4], aow, mask, ao_loc, False)
# *2 because wv[5] is scaled by 0.5 in _rks_mgga_wv1
wv[:, 5] *= 2
aow = [numint._scale_ao(ao[1], wv[i, 5]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, [ao[XX], ao[XY], ao[XZ]], aow, mask,
ao_loc, False)
aow = [numint._scale_ao(ao[2], wv[i, 5]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, [ao[YX], ao[YY], ao[YZ]], aow, mask,
ao_loc, False)
aow = [numint._scale_ao(ao[3], wv[i, 5]) for i in range(3)]
_d1d2_dot_(vmat[ia], mol, [ao[ZX], ao[ZY], ao[ZZ]], aow, mask,
ao_loc, False)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmat[ia, :, :, :, p0:p1] += ipip[:, :, :, p0:p1]
vmat[ia, :, :, :, p0:p1] += ipip[:, :, p0:p1].transpose(1, 0, 3, 2)
return vmat