def dot_dm_ket(ket_re, ket_im):
# DM * ket: e.g. ir denotes dm_im | ao_re >
c0_rr = _dot_ao_dm(cell, ket_re, dm_re, nao, ngrids, non0tab)
c0_ri = _dot_ao_dm(cell, ket_im, dm_re, nao, ngrids, non0tab)
c0_ir = _dot_ao_dm(cell, ket_re, dm_im, nao, ngrids, non0tab)
c0_ii = _dot_ao_dm(cell, ket_im, dm_im, nao, ngrids, non0tab)
return c0_ri, c0_rr, c0_ir, c0_ii
def dot_dm_ket(ket_re, ket_im):
# DM * ket: e.g. ir denotes dm_im | ao_re >
c0_rr = _dot_ao_dm(cell, ket_re, dm_re, nao, ngrids, non0tab)
c0_ri = _dot_ao_dm(cell, ket_im, dm_re, nao, ngrids, non0tab)
c0_ir = _dot_ao_dm(cell, ket_re, dm_im, nao, ngrids, non0tab)
c0_ii = _dot_ao_dm(cell, ket_im, dm_im, nao, ngrids, non0tab)
return c0_ri, c0_rr, c0_ir, c0_ii
def _dm2c_to_rho2x2(mol, ao, dm, non0tab, shls_slice, ao_loc, out=None):
aoa, aob = ao
out = _dot_ao_dm(mol, aoa, dm, non0tab, shls_slice, ao_loc, out=out)
rhoaa = numpy.einsum('pi,pi->p', aoa.real, out.real)
rhoaa+= numpy.einsum('pi,pi->p', aoa.imag, out.imag)
rhoba = numpy.einsum('pi,pi->p', aob, out.conj())
out = _dot_ao_dm(mol, aob, dm, non0tab, shls_slice, ao_loc, out=out)
rhoab = numpy.einsum('pi,pi->p', aoa, out.conj())
rhobb = numpy.einsum('pi,pi->p', aob.real, out.real)
rhobb+= numpy.einsum('pi,pi->p', aob.imag, out.imag)
return rhoaa, rhoab, rhoba, rhobb
def _dm2c_to_rho2x2(mol, ao, dm, non0tab, shls_slice, ao_loc, out=None):
aoa, aob = ao
out = _dot_ao_dm(mol, aoa, dm, non0tab, shls_slice, ao_loc, out=out)
rhoaa = numpy.einsum('pi,pi->p', aoa.real, out.real)
rhoaa += numpy.einsum('pi,pi->p', aoa.imag, out.imag)
rhoba = numpy.einsum('pi,pi->p', aob, out.conj())
out = _dot_ao_dm(mol, aob, dm, non0tab, shls_slice, ao_loc, out=out)
rhoab = numpy.einsum('pi,pi->p', aoa, out.conj())
rhobb = numpy.einsum('pi,pi->p', aob.real, out.real)
rhobb += numpy.einsum('pi,pi->p', aob.imag, out.imag)
return rhoaa, rhoab, rhoba, rhobb
def _grid_ao2mo(mol, ao, mo_coeff, non0tab=None, shls_slice=None, ao_loc=None):
''' ao[deriv,grid,AO].mo_coeff[AO,MO]->mo[deriv,grid,MO]
ASSUMES that ao is in data layout (deriv,AO,grid) in row-major order!
mo is returned in data layout (deriv,MO,grid) in row-major order '''
nderiv, ngrid, nao = ao.shape
nmo = mo_coeff.shape[-1]
mo = np.empty((nderiv, nmo, ngrid), dtype=mo_coeff.dtype,
order='C').transpose(0, 2, 1)
if shls_slice is None: shls_slice = (0, mol.nbas)
if ao_loc is None: ao_loc = mol.ao_loc_nr()
for ideriv in range(nderiv):
ao_i = ao[ideriv, :, :]
mo[ideriv] = _dot_ao_dm(mol,
ao_i,
mo_coeff,
non0tab,
shls_slice,
ao_loc,
out=mo[ideriv])
return mo
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[0].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()
dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)
vmata = numpy.zeros((mol.natm,3,3,nao,nao))
vmatb = numpy.zeros((mol.natm,3,3,nao,nao))
ipipa = numpy.zeros((3,3,nao,nao))
ipipb = 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):
rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
vrho = vxc[0]
u_u, u_d, d_d = fxc[0].T
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,0])
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,1])
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
# *2 for \nabla|ket> in rho1
rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1]) * 2
rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1]) * 2
wv = u_u * rho1a + u_d * rho1b
wv *= weight
# aow ~ rho1 ~ d/dR1
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)
wv = u_d * rho1a + d_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
elif xctype == 'GGA':
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva)
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb)
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)]
ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva[0])
rks_grad._d1_dot_(vmata[ia,0], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wva[1])
rks_grad._d1_dot_(vmata[ia,1], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wva[2])
rks_grad._d1_dot_(vmata[ia,2], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb[0])
rks_grad._d1_dot_(vmatb[ia,0], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wvb[1])
rks_grad._d1_dot_(vmatb[ia,1], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wvb[2])
rks_grad._d1_dot_(vmatb[ia,2], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,p0:p1].transpose(1,0,3,2)
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,p0:p1].transpose(1,0,3,2)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
return vmata, vmatb
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)
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, 'LDA')
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
ao_dm0 = numint._dot_ao_dm(mol, ao[0], dm0, mask, shls_slice,
ao_loc)
aow1 = numpy.einsum('xpi,p->xpi', ao[1:], weight * vrho)
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])
aow = numpy.einsum('pi,xp->xpi', ao[0], weight * frr * rho1)
aow[:, :, p0:p1] += aow1[:, :, p0:p1]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
True)
ao_dm0 = aow = aow1 = None
for ia in range(mol.natm):
vmat[ia] = -vmat[ia] - vmat[ia].transpose(0, 2, 1)
elif xctype == 'GGA':
ao_deriv = 2
v_ip = numpy.zeros((3, nao, nao))
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, 'GGA')
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 = numpy.einsum('npi,Xnp->Xpi', ao[:4], wv)
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], 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)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
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, 'LDA')
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight * vrho)
_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
aow = numpy.einsum('pi,xp->xpi', ao[0], weight * frr * rho1)
_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, 'GGA')
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)
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 = rks_grad._make_dR_dao_w(ao, wv[0])
rks_grad._d1_dot_(vmat[ia, 0], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wv[1])
rks_grad._d1_dot_(vmat[ia, 1], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wv[2])
rks_grad._d1_dot_(vmat[ia, 2], mol, aow, ao[0], mask, ao_loc,
True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wv)
_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':
raise NotImplementedError('meta-GGA')
return vmat
def _dot_spinor_dm(mol, ket, dm, non0tab, shls_slice, ao_loc):
ket_a, ket_b = ket
outa = _dot_ao_dm(mol, ket_a, dm, non0tab, shls_slice, ao_loc)
outb = _dot_ao_dm(mol, ket_b, dm, non0tab, shls_slice, ao_loc)
return outa, outb
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[0].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()
dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)
vmata = numpy.zeros((mol.natm,3,nao,nao))
vmatb = numpy.zeros((mol.natm,3,nao,nao))
max_memory = max(2000, max_memory-(vmata.size+vmatb.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):
rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
vrho = vxc[0]
u_u, u_d, d_d = fxc[0].T
ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
aow1a = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
aow1b = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
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
rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1])
rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1])
wv = u_u * rho1a + u_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
aow[:,:,p0:p1] += aow1a[:,:,p0:p1]
rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
wv = u_d * rho1a + d_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
aow[:,:,p0:p1] += aow1b[:,:,p0:p1]
rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0a = ao_dm0b = aow = aow1a = aow1b = None
for ia in range(mol.natm):
vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)
elif xctype == 'GGA':
ao_deriv = 2
vipa = numpy.zeros((3,nao,nao))
vipb = numpy.zeros((3,nao,nao))
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
rks_grad._gga_grad_sum_(vipa, mol, ao, wva, mask, ao_loc)
rks_grad._gga_grad_sum_(vipb, mol, ao, wvb, mask, ao_loc)
ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)]
ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,p0:p1] += vipa[:,p0:p1]
vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
vmatb[ia,:,p0:p1] += vipb[:,p0:p1]
vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
return vmata, vmatb
def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', verbose=None):
'''Collocate the *real* density (opt. gradients) on the real-space grid.
Args:
mol : instance of :class:`Mole` or :class:`Cell`
ao : ([4,] nx*ny*nz, nao=cell.nao_nr()) ndarray
The value of the AO crystal orbitals on the real-space grid by default.
If xctype='GGA', also contains the value of the gradient in the x, y,
and z directions.
Returns:
rho : ([4,] nx*ny*nz) ndarray
The value of the density on the real-space grid. If xctype='GGA',
also contains the value of the gradient in the x, y, and z
directions.
See Also:
pyscf.dft.numint.eval_rho
'''
assert(ao.flags.c_contiguous)
if xctype == 'GGA':
ngrids, nao = ao[0].shape
else:
ngrids, nao = ao.shape
if non0tab is None:
non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
dtype=numpy.int8)
# if xctype == 'GGA':
# rho = numpy.empty((4,ngrids))
# c0 = _dot_ao_dm(mol, ao[0], dm, nao, ngrids, non0tab)
# rho[0] = numpy.einsum('pi,pi->p', ao[0], c0)
# for i in range(1, 4):
# c1 = _dot_ao_dm(mol, ao[i], dm, nao, ngrids, non0tab)
# rho[i] = numpy.einsum('pi,pi->p', ao[0], c1) * 2 # *2 for +c.c.
# else:
# c0 = _dot_ao_dm(mol, ao, dm, nao, ngrids, non0tab)
# rho = numpy.einsum('pi,pi->p', ao, c0)
# return rho
# if xctype == 'GGA':
# ngrids, nao = ao[0].shape
# else:
# ngrids, nao = ao.shape
# if non0tab is None:
# print "this PATH"
# non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
# dtype=numpy.int8)
# #if ao[0].dtype==numpy.complex128: # complex orbitals
# if True:
# #dm_re = numpy.ascontiguousarray(dm.real)
# dm_re = dm
# rho = numpy.empty((4,ngrids))
# #ao_re = numpy.ascontiguousarray(ao[0].real)
# ao_re = ao[0]
# c0_rr = _dot_ao_dm(mol, ao_re, dm_re, nao, ngrids, non0tab)
# rho[0] = (numpy.einsum('pi,pi->p', ao_re, c0_rr))
# for i in range(1, 4):
# c1 = _dot_ao_dm(mol, ao[i], dm, nao, ngrids, non0tab)
# rho[i] = numpy.einsum('pi,pi->p', ao[0], c1) * 2 # *2 for +c.c.
# for i in range(1, 4):
# # #ao_re = numpy.ascontiguousarray(ao[i].real)
# ao_re = ao[i]
# c1_rr = _dot_ao_dm(mol, ao_re, dm_re, nao, ngrids, non0tab)
# rho[i] = (numpy.einsum('pi,pi->p', ao_re, c1_rr)) *2
# return rho
# complex orbitals or density matrix
if numpy.iscomplexobj(ao) or numpy.iscomplexobj(dm):
dm_re = numpy.ascontiguousarray(dm.real)
dm_im = numpy.ascontiguousarray(dm.imag)
if xctype == 'GGA':
rho = numpy.empty((4,ngrids))
ao0_re = numpy.ascontiguousarray(ao[0].real)
ao0_im = numpy.ascontiguousarray(ao[0].imag)
# DM * ket: e.g. ir denotes dm_im | ao_re >
c0_rr = _dot_ao_dm(mol, ao0_re, dm_re, nao, ngrids, non0tab)
c0_ri = _dot_ao_dm(mol, ao0_im, dm_re, nao, ngrids, non0tab)
c0_ir = _dot_ao_dm(mol, ao0_re, dm_im, nao, ngrids, non0tab)
c0_ii = _dot_ao_dm(mol, ao0_im, dm_im, nao, ngrids, non0tab)
# bra * DM
rho[0] = (numpy.einsum('pi,pi->p', ao0_im, c0_ri) +
numpy.einsum('pi,pi->p', ao0_re, c0_rr) +
numpy.einsum('pi,pi->p', ao0_im, c0_ir) -
numpy.einsum('pi,pi->p', ao0_re, c0_ii))
for i in range(1, 4):
# ao_re = numpy.ascontiguousarray(ao[i].real)
# ao_im = numpy.ascontiguousarray(ao[i].imag)
ao_re = numpy.ascontiguousarray(ao[i].real)
ao_im = numpy.ascontiguousarray(ao[i].imag)
c1_rr = _dot_ao_dm(mol, ao_re, dm_re, nao, ngrids, non0tab)
c1_ri = _dot_ao_dm(mol, ao_im, dm_re, nao, ngrids, non0tab)
c1_ir = _dot_ao_dm(mol, ao_re, dm_im, nao, ngrids, non0tab)
c1_ii = _dot_ao_dm(mol, ao_im, dm_im, nao, ngrids, non0tab)
rho[i] = (numpy.einsum('pi,pi->p', ao0_im, c1_ri) +
numpy.einsum('pi,pi->p', ao0_re, c1_rr) +
numpy.einsum('pi,pi->p', ao0_im, c1_ir) -
numpy.einsum('pi,pi->p', ao0_re, c1_ii)) * 2 # *2 for +c.c.
else:
ao_re = numpy.ascontiguousarray(ao.real)
ao_im = numpy.ascontiguousarray(ao.imag)
# DM * ket: e.g. ir denotes dm_im | ao_re >
c0_rr = _dot_ao_dm(mol, ao_re, dm_re, nao, ngrids, non0tab)
c0_ri = _dot_ao_dm(mol, ao_im, dm_re, nao, ngrids, non0tab)
c0_ir = _dot_ao_dm(mol, ao_re, dm_im, nao, ngrids, non0tab)
c0_ii = _dot_ao_dm(mol, ao_im, dm_im, nao, ngrids, non0tab)
# bra * DM
rho = (numpy.einsum('pi,pi->p', ao_im, c0_ri) +
numpy.einsum('pi,pi->p', ao_re, c0_rr) +
numpy.einsum('pi,pi->p', ao_im, c0_ir) -
numpy.einsum('pi,pi->p', ao_re, c0_ii))
# real orbitals and real DM
else:
rho = pyscf.dft.numint.eval_rho(mol, ao, dm, non0tab, xctype, verbose)
return rho
def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0,
with_lapl=True, verbose=None):
'''Calculate the electron density and magnetization spin density in the
framework of 2-component real basis.
''' + numint.eval_rho.__doc__
nao = ao.shape[-1]
assert dm.ndim == 2 and nao * 2 == dm.shape[0]
ngrids, nao = ao.shape[-2:]
xctype = xctype.upper()
if non0tab is None:
non0tab = np.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
dtype=np.uint8)
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc
if xctype == 'LDA':
c0a = _dot_ao_dm(mol, ao, dm[:nao], non0tab, shls_slice, ao_loc)
c0b = _dot_ao_dm(mol, ao, dm[nao:], non0tab, shls_slice, ao_loc)
rho_m = _contract_rho_m((ao, ao), (c0a, c0b), hermi, True)
elif xctype == 'GGA':
# first 4 ~ (rho, m), second 4 ~ (0th order, dx, dy, dz)
if hermi:
rho_m = np.empty((4, 4, ngrids))
else:
rho_m = np.empty((4, 4, ngrids), dtype=np.complex128)
c0a = _dot_ao_dm(mol, ao[0], dm[:nao], non0tab, shls_slice, ao_loc)
c0b = _dot_ao_dm(mol, ao[0], dm[nao:], non0tab, shls_slice, ao_loc)
c0 = (c0a, c0b)
rho_m[:,0] = _contract_rho_m((ao[0], ao[0]), c0, hermi, True)
for i in range(1, 4):
rho_m[:,i] = _contract_rho_m((ao[i], ao[i]), c0, hermi, False)
if hermi:
rho_m[:,1:4] *= 2 # *2 for |ao> dm < dx ao| + |dx ao> dm < ao|
else:
for i in range(1, 4):
c1a = _dot_ao_dm(mol, ao[i], dm[:nao], non0tab, shls_slice, ao_loc)
c1b = _dot_ao_dm(mol, ao[i], dm[nao:], non0tab, shls_slice, ao_loc)
rho_m[:,i] += _contract_rho_m((ao[0], ao[0]), (c1a, c1b), hermi, False)
else: # meta-GGA
if hermi:
dtype = np.double
else:
dtype = np.complex128
if with_lapl:
rho_m = np.empty((4, 6, ngrids), dtype=dtype)
tau_idx = 5
else:
rho_m = np.empty((4, 5, ngrids), dtype=dtype)
tau_idx = 4
c0a = _dot_ao_dm(mol, ao[0], dm[:nao], non0tab, shls_slice, ao_loc)
c0b = _dot_ao_dm(mol, ao[0], dm[nao:], non0tab, shls_slice, ao_loc)
c0 = (c0a, c0b)
rho_m[:,0] = _contract_rho_m((ao[0], ao[0]), c0, hermi, True)
rho_m[:,tau_idx] = 0
for i in range(1, 4):
c1a = _dot_ao_dm(mol, ao[i], dm[:nao], non0tab, shls_slice, ao_loc)
c1b = _dot_ao_dm(mol, ao[i], dm[nao:], non0tab, shls_slice, ao_loc)
rho_m[:,tau_idx] += _contract_rho_m((ao[i], ao[i]), (c1a, c1b), hermi, True)
rho_m[:,i] = _contract_rho_m((ao[i], ao[i]), c0, hermi, False)
if hermi:
rho_m[:,i] *= 2
else:
rho_m[:,i] += _contract_rho_m((ao[0], ao[0]), (c1a, c1b), hermi, False)
if with_lapl:
# TODO: rho_m[:,4] = \nabla^2 rho
raise NotImplementedError
# tau = 1/2 (\nabla f)^2
rho_m[:,tau_idx] *= .5
return rho_m
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
def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
"""" This functions is slightly different from hessian.rks._get_vxc_deriv1 in that <\nabla u|Vxc|v> is removed"""
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 = np.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, 'LDA')
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
frr = fxc[0]
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:]
rho1 = np.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0[:,p0:p1])
aow = np.einsum('pi,xp->xpi', ao[0], weight*frr*rho1)
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0 = aow = None
for ia in range(mol.natm):
vmat[ia] = -vmat[ia] - vmat[ia].transpose(0,2,1)
elif xctype == 'GGA':
ao_deriv = 2
# v_ip = np.zeros((3,nao,nao))
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, 'GGA')
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 = rks_hess._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 = np.einsum('npi,Xnp->Xpi', ao[:4], wv)
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0 = aow = None
for ia in range(mol.natm):
vmat[ia] = -vmat[ia] - vmat[ia].transpose(0,2,1)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
return vmat
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[0].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()
dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)
vmata = numpy.zeros((mol.natm,3,nao,nao))
vmatb = numpy.zeros((mol.natm,3,nao,nao))
max_memory = max(2000, max_memory-(vmata.size+vmatb.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):
rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
vrho = vxc[0]
u_u, u_d, d_d = fxc[0].T
ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
aow1a = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
aow1b = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
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
rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1])
rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1])
wv = u_u * rho1a + u_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
aow[:,:,p0:p1] += aow1a[:,:,p0:p1]
rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
wv = u_d * rho1a + d_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
aow[:,:,p0:p1] += aow1b[:,:,p0:p1]
rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0a = ao_dm0b = aow = aow1a = aow1b = None
for ia in range(mol.natm):
vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)
elif xctype == 'GGA':
ao_deriv = 2
vipa = numpy.zeros((3,nao,nao))
vipb = numpy.zeros((3,nao,nao))
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
rks_grad._gga_grad_sum_(vipa, mol, ao, wva, mask, ao_loc)
rks_grad._gga_grad_sum_(vipb, mol, ao, wvb, mask, ao_loc)
ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)]
ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,p0:p1] += vipa[:,p0:p1]
vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
vmatb[ia,:,p0:p1] += vipb[:,p0:p1]
vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
return vmata, vmatb
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[0].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()
dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)
vmata = numpy.zeros((mol.natm,3,3,nao,nao))
vmatb = numpy.zeros((mol.natm,3,3,nao,nao))
ipipa = numpy.zeros((3,3,nao,nao))
ipipb = 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):
rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
vrho = vxc[0]
u_u, u_d, d_d = fxc[0].T
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,0])
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,1])
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
# *2 for \nabla|ket> in rho1
rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1]) * 2
rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1]) * 2
wv = u_u * rho1a + u_d * rho1b
wv *= weight
# aow ~ rho1 ~ d/dR1
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)
wv = u_d * rho1a + d_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
elif xctype == 'GGA':
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva)
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb)
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)]
ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva[0])
rks_grad._d1_dot_(vmata[ia,0], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wva[1])
rks_grad._d1_dot_(vmata[ia,1], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wva[2])
rks_grad._d1_dot_(vmata[ia,2], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb[0])
rks_grad._d1_dot_(vmatb[ia,0], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wvb[1])
rks_grad._d1_dot_(vmatb[ia,1], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wvb[2])
rks_grad._d1_dot_(vmatb[ia,2], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
vmata[ia,:,:,:,p0:p1] += ipipa[:,:,p0:p1].transpose(1,0,3,2)
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,p0:p1].transpose(1,0,3,2)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
return vmata, vmatb
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, 'LDA')
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho)
_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
aow = numpy.einsum('pi,xp->xpi', ao[0], weight*frr*rho1)
_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, 'GGA')
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)
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 = rks_grad._make_dR_dao_w(ao, wv[0])
rks_grad._d1_dot_(vmat[ia,0], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wv[1])
rks_grad._d1_dot_(vmat[ia,1], mol, aow, ao[0], mask, ao_loc, True)
aow = rks_grad._make_dR_dao_w(ao, wv[2])
rks_grad._d1_dot_(vmat[ia,2], mol, aow, ao[0], mask, ao_loc, True)
aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wv)
_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':
raise NotImplementedError('meta-GGA')
return vmat
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)
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, 'LDA')
vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
vrho = vxc[0]
frr = fxc[0]
ao_dm0 = numint._dot_ao_dm(mol, ao[0], dm0, mask, shls_slice, ao_loc)
aow1 = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho)
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])
aow = numpy.einsum('pi,xp->xpi', ao[0], weight*frr*rho1)
aow[:,:,p0:p1] += aow1[:,:,p0:p1]
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc, True)
ao_dm0 = aow = aow1 = None
for ia in range(mol.natm):
vmat[ia] = -vmat[ia] - vmat[ia].transpose(0,2,1)
elif xctype == 'GGA':
ao_deriv = 2
v_ip = numpy.zeros((3,nao,nao))
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, 'GGA')
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 = numpy.einsum('npi,Xnp->Xpi', ao[:4], wv)
rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], 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)
elif xctype == 'MGGA':
raise NotImplementedError('meta-GGA')
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[0].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()
dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)
vmata = numpy.zeros((mol.natm, 3, 3, nao, nao))
vmatb = numpy.zeros((mol.natm, 3, 3, nao, nao))
ipipa = numpy.zeros((3, 3, nao, nao))
ipipb = 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):
rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask,
xctype)
rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask,
xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]
vrho = vxc[0]
u_u, u_d, d_d = fxc[0].T
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight * vrho[:, 0])
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight * vrho[:, 1])
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice,
ao_loc)
ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice,
ao_loc)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
# *2 for \nabla|ket> in rho1
rho1a = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
ao_dm0a[:, p0:p1]) * 2
rho1b = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
ao_dm0b[:, p0:p1]) * 2
wv = u_u * rho1a + u_d * rho1b
wv *= weight
# aow ~ rho1 ~ d/dR1
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
wv = u_d * rho1a + d_d * rho1b
wv *= weight
aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]
elif xctype == 'GGA':
ao_deriv = 2
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask,
xctype)
rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask,
xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_gga_wv0((rhoa, rhob), vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva)
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb)
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
ao_dm0a = [
numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)
]
ao_dm0b = [
numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(
ao, ao_dm0a, ia, aoslices)
wvb = dR_rho1b = rks_hess._make_dR_rho1(
ao, ao_dm0b, ia, aoslices)
wva[0], wvb[0] = numint._uks_gga_wv1(
(rhoa, rhob), (dR_rho1a[0], dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_gga_wv1(
(rhoa, rhob), (dR_rho1a[1], dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_gga_wv1(
(rhoa, rhob), (dR_rho1a[2], dR_rho1b[2]), vxc, fxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva[0])
rks_grad._d1_dot_(vmata[ia, 0], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wva[1])
rks_grad._d1_dot_(vmata[ia, 1], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wva[2])
rks_grad._d1_dot_(vmata[ia, 2], mol, aow, ao[0], mask, ao_loc,
True)
aow = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
aow = rks_grad._make_dR_dao_w(ao, wvb[0])
rks_grad._d1_dot_(vmatb[ia, 0], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wvb[1])
rks_grad._d1_dot_(vmatb[ia, 1], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wvb[2])
rks_grad._d1_dot_(vmatb[ia, 2], mol, aow, ao[0], mask, ao_loc,
True)
aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
ao_dm0a = ao_dm0b = aow = None
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
vmata[ia, :, :, :, p0:p1] += ipipa[:, :,
p0:p1].transpose(1, 0, 3, 2)
vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]
vmatb[ia, :, :, :, p0:p1] += ipipb[:, :,
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):
rhoa = ni.eval_rho2(mol, ao[:10], mo_coeff[0], mo_occ[0], mask,
xctype)
rhob = ni.eval_rho2(mol, ao[:10], mo_coeff[1], mo_occ[1], mask,
xctype)
vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]
wva, wvb = numint._uks_mgga_wv0((rhoa, rhob), vxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva)
rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
aow = rks_grad._make_dR_dao_w(ao, wvb)
rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)
aow = [numint._scale_ao(ao[i], wva[5]) for i in range(4, 10)]
rks_hess._d1d2_dot_(ipipa, mol, [aow[0], aow[1], aow[2]],
[ao[XX], ao[XY], ao[XZ]], mask, ao_loc, False)
rks_hess._d1d2_dot_(ipipa, mol, [aow[1], aow[3], aow[4]],
[ao[YX], ao[YY], ao[YZ]], mask, ao_loc, False)
rks_hess._d1d2_dot_(ipipa, mol, [aow[2], aow[4], aow[5]],
[ao[ZX], ao[ZY], ao[ZZ]], mask, ao_loc, False)
aow = [numint._scale_ao(ao[i], wvb[5]) for i in range(4, 10)]
rks_hess._d1d2_dot_(ipipb, mol, [aow[0], aow[1], aow[2]],
[ao[XX], ao[XY], ao[XZ]], mask, ao_loc, False)
rks_hess._d1d2_dot_(ipipb, mol, [aow[1], aow[3], aow[4]],
[ao[YX], ao[YY], ao[YZ]], mask, ao_loc, False)
rks_hess._d1d2_dot_(ipipb, mol, [aow[2], aow[4], aow[5]],
[ao[ZX], ao[ZY], ao[ZZ]], mask, ao_loc, False)
ao_dm0a = [
numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
for i in range(4)
]
ao_dm0b = [
numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
for i in range(4)
]
for ia in range(mol.natm):
wva = dR_rho1a = rks_hess._make_dR_rho1(
ao, ao_dm0a, ia, aoslices, xctype)
wvb = dR_rho1b = rks_hess._make_dR_rho1(
ao, ao_dm0b, ia, aoslices, xctype)
wva[0], wvb[0] = numint._uks_mgga_wv1(
(rhoa, rhob), (dR_rho1a[0], dR_rho1b[0]), vxc, fxc, weight)
wva[1], wvb[1] = numint._uks_mgga_wv1(
(rhoa, rhob), (dR_rho1a[1], dR_rho1b[1]), vxc, fxc, weight)
wva[2], wvb[2] = numint._uks_mgga_wv1(
(rhoa, rhob), (dR_rho1a[2], dR_rho1b[2]), vxc, fxc, weight)
aow = rks_grad._make_dR_dao_w(ao, wva[0])
rks_grad._d1_dot_(vmata[ia, 0], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wva[1])
rks_grad._d1_dot_(vmata[ia, 1], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wva[2])
rks_grad._d1_dot_(vmata[ia, 2], mol, aow, ao[0], mask, ao_loc,
True)
aow = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
aow = rks_grad._make_dR_dao_w(ao, wvb[0])
rks_grad._d1_dot_(vmatb[ia, 0], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wvb[1])
rks_grad._d1_dot_(vmatb[ia, 1], mol, aow, ao[0], mask, ao_loc,
True)
aow = rks_grad._make_dR_dao_w(ao, wvb[2])
rks_grad._d1_dot_(vmatb[ia, 2], mol, aow, ao[0], mask, ao_loc,
True)
aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
False)
# *2 because wv[5] is scaled by 0.5 in _rks_mgga_wv1
wv = wva[:, 5] * 2
aow = [numint._scale_ao(ao[1], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmata[ia], mol, [ao[XX], ao[XY], ao[XZ]],
aow, mask, ao_loc, False)
aow = [numint._scale_ao(ao[2], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmata[ia], mol, [ao[YX], ao[YY], ao[YZ]],
aow, mask, ao_loc, False)
aow = [numint._scale_ao(ao[3], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmata[ia], mol, [ao[ZX], ao[ZY], ao[ZZ]],
aow, mask, ao_loc, False)
wv = wvb[:, 5] * 2
aow = [numint._scale_ao(ao[1], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[XX], ao[XY], ao[XZ]],
aow, mask, ao_loc, False)
aow = [numint._scale_ao(ao[2], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[YX], ao[YY], ao[YZ]],
aow, mask, ao_loc, False)
aow = [numint._scale_ao(ao[3], wv[i]) for i in range(3)]
rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[ZX], ao[ZY], ao[ZZ]],
aow, mask, ao_loc, False)
for ia in range(mol.natm):
p0, p1 = aoslices[ia][2:]
vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
vmata[ia, :, :, :, p0:p1] += ipipa[:, :,
p0:p1].transpose(1, 0, 3, 2)
vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]
vmatb[ia, :, :, :, p0:p1] += ipipb[:, :,
p0:p1].transpose(1, 0, 3, 2)
return vmata, vmatb
