def get_j_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngs = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
ni = numint._KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpt_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpt_band)
dms = dm_kpts.reshape(-1,nkpts,nao,nao)
nset = dms.shape[0]
vjR = [get_vjR(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpt_band is not None:
vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
for i in range(nset)]
if dm_kpts.ndim == 3: # One set of dm_kpts for KRHF
vj_kpts = vj_kpts[0]
return lib.asarray(vj_kpts)
else:
vj_kpts = []
for i in range(nset):
vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
for aoR_k in aoR_kpts]
vj_kpts.append(lib.asarray(vj))
return lib.asarray(vj_kpts).reshape(dm_kpts.shape)
def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpt_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpt_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
vkR_k1k2 = get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2)
ngs, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
#:vk = (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq', dm, aoR_k1.conj(),
#: vkR_k1k2, aoR_k2)
aoR_dm_k2 = np.dot(aoR_k2, dm)
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dm_k2)
vk = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), tmp_Rq)
return vj, vk
if dm.ndim == 2:
vj, vk = contract(dm)
else:
jk = [contract(x) for x in dm.reshape(-1,nao,nao)]
vj = lib.asarray([x[0] for x in jk])
vk = lib.asarray([x[1] for x in jk])
return vj.reshape(dm.shape), vk.reshape(dm.shape)
def get_jk_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngs = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
dms = dm_kpts.reshape(-1,nkpts,nao,nao)
nset = dms.shape[0]
ni = numint._KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpt_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpt_band)
# J
vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpt_band is not None:
vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
for i in range(nset)]
else:
vj_kpts = []
for i in range(nset):
vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
for aoR_k in aoR_kpts]
vj_kpts.append(lib.asarray(vj))
vj_kpts = lib.asarray(vj_kpts)
vjR = None
# K
weight = 1./nkpts * (cell.vol/ngs)
vk_kpts = np.zeros_like(vj_kpts)
if kpt_band is not None:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
vkR_k1k2 = get_vkR(mf, cell, aoR_kband, aoR_kpts[k2],
kpt_band, kpt2)
#:vk_kpts = 1./nkpts * (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq',
#: dm_kpts[k2], aoR_kband.conj(),
#: vkR_k1k2, aoR_kpts[k2])
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i] += weight * lib.dot(aoR_kband.T.conj(), tmp_Rq)
vkR_k1k2 = None
if dm_kpts.ndim == 3:
vj_kpts = vj_kpts[0]
vk_kpts = vk_kpts[0]
return lib.asarray(vj_kpts), lib.asarray(vk_kpts)
else:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
for k1, kpt1 in enumerate(kpts):
vkR_k1k2 = get_vkR(mf, cell, aoR_kpts[k1], aoR_kpts[k2],
kpt1, kpt2)
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i,k1] += weight * lib.dot(aoR_kpts[k1].T.conj(), tmp_Rq)
vkR_k1k2 = None
return vj_kpts.reshape(dm_kpts.shape), vk_kpts.reshape(dm_kpts.shape)
def get_hcore(self, cell=None, kpts=None):
if cell is None: cell = self.cell
if kpts is None: kpts = self.kpts
if cell.pseudo is None:
nuc = self.with_df.get_nuc(kpts)
else:
nuc = self.with_df.get_pp(kpts)
t = cell.pbc_intor('cint1e_kin_sph', 1, 1, kpts)
return lib.asarray(nuc) + lib.asarray(t)
def eig(self, h_kpts, s_kpts):
nkpts = len(h_kpts)
eig_kpts = []
mo_coeff_kpts = []
for k in range(nkpts):
e, c = hf.RHF.eig(self, h_kpts[k], s_kpts[k])
eig_kpts.append(e)
mo_coeff_kpts.append(c)
return lib.asarray(eig_kpts), lib.asarray(mo_coeff_kpts)
def get_hcore(self, cell=None, kpts=None):
if cell is None: cell = self.cell
if kpts is None: kpts = self.kpts
if cell.pseudo:
nuc = lib.asarray(self.with_df.get_pp(kpts))
else:
nuc = lib.asarray(self.with_df.get_nuc(kpts))
if len(cell._ecpbas) > 0:
nuc += lib.asarray(ecp.ecp_int(cell, kpts))
t = lib.asarray(cell.pbc_intor('cint1e_kin_sph', 1, 1, kpts))
return nuc + t
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpt_band=None):
"""Get the Coulomb (J) AO matrix at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
Density matrix at each k-point. If a list of k-point DMs, eg,
UHF alpha and beta DM, the alpha and beta DMs are contracted
separately.
kpts : (nkpts, 3) ndarray
Kwargs:
kpt_band : (3,) ndarray
An arbitrary "band" k-point at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
or list of vj if the input dm_kpts is a list of DMs
"""
cell = mydf.cell
gs = mydf.gs
dm_kpts = lib.asarray(dm_kpts, order="C")
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
coulG = tools.get_coulG(cell, gs=gs)
ngs = len(coulG)
vR = rhoR = np.zeros((nset, ngs))
for k, aoR in mydf.aoR_loop(gs, kpts):
for i in range(nset):
rhoR[i] += numint.eval_rho(cell, aoR, dms[i, k])
for i in range(nset):
rhoR[i] *= 1.0 / nkpts
rhoG = tools.fft(rhoR[i], gs)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, gs).real
if kpt_band is not None:
for k, aoR_kband in mydf.aoR_loop(gs, kpts, kpt_band):
pass
vj_kpts = [cell.vol / ngs * lib.dot(aoR_kband.T.conj() * vR[i], aoR_kband) for i in range(nset)]
if dm_kpts.ndim == 3: # One set of dm_kpts for KRHF
vj_kpts = vj_kpts[0]
return lib.asarray(vj_kpts)
else:
vj_kpts = []
weight = cell.vol / ngs
for k, aoR in mydf.aoR_loop(gs, kpts):
for i in range(nset):
vj_kpts.append(weight * lib.dot(aoR.T.conj() * vR[i], aoR))
vj_kpts = lib.asarray(vj_kpts).reshape(nkpts, nset, nao, nao)
return vj_kpts.transpose(1, 0, 2, 3).reshape(dm_kpts.shape)
def kernel(method, efg_nuc=None):
log = lib.logger.Logger(method.stdout, method.verbose)
cell = method.cell
if efg_nuc is None:
efg_nuc = range(cell.natm)
dm = method.make_rdm1()
if isinstance(method, scf.khf.KSCF):
if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
dm = dm[0] + dm[1] # KUHF density matrix
elif isinstance(method, scf.hf.SCF):
if isinstance(dm[0], numpy.ndarray) and dm[0].ndim == 2:
dm = dm[0] + dm[1] # UHF density matrix
else:
mo = method.mo_coeff
if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
dm_a = [lib.einsum('pi,ij,qj->pq', c, dm[0][k], c.conj())
for k, c in enumerate(mo)]
dm_b = [lib.einsum('pi,ij,qj->pq', c, dm[1][k], c.conj())
for k, c in enumerate(mo)]
dm = lib.asarray(dm_a) + lib.asarray(dm_b)
else:
dm = lib.asarray([lib.einsum('pi,ij,qj->pq', c, dm[k], c.conj())
for k, c in enumerate(mo)])
if isinstance(method, scf.hf.SCF):
with_df = getattr(method, 'with_df', None)
with_x2c = getattr(method, 'with_x2c', None)
else:
with_df = getattr(method._scf, 'with_df', None)
with_x2c = getattr(method._scf, 'with_x2c', None)
if with_x2c:
raise NotImplementedError
log.info('\nElectric Field Gradient Tensor Results')
if isinstance(with_df, df.fft.FFTDF):
efg_e = _fft_quad_integrals(with_df, dm, efg_nuc)
else:
efg_e = _aft_quad_integrals(with_df, dm, efg_nuc)
efg = []
for i, atm_id in enumerate(efg_nuc):
efg_nuc = _get_quad_nuc(cell, atm_id)
v = efg_nuc - efg_e[i]
efg.append(v)
rhf_efg._analyze(cell, atm_id, v, log)
return numpy.asarray(efg)
def vind(mo1):
mo1aa, mo1ab, mo1ba, mo1bb = _split_mo1(mo1)
dm1aa = _dm1_mo2ao(mo1aa, orbva, orboa)
dm1ab = _dm1_mo2ao(mo1ab, orbva, orbob)
dm1ba = _dm1_mo2ao(mo1ba, orbvb, orboa)
dm1bb = _dm1_mo2ao(mo1bb, orbvb, orbob)
dm1 = lib.asarray([dm1aa+dm1aa.transpose(0,2,1),
dm1ab+dm1ba.transpose(0,2,1),
dm1ba+dm1ab.transpose(0,2,1),
dm1bb+dm1bb.transpose(0,2,1)])
v1 = vresp(dm1)
v1aa = _ao2mo.nr_e2(v1[0], mo_va_oa, (0,nvira,nvira,nvira+nocca))
v1ab = _ao2mo.nr_e2(v1[1], mo_va_ob, (0,nvira,nvira,nvira+noccb))
v1ba = _ao2mo.nr_e2(v1[2], mo_vb_oa, (0,nvirb,nvirb,nvirb+nocca))
v1bb = _ao2mo.nr_e2(v1[3], mo_vb_ob, (0,nvirb,nvirb,nvirb+noccb))
v1aa = v1aa.reshape(nset,nvira,nocca)
v1ab = v1ab.reshape(nset,nvira,noccb)
v1ba = v1ba.reshape(nset,nvirb,nocca)
v1bb = v1bb.reshape(nset,nvirb,noccb)
v1aa *= eai_aa
v1ab *= eai_ab
v1ba *= eai_ba
v1bb *= eai_bb
v1mo = numpy.hstack((v1aa.reshape(nset,-1),
v1ab.reshape(nset,-1),
v1ba.reshape(nset,-1),
v1bb.reshape(nset,-1)))
return v1mo.ravel()
def get_fock(
mf,
h1e_kpts,
s_kpts,
vhf_kpts,
dm_kpts,
cycle=-1,
adiis=None,
diis_start_cycle=None,
level_shift_factor=None,
damp_factor=None,
):
if diis_start_cycle is None:
diis_start_cycle = mf.diis_start_cycle
if level_shift_factor is None:
level_shift_factor = mf.level_shift
if damp_factor is None:
damp_factor = mf.damp
if isinstance(level_shift_factor, (tuple, list, np.ndarray)):
shifta, shiftb = level_shift_factor
else:
shifta = shiftb = level_shift_factor
f_kpts = h1e_kpts + vhf_kpts
if adiis and cycle >= diis_start_cycle:
f_kpts = adiis.update(s_kpts, dm_kpts, f_kpts)
if abs(level_shift_factor) > 1e-4:
f_kpts = [hf.level_shift(s, dm_kpts[0, k], f_kpts[0, k], shifta) for k, s in enumerate(s_kpts)] + [
hf.level_shift(s, dm_kpts[1, k], f_kpts[1, k], shiftb) for k, s in enumerate(s_kpts)
]
return lib.asarray(f_kpts)
def h_op(x1):
xtmp = []
p0 = 0
for k in range(nkpts):
nocc = len(occidx[k])
nvir = len(viridx[k])
xtmp.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
p0 += nocc * nvir
x1 = xtmp
dm1 = []
for k in range(nkpts):
d1 = reduce(numpy.dot, (mo_coeff[k][:,viridx[k]], x1[k],
mo_coeff[k][:,occidx[k]].T.conj()))
dm1.append(d1+d1.T.conj())
dm1 = lib.asarray(dm1)
if hyb is None:
v1 = mf.get_veff(cell, dm1)
else:
v1 = mf._numint.nr_rks_fxc(cell, mf.grids, mf.xc, dm0, dm1,
0, 1, rho0, vxc, fxc)
if abs(hyb) < 1e-10:
v1 += mf.get_j(cell, dm1)
else:
vj, vk = mf.get_jk(cell, dm1)
v1 += vj - vk * hyb * .5
x2 = [0] * nkpts
for k in range(nkpts):
x2[k] = numpy.einsum('sp,sq->pq', fvv[k], x1[k].conj()) * 2
x2[k]-= numpy.einsum('sp,rp->rs', foo[k], x1[k].conj()) * 2
x2[k] += reduce(numpy.dot, (mo_coeff[k][:,occidx[k]].T.conj(), v1[k],
mo_coeff[k][:,viridx[k]])).T * 4
return numpy.hstack([x.ravel() for x in x2])
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
if kpts_band is not None:
return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
weight = 1./len(kpts)
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
vG = [0] * nset
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
aoao.reshape(-1,nao,nao)).conj()
vG[i] += rho * coulG[p0:p1]
for i in range(nset):
vG[i] *= weight
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
aoao.reshape(-1,nao,nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts):
vj_kpts = vj_kpts.real.copy()
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
h1e_kpts, s_kpts, vhf_kpts, dm_kpts = h1e, s1e, vhf, dm
if h1e_kpts is None: h1e_kpts = mf.get_hcore()
if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
f_kpts = h1e_kpts + vhf_kpts
if cycle < 0 and diis is None: # Not inside the SCF iteration
return f_kpts
if diis_start_cycle is None:
diis_start_cycle = mf.diis_start_cycle
if level_shift_factor is None:
level_shift_factor = mf.level_shift
if damp_factor is None:
damp_factor = mf.damp
if s_kpts is None: s_kpts = mf.get_ovlp()
if dm_kpts is None: dm_kpts = mf.make_rdm1()
if isinstance(level_shift_factor, (tuple, list, np.ndarray)):
shifta, shiftb = level_shift_factor
else:
shifta = shiftb = level_shift_factor
if diis and cycle >= diis_start_cycle:
f_kpts = diis.update(s_kpts, dm_kpts, f_kpts, mf, h1e_kpts, vhf_kpts)
if abs(level_shift_factor) > 1e-4:
f_kpts =([mol_hf.level_shift(s, dm_kpts[0,k], f_kpts[0,k], shifta)
for k, s in enumerate(s_kpts)],
[mol_hf.level_shift(s, dm_kpts[1,k], f_kpts[1,k], shiftb)
for k, s in enumerate(s_kpts)])
return lib.asarray(f_kpts)
def hop_uhf2ghf(x1):
x1ab = []
x1ba = []
ip = 0
for k in range(nkpts):
nv = nvira[k]
no = noccb[k]
x1ab.append(x1[ip:ip+nv*no].reshape(nv,no))
ip += nv * no
for k in range(nkpts):
nv = nvirb[k]
no = nocca[k]
x1ba.append(x1[ip:ip+nv*no].reshape(nv,no))
ip += nv * no
dm1ab = []
dm1ba = []
for k in range(nkpts):
d1ab = reduce(numpy.dot, (orbva[k], x1ab[k], orbob[k].T.conj()))
d1ba = reduce(numpy.dot, (orbvb[k], x1ba[k], orboa[k].T.conj()))
dm1ab.append(d1ab+d1ba.T.conj())
dm1ba.append(d1ba+d1ab.T.conj())
v1ao = vresp1(lib.asarray([dm1ab,dm1ba]))
x2ab = [0] * nkpts
x2ba = [0] * nkpts
for k in range(nkpts):
x2ab[k] = numpy.einsum('pr,rq->pq', fvva[k], x1ab[k])
x2ab[k]-= numpy.einsum('sq,ps->pq', foob[k], x1ab[k])
x2ba[k] = numpy.einsum('pr,rq->pq', fvvb[k], x1ba[k])
x2ba[k]-= numpy.einsum('qs,ps->pq', fooa[k], x1ba[k])
x2ab[k] += reduce(numpy.dot, (orbva[k].T.conj(), v1ao[0][k], orbob[k]))
x2ba[k] += reduce(numpy.dot, (orbvb[k].T.conj(), v1ao[1][k], orboa[k]))
return numpy.hstack([x.real.ravel() for x in (x2ab+x2ba)])
def eval_mat(self, cell, ao_kpts, weight, rho, vxc,
non0tab=None, xctype='LDA', spin=0, verbose=None):
nkpts = len(ao_kpts)
mat = [eval_mat(cell, ao_kpts[k], weight, rho, vxc,
non0tab, xctype, spin, verbose)
for k in range(nkpts)]
return lib.asarray(mat)
def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpt_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpt_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
ngs, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
return vj
if dm.ndim == 2:
vj = contract(dm)
else:
vj = lib.asarray([contract(x) for x in dm.reshape(-1,nao,nao)])
return vj.reshape(dm.shape)
def get_init_guess(self, cell=None, key='minao'):
if cell is None: cell = self.cell
dm = hf.RHF.get_init_guess(self, cell, key)
if key.lower() == 'chkfile':
dm_kpts = dm
else:
dm_kpts = lib.asarray([dm]*len(self.kpts))
return dm_kpts
def get_hcore(self, cell=None, kpts=None):
if cell is None: cell = self.cell
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
xcell, contr_coeff = self.get_xmol(cell)
with_df = aft.AFTDF(xcell)
c = lib.param.LIGHT_SPEED
assert('1E' in self.approx.upper())
if 'ATOM' in self.approx.upper():
atom_slices = xcell.offset_nr_by_atom()
nao = xcell.nao_nr()
x = numpy.zeros((nao,nao))
vloc = numpy.zeros((nao,nao))
wloc = numpy.zeros((nao,nao))
for ia in range(xcell.natm):
ish0, ish1, p0, p1 = atom_slices[ia]
shls_slice = (ish0, ish1, ish0, ish1)
t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
with xcell.with_rinv_as_nucleus(ia):
z = -xcell.atom_charge(ia)
v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
vloc[p0:p1,p0:p1] = v1
wloc[p0:p1,p0:p1] = w1
x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
else:
raise NotImplementedError
t = xcell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
s = xcell.pbc_intor('int1e_ovlp', 1, lib.HERMITIAN, kpts_lst)
v = with_df.get_nuc(kpts_lst)
#w = get_pnucp(with_df, kpts_lst)
if self.basis is not None:
s22 = s
s21 = pbcgto.intor_cross('int1e_ovlp', xcell, cell, kpts=kpts_lst)
h1_kpts = []
for k in range(len(kpts_lst)):
# The treatment of pnucp local part has huge effects to hcore
#h1 = x2c._get_hcore_fw(t[k], vloc, wloc, s[k], x, c) - vloc + v[k]
#h1 = x2c._get_hcore_fw(t[k], v[k], w[k], s[k], x, c)
h1 = x2c._get_hcore_fw(t[k], v[k], wloc, s[k], x, c)
if self.basis is not None:
c = lib.cho_solve(s22[k], s21[k])
h1 = reduce(numpy.dot, (c.T, h1, c))
if self.xuncontract and contr_coeff is not None:
h1 = reduce(numpy.dot, (contr_coeff.T, h1, contr_coeff))
h1_kpts.append(h1)
if kpts is None or numpy.shape(kpts) == (3,):
h1_kpts = h1_kpts[0]
return lib.asarray(h1_kpts)
def make_rdm1(mo_coeff_kpts, mo_occ_kpts):
'''One particle density matrices for all k-points.
Returns:
dm_kpts : (nkpts, nao, nao) ndarray
'''
nkpts = len(mo_occ_kpts)
dm_kpts = [hf.make_rdm1(mo_coeff_kpts[k], mo_occ_kpts[k])
for k in range(nkpts)]
return lib.asarray(dm_kpts)
def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
kpts=None, kpt_band=None):
'''Coulomb + XC functional for UKS. See pyscf/pbc/dft/uks.py
:func:`get_veff` fore more details.
'''
if cell is None: cell = ks.cell
if dm is None: dm = ks.make_rdm1()
if kpts is None: kpts = ks.kpts
t0 = (time.clock(), time.time())
if ks.grids.coords is None:
ks.grids.build()
small_rho_cutoff = ks.small_rho_cutoff
t0 = logger.timer(ks, 'setting up grids', *t0)
else:
small_rho_cutoff = 0
dm = np.asarray(dm)
nao = dm.shape[-1]
# ndim = 4 : dm.shape = (alpha_beta, nkpts, nao, nao)
ground_state = (dm.ndim == 4 and kpt_band is None)
nkpts = len(kpts)
if hermi == 2: # because rho = 0
n, ks._exc, vx = 0, 0, 0
else:
n, ks._exc, vx = ks._numint.nr_uks(cell, ks.grids, ks.xc, dm, 1,
kpts, kpt_band)
logger.debug(ks, 'nelec by numeric integration = %s', n)
t0 = logger.timer(ks, 'vxc', *t0)
hyb = ks._numint.hybrid_coeff(ks.xc, spin=(cell.spin>0)+1)
if abs(hyb) < 1e-10:
vj = ks.get_j(cell, dm, hermi, kpts, kpt_band)
vhf = lib.asarray([vj[0]+vj[1]] * 2)
else:
vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpt_band)
vhf = pbcuhf._makevhf(vj, vk*hyb)
if ground_state:
ks._exc -= (np.einsum('Kij,Kji', dm[0], vk[0]) +
np.einsum('Kij,Kji', dm[1], vk[1])).real * .5 * hyb * (1./nkpts)
if ground_state:
ks._ecoul = np.einsum('Kij,Kji', dm[0]+dm[1], vj[0]+vj[1]).real * .5 * (1./nkpts)
if small_rho_cutoff > 1e-20 and ground_state:
# Filter grids the first time setup grids
idx = ks._numint.large_rho_indices(cell, dm, ks.grids,
small_rho_cutoff, kpts)
logger.debug(ks, 'Drop grids %d',
ks.grids.weights.size - np.count_nonzero(idx))
ks.grids.coords = np.asarray(ks.grids.coords [idx], order='C')
ks.grids.weights = np.asarray(ks.grids.weights[idx], order='C')
ks._numint.non0tab = None
return vhf + vx
def dot_eri_dm(eri, dm, hermi=0):
"""Compute J, K matrices in terms of the given 2-electron integrals and
density matrix. eri or dm can be complex.
Args:
eri : ndarray
complex integral array with N^4 elements (N is the number of orbitals)
dm : ndarray or list of ndarrays
A density matrix or a list of density matrices
Kwargs:
hermi : int
Whether J, K matrix is hermitian
| 0 : no hermitian or symmetric
| 1 : hermitian
| 2 : anti-hermitian
Returns:
Depending on the given dm, the function returns one J and one K matrix,
or a list of J matrices and a list of K matrices, corresponding to the
input density matrices.
"""
dm = np.asarray(dm)
if np.iscomplexobj(dm) or np.iscomplexobj(eri):
nao = dm.shape[-1]
eri = eri.reshape((nao,) * 4)
def contract(dm):
vj = np.einsum("ijkl,ji->kl", eri, dm)
vk = np.einsum("ijkl,jk->il", eri, dm)
return vj, vk
if isinstance(dm, np.ndarray) and dm.ndim == 2:
vj, vk = contract(dm)
else:
vjk = [contract(dmi) for dmi in dm]
vj = lib.asarray([v[0] for v in vjk]).reshape(dm.shape)
vk = lib.asarray([v[1] for v in vjk]).reshape(dm.shape)
else:
vj, vk = pyscf.scf.hf.dot_eri_dm(eri, dm, hermi)
return vj, vk
def make_h1_fcsd(mol, mo_coeff, mo_occ, atmlst):
'''FC + SD'''
orboa = mo_coeff[0][:,mo_occ[0]> 0]
orbva = mo_coeff[0][:,mo_occ[0]==0]
orbob = mo_coeff[1][:,mo_occ[1]> 0]
orbvb = mo_coeff[1][:,mo_occ[1]==0]
nao = mo_coeff[0].shape[0]
h1aa = []
h1ab = []
h1ba = []
h1bb = []
for ia in atmlst:
h1ao = rhf_ssc._get_integrals_fcsd(mol, ia)
# *.5 due to s = 1/2 * pauli-matrix
h1aa.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbva.conj(), orboa) * .5)
h1ab.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbva.conj(), orbob) * .5)
h1ba.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbvb.conj(), orboa) * .5)
h1bb.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbvb.conj(), orbob) *-.5)
return (lib.asarray(h1aa), lib.asarray(h1ab),
lib.asarray(h1ba), lib.asarray(h1bb))
def get_init_guess(self, cell=None, key='minao'):
if cell is None: cell = self.cell
dm = uhf.UHF.get_init_guess(self, cell, key)
if key.lower() == 'chkfile':
dm_kpts = dm
else:
nao = dm.shape[1]
nkpts = len(self.kpts)
dm_kpts = lib.asarray([dm]*nkpts).reshape(nkpts,2,nao,nao)
dm_kpts = dm_kpts.transpose(1,0,2,3)
dm[1,:] *= .98 # To break spin symmetry
return dm_kpts
def h_op(x1):
p0 = 0
x1a = []
for k in range(nkpts):
nocc = len(occidxa[k])
nvir = len(viridxa[k])
x1a.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
p0 += nocc * nvir
x1b = []
for k in range(nkpts):
nocc = len(occidxb[k])
nvir = len(viridxb[k])
x1b.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
p0 += nocc * nvir
dm1a = []
for k in range(nkpts):
d1 = reduce(numpy.dot, (moa[k][:,viridxa[k]], x1a[k],
moa[k][:,occidxa[k]].T.conj()))
dm1a.append(d1+d1.T.conj())
dm1b = []
for k in range(nkpts):
d1 = reduce(numpy.dot, (mob[k][:,viridxa[k]], x1b[k],
mob[k][:,occidxa[k]].T.conj()))
dm1b.append(d1+d1.T.conj())
dm1 = lib.asarray([dm1a,dm1b])
dm1a = dm1b = None
if hyb is None:
v1 = mf.get_veff(cell, dm1)
else:
v1 = mf._numint.nr_uks_fxc(cell, mf.grids, mf.xc, dm0, dm1,
0, 1, rho0, vxc, fxc)
if abs(hyb) < 1e-10:
vj = mf.get_j(cell, dm1)
v1 += vj[0] + vj[1]
else:
vj, vk = mf.get_jk(cell, dm1)
v1 += vj[0]+vj[1] - vk * hyb * .5
x2a = [0] * nkpts
x2b = [0] * nkpts
for k in range(nkpts):
x2a[k] = numpy.einsum('sp,sq->pq', focka[k][viridxa[k][:,None],viridxa[k]], x1a[k].conj())
x2a[k]-= numpy.einsum('sp,rp->rs', focka[k][occidxa[k][:,None],occidxa[k]], x1a[k].conj())
x2b[k] = numpy.einsum('sp,sq->pq', fockb[k][viridxb[k][:,None],viridxb[k]], x1b[k].conj())
x2b[k]-= numpy.einsum('sp,rp->rs', fockb[k][occidxb[k][:,None],occidxb[k]], x1b[k].conj())
x2a[k] += reduce(numpy.dot, (moa[k][:,occidxa[k]].T.conj(), v1[0][k],
moa[k][:,viridxa[k]])).T
x2b[k] += reduce(numpy.dot, (mob[k][:,occidxb[k]].T.conj(), v1[1][k],
mob[k][:,viridxb[k]])).T
return numpy.hstack([x.ravel() for x in (x2a+x2b)])
def get_hcore(mf, cell=None, kpts=None):
'''Get the core Hamiltonian AO matrices at sampled k-points.
Args:
kpts : (nkpts, 3) ndarray
Returns:
hcore : (nkpts, nao, nao) ndarray
'''
if cell is None: cell = mf.cell
if kpts is None: kpts = mf.kpts
return lib.asarray([pbchf.get_hcore(cell, k) for k in kpts])
def get_ovlp(mf, cell=None, kpts=None):
'''Get the overlap AO matrices at sampled k-points.
Args:
kpts : (nkpts, 3) ndarray
Returns:
ovlp_kpts : (nkpts, nao, nao) ndarray
'''
if cell is None: cell = mf.cell
if kpts is None: kpts = mf.kpts
return lib.asarray(cell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=kpts))
def make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs):
'''Alpha and beta spin one particle density matrices for all k-points.
Returns:
dm_kpts : (2, nkpts, nao, nao) ndarray
'''
nkpts = len(mo_occ_kpts[0])
nao, nmo = mo_coeff_kpts[0][0].shape
def make_dm(mos, occs):
return [np.dot(mos[k]*occs[k], mos[k].T.conj()) for k in range(nkpts)]
dm_kpts =(make_dm(mo_coeff_kpts[0], mo_occ_kpts[0]) +
make_dm(mo_coeff_kpts[1], mo_occ_kpts[1]))
return lib.asarray(dm_kpts).reshape(2,nkpts,nao,nao)
def get_fock(mf, h1e_kpts, s_kpts, vhf_kpts, dm_kpts, cycle=-1, adiis=None,
diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
if diis_start_cycle is None:
diis_start_cycle = mf.diis_start_cycle
if level_shift_factor is None:
level_shift_factor = mf.level_shift
if damp_factor is None:
damp_factor = mf.damp
f_kpts = h1e_kpts + vhf_kpts
if adiis and cycle >= diis_start_cycle:
f_kpts = adiis.update(s_kpts, dm_kpts, f_kpts)
if abs(level_shift_factor) > 1e-4:
f_kpts = [hf.level_shift(s, dm_kpts[k], f_kpts[k], level_shift_factor)
for k, s in enumerate(s_kpts)]
return lib.asarray(f_kpts)
def get_j_for_bands(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
dmsR = dms.real.reshape(nset,nkpts,nao**2)
dmsI = dms.imag.reshape(nset,nkpts,nao**2)
kpt_allow = numpy.zeros(3)
mesh = mydf.mesh
coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
ngrids = len(coulG)
vG = numpy.zeros((nset,ngrids), dtype=numpy.complex128)
max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
dmsC = dms.conj()
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
#:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
for k, aoao in enumerate(aoaoks):
for i in range(nset):
rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
aoao.reshape(-1,nao,nao)).conj()
vG[i,p0:p1] += rho * coulG[p0:p1]
aoao = aoaoks = p0 = p1 = None
weight = 1./len(kpts)
vG *= weight
t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_band,
max_memory=max_memory):
for k, aoao in enumerate(aoaoks):
for i in range(nset):
vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i,p0:p1],
aoao.reshape(-1,nao,nao))
aoao = aoaoks = p0 = p1 = None
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real.copy()
t1 = log.timer_debug1('get_j pass 2', *t1)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(civec_strs, (tuple, list)):
nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
ci_coeff = lib.asarray(civec_strs)
else:
ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
na = len(strsa)
nb = len(strsb)
ci0 = ci_coeff.reshape(-1,na,nb)
civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
civec_a_max = civec_a_max[ci_aidx]
civec_b_max = civec_b_max[ci_bidx]
strsa = strsa[ci_aidx]
strsb = strsb[ci_bidx]
eri = ao2mo.restore(1, eri, norb)
eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)
strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb, nelec[1])
strsa = numpy.append(strsa, strsa_add)
strsb = numpy.append(strsb, strsb_add)
aidx = numpy.argsort(strsa)
bidx = numpy.argsort(strsb)
ci_strs = (strsa[aidx], strsb[bidx])
aidx = numpy.where(aidx < len(ci_aidx))[0]
bidx = numpy.where(bidx < len(ci_bidx))[0]
ma = len(strsa)
mb = len(strsb)
cs = []
for i in range(ci0.shape[0]):
ci1 = numpy.zeros((ma,mb))
tmp = lib.take_2d(ci0[i], ci_aidx, ci_bidx)
lib.takebak_2d(ci1, tmp, aidx, bidx)
cs.append(_as_SCIvector(ci1, ci_strs))
if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
cs = cs[0]
return cs
def eig(self, fock, s):
e_a, c_a = self._eigh(fock[0], s)
e_b, c_b = self._eigh(fock[1], s)
return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
def get_hcore(self, cell=None, kpts=None):
hcore = khf.KSCF.get_hcore(self, cell, kpts)
return lib.asarray([scipy.linalg.block_diag(h, h) for h in hcore])
def eig(self, h_kpts, s_kpts):
e_a, c_a = khf.KRHF.eig(self, h_kpts[0], s_kpts)
e_b, c_b = khf.KRHF.eig(self, h_kpts[1], s_kpts)
return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
def get_veff(ks,
cell=None,
dm=None,
dm_last=0,
vhf_last=0,
hermi=1,
kpt=None,
kpt_band=None):
'''Coulomb + XC functional for UKS. See pyscf/pbc/dft/uks.py
:func:`get_veff` fore more details.
'''
if cell is None: cell = ks.cell
if dm is None: dm = ks.make_rdm1()
if kpt is None: kpt = ks.kpt
t0 = (time.clock(), time.time())
if ks.grids.coords is None:
ks.grids.build(with_non0tab=True)
small_rho_cutoff = ks.small_rho_cutoff
t0 = logger.timer(ks, 'setting up grids', *t0)
else:
small_rho_cutoff = 0
if not isinstance(dm, numpy.ndarray):
dm = numpy.asarray(dm)
if dm.ndim == 2: # RHF DM
dm = numpy.asarray((dm * .5, dm * .5))
if hermi == 2: # because rho = 0
n, ks._exc, vx = (0, 0), 0, 0
else:
n, ks._exc, vx = ks._numint.nr_uks(cell, ks.grids, ks.xc, dm, 0, kpt,
kpt_band)
logger.debug(ks, 'nelec by numeric integration = %s', n)
t0 = logger.timer(ks, 'vxc', *t0)
# ndim = 3 : dm.shape = ([alpha,beta], nao, nao)
ground_state = (dm.ndim == 3 and dm.shape[0] == 2)
hyb = ks._numint.hybrid_coeff(ks.xc, spin=cell.spin)
if abs(hyb) < 1e-10:
vj = ks.get_j(cell, dm, hermi, kpt, kpt_band)
vhf = lib.asarray([vj[0] + vj[1]] * 2)
else:
vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpt_band)
vhf = pbcuhf._makevhf(vj, vk * hyb)
if ground_state:
ks._exc -= (numpy.einsum('ij,ji', dm[0], vk[0]) +
numpy.einsum('ij,ji', dm[1], vk[1])).real * .5 * hyb
if ground_state:
ks._ecoul = numpy.einsum('ij,ji', dm[0] + dm[1],
vj[0] + vj[1]).real * .5
nelec = cell.nelec
if (small_rho_cutoff > 1e-20 and ground_state
and abs(n[0] - nelec[0]) < 0.01 * n[0]
and abs(n[1] - nelec[1]) < 0.01 * n[1]):
# Filter grids the first time setup grids
idx = ks._numint.large_rho_indices(cell, dm, ks.grids,
small_rho_cutoff, kpt)
logger.debug(ks, 'Drop grids %d',
ks.grids.weights.size - numpy.count_nonzero(idx))
ks.grids.coords = numpy.asarray(ks.grids.coords[idx], order='C')
ks.grids.weights = numpy.asarray(ks.grids.weights[idx], order='C')
ks.grids.non0tab = ks.grids.make_mask(cell, ks.grids.coords)
return vhf + vx
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
mesh = mydf.mesh
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
swap_2e = (kpts_band is None)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
kk_table = kpts_band.reshape(-1, 1, 3) - kpts.reshape(1, -1, 3)
kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool)
vkR = numpy.zeros((nset, nband, nao, nao))
vkI = numpy.zeros((nset, nband, nao, nao))
dmsR = numpy.asarray(dms.real, order='C')
dmsI = numpy.asarray(dms.imag, order='C')
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
# K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
def make_kpt(kpt): # kpt = kptj - kpti
# search for all possible ki and kj that has ki-kj+kpt=0
kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
nkptj = len(kptj_idx)
log.debug1('kpt = %s', kpt)
log.debug2('kpti_idx = %s', kpti_idx)
log.debug2('kptj_idx = %s', kptj_idx)
kk_todo[kpti_idx, kptj_idx] = False
if swap_2e and not is_zero(kpt):
kk_todo[kptj_idx, kpti_idx] = False
max_memory1 = max_memory * (nkptj + 1) / (nkptj + 5)
#blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)
#bufR = numpy.empty((blksize*nao**2))
#bufI = numpy.empty((blksize*nao**2))
# Use DF object to mimic KRHF/KUHF object in function get_coulG
mydf.exxdiv = exxdiv
vkcoulG = mydf.weighted_coulG(kpt, True, mesh)
kptjs = kpts[kptj_idx]
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
#buf1R = numpy.empty((blksize*nao**2))
#buf1I = numpy.empty((blksize*nao**2))
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt,
kptjs,
max_memory=max_memory1):
coulG = numpy.sqrt(vkcoulG[p0:p1])
nG = p1 - p0
bufR = numpy.empty((nG * nao**2))
bufI = numpy.empty((nG * nao**2))
buf1R = numpy.empty((nG * nao**2))
buf1I = numpy.empty((nG * nao**2))
for k, aoao in enumerate(aoaoks):
ki = kpti_idx[k]
kj = kptj_idx[k]
# case 1: k_pq = (pi|iq)
#:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,jk->il', v4, dm)
pLqR = numpy.ndarray((nao, nG, nao), buffer=bufR)
pLqI = numpy.ndarray((nao, nG, nao), buffer=bufI)
pLqR[:] = aoao.real.reshape(nG, nao, nao).transpose(1, 0, 2)
pLqI[:] = aoao.imag.reshape(nG, nao, nao).transpose(1, 0, 2)
pLqR *= coulG.reshape(1, nG, 1)
pLqI *= coulG.reshape(1, nG, 1)
iLkR = numpy.ndarray((nao * nG, nao), buffer=buf1R)
iLkI = numpy.ndarray((nao * nG, nao), buffer=buf1I)
for i in range(nset):
iLkR, iLkI = zdotNN(pLqR.reshape(-1, nao),
pLqI.reshape(-1, nao), dmsR[i, kj],
dmsI[i, kj], 1, iLkR, iLkI)
zdotNC(iLkR.reshape(nao, -1), iLkI.reshape(nao, -1),
pLqR.reshape(nao, -1).T,
pLqI.reshape(nao, -1).T, 1, vkR[i, ki], vkI[i, ki],
1)
# case 2: k_pq = (iq|pi)
#:v4 = numpy.einsum('iLj,lLk->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,li->kj', v4, dm)
if swap_2e and not is_zero(kpt):
iLkR = iLkR.reshape(nao, -1)
iLkI = iLkI.reshape(nao, -1)
for i in range(nset):
iLkR, iLkI = zdotNN(dmsR[i, ki], dmsI[i, ki],
pLqR.reshape(nao, -1),
pLqI.reshape(nao, -1), 1, iLkR,
iLkI)
zdotCN(
pLqR.reshape(-1, nao).T,
pLqI.reshape(-1, nao).T, iLkR.reshape(-1, nao),
iLkI.reshape(-1, nao), 1, vkR[i, kj], vkI[i, kj],
1)
for ki, kpti in enumerate(kpts_band):
for kj, kptj in enumerate(kpts):
if kk_todo[ki, kj]:
make_kpt(kptj - kpti)
t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)
if (gamma_point(kpts) and gamma_point(kpts_band)
and not numpy.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
vk_kpts *= 1. / nkpts
# G=0 was not included in the non-uniform grids
if cell.dimension != 3 and exxdiv:
assert (exxdiv.lower() == 'ewald')
_ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def uks_j_xc(mydf,
dm_kpts,
xc_code,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None,
with_j=WITH_J,
j_in_xc=J_IN_XC):
log = lib.logger.Logger(mydf.stdout, mydf.verbose)
cell = mydf.cell
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
dms = None
#TODO: Handle multiple sets of KUKS density matrices (2,nset,nkpts,nao,nao)
assert (nset == 2) # alpha and beta density matrices in KUKS
ni = mydf._numint
xctype = ni._xc_type(xc_code)
if xctype == 'LDA':
deriv = 0
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k])
aow = numpy.einsum('pi,p->pi', ao_l[k], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
add_j_(v, ao_l, ao_h, idx_l, idx_h, wv[:, 0])
elif xctype == 'GGA':
deriv = 1
if RHOG_HIGH_DERIV:
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)
else:
Gv = cell.Gv
ngrids = Gv.shape[0]
rhoG = numpy.empty((2, 4, ngrids), dtype=numpy.complex128)
rhoG[:, :1] = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
rhoG[:, 1:] = numpy.einsum('np,px->nxp', 1j * rhoG[:, 0], Gv)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
wva, wvb = wv
for k in range(nkpts):
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva)
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
v[0, k, idx_l[:, None], idx_h] += v1
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb)
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
v[1, k, idx_l[:, None], idx_h] += v1
else: # MGGA
deriv = 2
#TODO: RHOG_HIGH_DERIV:
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)
def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
for k in range(nkpts):
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
v[0, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
v[1, k, idx_l[:, None],
idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
wva, wvb = wv
for k in range(nkpts):
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva[:4])
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][1], wva[4], out=aow)
v1 += lib.dot(ao_l[k][1].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][2], wva[4], out=aow)
v1 += lib.dot(ao_l[k][2].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][3], wva[4], out=aow)
v1 += lib.dot(ao_l[k][3].conj().T, aow)
v[0, k, idx_l[:, None], idx_h] += v1
aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb[:4])
v1 = lib.dot(aow.conj().T, ao_h[k][0])
aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
v1 += lib.dot(ao_l[k][0].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][1], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][1].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][2], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][2].conj().T, aow)
aow = numpy.einsum('pi,p->pi', ao_h[k][3], wvb[4], out=aow)
v1 += lib.dot(ao_l[k][3].conj().T, aow)
v[1, k, idx_l[:, None], idx_h] += v1
mesh = cell.mesh
coulG = tools.get_coulG(cell,
mesh=mesh,
low_dim_ft_type=mydf.low_dim_ft_type)
ngrids = coulG.size
vG = numpy.einsum('ng,g->ng', rhoG[:, 0].reshape(-1, ngrids), coulG)
vG = vG.reshape(2, *mesh)
weight = cell.vol / ngrids
# *(1./weight) because rhoR is scaled by weight in _eval_rhoG. When
# computing rhoR with IFFT, the weight factor is not needed.
rhoR = tools.ifft(rhoG.reshape(-1, ngrids), mesh) * (1. / weight)
rhoR = rhoR.real.reshape(2, -1, ngrids)
nelec = numpy.zeros(2)
excsum = 0
exc, vxc = ni.eval_xc(xc_code, rhoR, 1, deriv=1)[:2]
if xctype == 'LDA':
vrho = vxc[0]
wva = vrho[:, 0].reshape(1, ngrids)
wvb = vrho[:, 1].reshape(1, ngrids)
elif xctype == 'GGA':
vrho, vsigma = vxc[:2]
wva = numpy.empty((4, ngrids))
wvb = numpy.empty((4, ngrids))
wva[0] = vrho[:, 0]
wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu
wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud
wvb[0] = vrho[:, 1]
wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd
wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud
else:
vrho, vsigma, vlapl, vtau = vxc
wva = numpy.empty((5, ngrids))
wvb = numpy.empty((5, ngrids))
wva[0] = vrho[:, 0]
wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu
wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud
wvb[0] = vrho[:, 1]
wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd
wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud
if vlapl is None:
wvb[4] = .5 * vtau[:, 1]
wva[4] = .5 * vtau[:, 0]
else:
wva[4] = (.5 * vtau[:, 0] + 2 * vlapl[:, 0])
wvb[4] = (.5 * vtau[:, 1] + 2 * vlapl[:, 1])
nelec[0] += rhoR[0, 0].sum() * weight
nelec[1] += rhoR[1, 0].sum() * weight
excsum += (rhoR[0, 0] * exc).sum() * weight
excsum += (rhoR[1, 0] * exc).sum() * weight
wv_freq = tools.fft(numpy.vstack((wva, wvb)), mesh) * weight
wv_freq = wv_freq.reshape(2, -1, *mesh)
if j_in_xc:
wv_freq[:, 0] += vG
vR = tools.ifft(vG.reshape(-1, ngrids), mesh)
ecoul = numpy.einsum('ng,ng->', rhoR[:, 0].real, vR.real) * .5
log.debug('Coulomb energy %s', ecoul)
excsum += ecoul
rhoR = rhoG = None
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
if gamma_point(kpts_band):
veff = numpy.zeros((2, nkpts, nao, nao))
vj = numpy.zeros((2, nkpts, nao, nao))
else:
veff = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)
vj = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)
for grids_high, grids_low in mydf.tasks:
cell_high = grids_high.cell
mesh = grids_high.mesh
coords_idx = grids_high.coords_idx
ngrids0 = numpy.prod(mesh)
ngrids1 = grids_high.coords.shape[0]
log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0)
gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int)
gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int)
gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int)
sub_wvG = wv_freq[:, :, gx[:, None, None], gy[:, None],
gz].reshape(-1, ngrids0)
wv = tools.ifft(sub_wvG, mesh).real.reshape(2, -1, ngrids0)
wv = wv[:, :, coords_idx]
if with_j:
sub_vG = vG[:, gx[:, None, None], gy[:, None],
gz].reshape(-1, ngrids0)
vR = tools.ifft(sub_vG, mesh).real.reshape(2, ngrids0)
vR = vR[:, coords_idx]
idx_h = grids_high.ao_idx
if grids_low is None:
for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv):
ao_h = ao_h_etc[0]
add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
if with_j:
add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
ao_h = ao_h_etc = None
else:
idx_l = grids_low.ao_idx
for ao_h_etc, ao_l_etc in zip(
mydf.aoR_loop(grids_high, kpts, deriv),
mydf.aoR_loop(grids_low, kpts, deriv)):
p0, p1 = ao_h_etc[1:3]
ao_h = ao_h_etc[0][0]
ao_l = ao_l_etc[0][0]
add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
add_xc_(veff, ao_h, ao_l, idx_h, idx_l, wv[:, :, p0:p1])
add_xc_(veff, ao_l, ao_h, idx_l, idx_h, wv[:, :, p0:p1])
if with_j:
add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
add_j_(vj, ao_h, ao_l, idx_h, idx_l, vR[:, p0:p1])
add_j_(vj, ao_l, ao_h, idx_l, idx_h, vR[:, p0:p1])
ao_h = ao_l = ao_h_etc = ao_l_etc = None
vj = _format_jks(vj, dm_kpts, input_band, kpts)
veff = _format_jks(veff, dm_kpts, input_band, kpts)
return nelec, excsum, veff, vj
def rotate_mo(self, mo_coeff, u, log=None):
return lib.asarray(
[numpy.dot(mo, u[k]) for k, mo in enumerate(mo_coeff)])
def get_j_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None):
cell = mydf.cell
mesh = mydf.mesh
ni = mydf._numint
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi=1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
coulG = tools.get_coulG(cell, mesh=mesh)
ngrids = len(coulG)
if gamma_point(kpts):
vR = rhoR = np.zeros((nset, ngrids))
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA')
ao = ao_ks = None
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i, k])
rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1. / nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngrids
vR *= weight
if gamma_point(kpts_band):
vj_kpts = np.zeros((3, nset, nband, nao, nao))
else:
vj_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
rho = None
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band, deriv=1):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao[0], vR[i, p0:p1])
vj_kpts[:, i, k] -= lib.einsum('axi,xj->aij', ao[1:].conj(),
aow)
vj_kpts = np.asarray(
[_format_jks(vj, dm_kpts, input_band, kpts) for vj in vj_kpts])
return vj_kpts
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
kpts : (nkpts, 3) ndarray
Kwargs:
hermi : int
Whether K matrix is hermitian
| 0 : not hermitian and not symmetric
| 1 : hermitian
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts
# Note in the _ewald_exxdiv_for_G0 implementation, the G=0 treatments are
# different for 1D/2D and 3D systems. The special treatments for 1D and 2D
# can only be used with AFTDF/GDF/MDF method. In the FFTDF method, 1D, 2D
# and 3D should use the ewald probe charge correction.
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None):
'''Get the Coulomb (J) AO matrix at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
Density matrix at each k-point. If a list of k-point DMs, eg,
UHF alpha and beta DM, the alpha and beta DMs are contracted
separately.
kpts : (nkpts, 3) ndarray
Kwargs:
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
or list of vj if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
ni = mydf._numint
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
coulG = tools.get_coulG(cell, mesh=mesh)
ngrids = len(coulG)
if hermi == 1 or gamma_point(kpts):
vR = rhoR = np.zeros((nset, ngrids))
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA')
ao = ao_ks = None
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i, k])
rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1. / nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngrids
vR *= weight
if gamma_point(kpts_band):
vj_kpts = np.zeros((nset, nband, nao, nao))
else:
vj_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
rho = None
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao, vR[i, p0:p1])
vj_kpts[i, k] += lib.dot(ao.conj().T, aow)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_k_e1_kpts(mydf,
dm_kpts,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
if input_band is None:
ao2_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1)
]
ao1_kpts = ao2_kpts
ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts]
else:
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
ao1_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C') for ao in
mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao,
max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / 3 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((3, nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('aig,jg->aijg',
ao1T[1:, p0:p1].conj() * expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(3, p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('aijg,jg->aig',
vR,
ao_dms[i],
out=vR_dm[:, i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[:, i, k1] -= weight * np.einsum(
'aig,jg->aij', vR_dm[:, i], ao1T[0])
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Ewald correction has no contribution to nuclear gradient unless range separted Coulomb is used
# The gradient correction part is not added in the vk matrix
if exxdiv == 'ewald' and cell.omega != 0:
raise NotImplementedError("Range Separated Coulomb")
# when cell.omega !=0: madelung constant will have a non-zero derivative
vk_kpts = np.asarray(
[_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts])
return vk_kpts
def get_jk_kpts(mf, cell, dm_kpts, kpts, kpts_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngrids = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
dms = dm_kpts.reshape(-1, nkpts, nao, nao)
nset = dms.shape[0]
ni = numint.KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpts_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpts_band)
# J
vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpts_band is not None:
vj_kpts = [
cell.vol / ngrids * lib.dot(aoR_kband.T.conj() * vjR[i], aoR_kband)
for i in range(nset)
]
else:
vj_kpts = []
for i in range(nset):
vj = [
cell.vol / ngrids * lib.dot(aoR_k.T.conj() * vjR[i], aoR_k)
for aoR_k in aoR_kpts
]
vj_kpts.append(lib.asarray(vj))
vj_kpts = lib.asarray(vj_kpts)
vjR = None
# K
weight = 1. / nkpts * (cell.vol / ngrids)
vk_kpts = np.zeros_like(vj_kpts)
if kpts_band is not None:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i, k2]) for i in range(nset)]
vkR_k1k2 = get_vkR(mf, cell, aoR_kband, aoR_kpts[k2], kpts_band,
kpt2)
#:vk_kpts = 1./nkpts * (cell.vol/ngrids) * np.einsum('rs,Rp,Rqs,Rr->pq',
#: dm_kpts[k2], aoR_kband.conj(),
#: vkR_k1k2, aoR_kpts[k2])
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i] += weight * lib.dot(aoR_kband.T.conj(), tmp_Rq)
vkR_k1k2 = None
if dm_kpts.ndim == 3:
vj_kpts = vj_kpts[0]
vk_kpts = vk_kpts[0]
return lib.asarray(vj_kpts), lib.asarray(vk_kpts)
else:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i, k2]) for i in range(nset)]
for k1, kpt1 in enumerate(kpts):
vkR_k1k2 = get_vkR(mf, cell, aoR_kpts[k1], aoR_kpts[k2], kpt1,
kpt2)
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i, k1] += weight * lib.dot(aoR_kpts[k1].T.conj(),
tmp_Rq)
vkR_k1k2 = None
return vj_kpts.reshape(dm_kpts.shape), vk_kpts.reshape(dm_kpts.shape)
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)),
kpts_band=None, exxdiv=None):
mydf = _sync_mydf(mydf)
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if hasattr(dm_kpts, 'mo_coeff'):
if dm_kpts.ndim == 3: # KRHF
mo_coeff = [dm_kpts.mo_coeff]
mo_occ = [dm_kpts.mo_occ ]
else: # KUHF
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
elif hasattr(dm_kpts[0], 'mo_coeff'):
mo_coeff = [dm.mo_coeff for dm in dm_kpts]
mo_occ = [dm.mo_occ for dm in dm_kpts]
else:
mo_coeff = None
kpts = numpy.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1./nkpts * (cell.vol/ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=dms.dtype)
else:
vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
coords = mydf.grids.coords
ao2_kpts = [numpy.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [numpy.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(min(nao, max(1, (max_memory-mem_now)*1e6/16/4/ngrids/nao)))
lib.logger.debug1(mydf, 'max_memory %s blksize %d', max_memory, blksize)
ao1_dtype = numpy.result_type(*ao1_kpts)
ao2_dtype = numpy.result_type(*ao2_kpts)
vR_dm = numpy.empty((nset,nao,ngrids), dtype=vk_kpts.dtype)
ao_dms_buf = [None] * nkpts
tasks = [(k1,k2) for k2 in range(nkpts) for k1 in range(nband)]
for k1, k2 in mpi.static_partition(tasks):
ao1T = ao1_kpts[k1]
ao2T = ao2_kpts[k2]
kpt1 = kpts_band[k1]
kpt2 = kpts[k2]
if ao2T.size == 0 or ao1T.size == 0:
continue
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2-kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2-kpt1, exxdiv, mydf, mesh)
if is_zero(kpt1-kpt2):
expmikr = numpy.array(1.)
else:
expmikr = numpy.exp(-1j * numpy.dot(coords, kpt2-kpt1))
if ao_dms_buf[k2] is None:
if mo_coeff is None:
ao_dms = [lib.dot(dm[k2], ao2T.conj()) for dm in dms]
else:
ao_dms = []
for i, dm in enumerate(dms):
occ = mo_occ[i][k2]
mo_scaled = mo_coeff[i][k2][:,occ>0] * numpy.sqrt(occ[occ>0])
ao_dms.append(lib.dot(mo_scaled.T, ao2T).conj())
ao_dms_buf[k2] = ao_dms
else:
ao_dms = ao_dms_buf[k2]
if mo_coeff is None:
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = numpy.einsum('ig,jg->ijg', ao1T[p0:p1].conj()*expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1-p0,nao,ngrids)
vG = None
if vR_dm.dtype == numpy.double:
vR = vR.real
for i in range(nset):
numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
vR = None
else:
for p0, p1 in lib.prange(0, nao, blksize):
for i in range(nset):
rho1 = numpy.einsum('ig,jg->ijg',
ao1T[p0:p1].conj()*expmikr,
ao_dms[i].conj())
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1-p0,-1,ngrids)
vG = None
if vR_dm.dtype == numpy.double:
vR = vR.real
numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i,k1] += weight * lib.dot(vR_dm[i], ao1T.T)
vk_kpts = mpi.reduce(lib.asarray(vk_kpts))
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = vk_kpts.real
if rank == 0:
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
kpts : (nkpts, 3) ndarray
Kwargs:
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
gs = mydf.gs
coords = cell.gen_uniform_grids(gs)
ngs = coords.shape[0]
if hasattr(dm_kpts, 'mo_coeff'):
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngs)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
ao2_kpts = mydf._numint.eval_ao(cell, coords, kpts, non0tab=mydf.non0tab)
ao2_kpts = [np.asarray(ao.T, order='C') for ao in ao2_kpts]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = mydf._numint.eval_ao(cell,
coords,
kpts_band,
non0tab=mydf.non0tab)
ao1_kpts = [np.asarray(ao.T, order='C') for ao in ao1_kpts]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
max_memory = mydf.max_memory - lib.current_memory()[0]
blksize = int(min(nao, max_memory * 1e6 / 16 / 2 / ngs / nao + 1))
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((nset, nao, ngs), dtype=vk_kpts.dtype)
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
mydf.exxdiv = exxdiv
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, gs)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
ao2T)
vG = tools.fft(rho1.reshape(-1, ngs), gs)
vG *= coulG
vR = tools.ifft(vG, gs).reshape(p1 - p0, -1, ngs)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), deriv=0):
log = lib.logger.Logger(mydf.stdout, mydf.verbose)
cell = mydf.cell
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
tasks = getattr(mydf, 'tasks', None)
if tasks is None:
mydf.tasks = tasks = multi_grids_tasks(cell, log)
log.debug('Multigrid ntasks %s', len(tasks))
assert (deriv <= 2)
if abs(dms - dms.transpose(0, 1, 3, 2).conj()).max() < 1e-9:
def dot_bra(bra, aodm):
rho = numpy.einsum('pi,pi->p', bra.real, aodm.real)
if aodm.dtype == numpy.complex:
rho += numpy.einsum('pi,pi->p', bra.imag, aodm.imag)
return rho
if deriv == 0:
xctype = 'LDA'
rhodim = 1
def make_rho(ao_l, ao_h, dm_lh, dm_hl):
c0 = lib.dot(ao_l, dm_lh)
rho = dot_bra(ao_h, c0)
return rho * 2
elif deriv == 1:
xctype = 'GGA'
rhodim = 4
def make_rho(ao_l, ao_h, dm_lh, dm_hl):
ngrids = ao_l[0].shape[0]
rho = numpy.empty((4, ngrids))
c0 = lib.dot(ao_l[0], dm_lh)
rho[0] = dot_bra(ao_h[0], c0)
for i in range(1, 4):
rho[i] = dot_bra(ao_h[i], c0)
c0 = lib.dot(ao_h[0], dm_hl)
for i in range(1, 4):
rho[i] += dot_bra(ao_l[i], c0)
return rho * 2 # *2 for dm_lh+dm_hl.T
elif deriv == 2:
xctype = 'MGGA'
rhodim = 6
def make_rho(ao_l, ao_h, dm_lh, dm_hl):
ngrids = ao_l[0].shape[0]
rho = numpy.empty((6, ngrids))
c = [lib.dot(ao_l[i], dm_lh) for i in range(4)]
rho[0] = dot_bra(ao_h[0], c[0])
rho[5] = 0
for i in range(1, 4):
rho[i] = dot_bra(ao_h[i], c[0])
rho[i] += dot_bra(ao_h[0], c[i])
rho[5] += dot_bra(ao_h[i], c[i]) * 2
XX, YY, ZZ = 4, 7, 9
ao2 = ao_h[XX] + ao_h[YY] + ao_h[ZZ]
rho[4] = dot_bra(ao2, c[0])
ao2 = lib.dot(ao_l[XX] + ao_l[YY] + ao_l[ZZ], dm_lh)
rho[4] += dot_bra(ao2, ao_h[0])
rho[4] += rho[5] * 2
rho[5] *= .5
return rho * 2 # *2 for dm_lh+dm_hl.T
else:
raise NotImplementedError('Non-hermitian density matrices')
ni = mydf._numint
nx, ny, nz = cell.mesh
rhoG = numpy.zeros((nset * rhodim, nx, ny, nz), dtype=numpy.complex)
for grids_high, grids_low in tasks:
cell_high = grids_high.cell
mesh = grids_high.mesh
coords_idx = grids_high.coords_idx
ngrids0 = numpy.prod(mesh)
ngrids1 = grids_high.coords.shape[0]
log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0)
idx_h = grids_high.ao_idx
dms_hh = numpy.asarray(dms[:, :, idx_h[:, None], idx_h], order='C')
if grids_low is not None:
idx_l = grids_low.ao_idx
dms_hl = numpy.asarray(dms[:, :, idx_h[:, None], idx_l], order='C')
dms_lh = numpy.asarray(dms[:, :, idx_l[:, None], idx_h], order='C')
rho = numpy.zeros((nset, rhodim, ngrids1))
if grids_low is None:
for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv):
ao_h, mask = ao_h_etc[0], ao_h_etc[2]
for k in range(nkpts):
for i in range(nset):
rho_sub = numint.eval_rho(cell_high, ao_h[k],
dms_hh[i, k], mask, xctype,
hermi)
rho[i, :, p0:p1] += rho_sub.real
ao_h = ao_h_etc = None
else:
for ao_h_etc, ao_l_etc in zip(
mydf.aoR_loop(grids_high, kpts, deriv),
mydf.aoR_loop(grids_low, kpts, deriv)):
p0, p1 = ao_h_etc[1:3]
ao_h, mask = ao_h_etc[0][0], ao_h_etc[0][2]
ao_l = ao_l_etc[0][0]
for k in range(nkpts):
for i in range(nset):
rho_sub = numint.eval_rho(cell_high, ao_h[k],
dms_hh[i, k], mask, xctype,
hermi)
rho[i, :, p0:p1] += rho_sub.real
rho_sub = make_rho(ao_l[k], ao_h[k], dms_lh[i, k],
dms_hl[i, k])
rho[i, :, p0:p1] += rho_sub.real
ao_h = ao_l = ao_h_etc = ao_l_etc = None
rho *= 1. / nkpts
rhoR = numpy.zeros((nset * rhodim, ngrids0))
rhoR[:, coords_idx] = rho.reshape(nset * rhodim, ngrids1)
gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int)
gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int)
gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int)
rho_freq = tools.fft(rhoR, mesh) * cell.vol / ngrids0
for i in range(nset * rhodim):
rhoG[i, gx[:, None, None], gy[:, None],
gz] += rho_freq[i].reshape(mesh)
return rhoG.reshape(nset, rhodim, ngrids0)
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (time.clock(), time.time())
mesh = mydf.mesh
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
swap_2e = (kpts_band is None)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
kk_table = kpts_band.reshape(-1, 1, 3) - kpts.reshape(1, -1, 3)
kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool)
vkR = numpy.zeros((nset, nband, nao, nao))
vkI = numpy.zeros((nset, nband, nao, nao))
dmsR = numpy.asarray(dms.real, order='C')
dmsI = numpy.asarray(dms.imag, order='C')
mem_now = lib.current_memory()[0]
max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
# K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
def make_kpt(kpt): # kpt = kptj - kpti
# search for all possible ki and kj that has ki-kj+kpt=0
kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
nkptj = len(kptj_idx)
log.debug1('kpt = %s', kpt)
log.debug2('kpti_idx = %s', kpti_idx)
log.debug2('kptj_idx = %s', kptj_idx)
kk_todo[kpti_idx, kptj_idx] = False
if swap_2e and not is_zero(kpt):
kk_todo[kptj_idx, kpti_idx] = False
max_memory1 = max_memory * (nkptj + 1) / (nkptj + 5)
#blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)
#bufR = numpy.empty((blksize*nao**2))
#bufI = numpy.empty((blksize*nao**2))
# Use DF object to mimic KRHF/KUHF object in function get_coulG
vkcoulG = mydf.weighted_coulG(kpt, exxdiv, mesh)
kptjs = kpts[kptj_idx]
weight = 1. / len(kpts)
perm_sym = swap_2e and not is_zero(kpt)
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt,
kptjs,
max_memory=max_memory1):
_update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), vkcoulG[p0:p1],
weight, kpti_idx, kptj_idx, perm_sym)
for ki, kpti in enumerate(kpts_band):
for kj, kptj in enumerate(kpts):
if kk_todo[ki, kj]:
make_kpt(kptj - kpti)
t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)
if (gamma_point(kpts) and gamma_point(kpts_band)
and not numpy.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
# Add ewald_exxdiv contribution because G=0 was not included in the
# non-uniform grids
if (exxdiv == 'ewald' and
(cell.dimension < 2 or # 0D and 1D are computed with inf_vacuum
(cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'))):
_ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_pp(mydf, kpts=None):
mydf = _sync_mydf(mydf)
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
if abs(kpts_lst).sum < 1e-9:
dtype = numpy.float64
else:
dtype = numpy.complex128
gs = mydf.gs
SI = cell.get_SI()
Gv = cell.get_Gv(gs)
vpplocG = pseudo.get_vlocG(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', SI, vpplocG)
vpplocG[0] = numpy.sum(pseudo.get_alphas(cell)) # from get_jvloc_G0 function
ngs = len(vpplocG)
nao = cell.nao_nr()
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpp = [lib.dot(aoR.T.conj()*vpplocR, aoR)
for k, aoR in mydf.mpi_aoR_loop(gs, kpts_lst)]
vpp = mpi.gather(lib.asarray(vpp, dtype=dtype))
# vppnonloc evaluated in reciprocal space
fakemol = gto.Mole()
fakemol._atm = numpy.zeros((1,gto.ATM_SLOTS), dtype=numpy.int32)
fakemol._bas = numpy.zeros((1,gto.BAS_SLOTS), dtype=numpy.int32)
ptr = gto.PTR_ENV_START
fakemol._env = numpy.zeros(ptr+10)
fakemol._bas[0,gto.NPRIM_OF ] = 1
fakemol._bas[0,gto.NCTR_OF ] = 1
fakemol._bas[0,gto.PTR_EXP ] = ptr+3
fakemol._bas[0,gto.PTR_COEFF] = ptr+4
# buf for SPG_lmi upto l=0..3 and nl=3
buf = numpy.empty((48,ngs), dtype=numpy.complex128)
def vppnl_by_k(kpt):
Gk = Gv + kpt
G_rad = lib.norm(Gk, axis=1)
aokG = ft_ao.ft_ao(cell, Gv, kpt=kpt) * (ngs/cell.vol)
vppnl = 0
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
p1 = 0
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
fakemol._bas[0,gto.ANG_OF] = l
fakemol._env[ptr+3] = .5*rl**2
fakemol._env[ptr+4] = rl**(l+1.5)*numpy.pi**1.25
pYlm_part = dft.numint.eval_ao(fakemol, Gk, deriv=0)
p0, p1 = p1, p1+nl*(l*2+1)
# pYlm is real, SI[ia] is complex
pYlm = numpy.ndarray((nl,l*2+1,ngs), dtype=numpy.complex128, buffer=buf[p0:p1])
for k in range(nl):
qkl = pseudo.pp._qli(G_rad*rl, l, k)
pYlm[k] = pYlm_part.T * qkl
#:SPG_lmi = numpy.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
#:SPG_lm_aoG = numpy.einsum('nmg,gp->nmp', SPG_lmi, aokG)
#:tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
#:vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
SPG_lmi = buf[:p1]
SPG_lmi *= SI[ia].conj()
SPG_lm_aoGs = lib.zdot(SPG_lmi, aokG)
p1 = 0
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
p0, p1 = p1, p1+nl*(l*2+1)
hl = numpy.asarray(hl)
SPG_lm_aoG = SPG_lm_aoGs[p0:p1].reshape(nl,l*2+1,-1)
tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
return vppnl * (1./ngs**2)
vppnl = []
for kpt in mpi.static_partition(kpts_lst):
vppnl.append(vppnl_by_k(kpt))
vppnl = mpi.gather(lib.asarray(vppnl, dtype=dtype))
if rank == 0:
vpp += vppnl
if kpts is None or numpy.shape(kpts) == (3,):
vpp = vpp[0]
return vpp
def eig(self, fock, s):
e_a, c_a = hf.SCF.eig(self, fock[0], s)
e_b, c_b = hf.SCF.eig(self, fock[1], s)
return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
def get_ovlp(self, cell=None, kpts=None):
s = khf.KSCF.get_ovlp(self, cell, kpts)
return lib.asarray([scipy.linalg.block_diag(x, x) for x in s])
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
cell = mydf.cell
log = logger.Logger(mydf.stdout, mydf.verbose)
if exxdiv is not None and exxdiv != 'ewald':
log.warn(
'GDF does not support exxdiv %s. '
'exxdiv needs to be "ewald" or None', exxdiv)
raise RuntimeError('GDF does not support exxdiv %s' % exxdiv)
t1 = (logger.process_clock(), logger.perf_counter())
if mydf._cderi is None or not mydf.has_kpts(kpts_band):
if mydf._cderi is not None:
log.warn(
'DF integrals for band k-points were not found %s. '
'DF integrals will be rebuilt to include band k-points.',
mydf._cderi)
mydf.build(kpts_band=kpts_band)
t1 = log.timer_debug1('Init get_k_kpts', *t1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
vkR = numpy.zeros((nset, nband, nao, nao))
vkI = numpy.zeros((nset, nband, nao, nao))
dmsR = numpy.asarray(dms.real, order='C')
dmsI = numpy.asarray(dms.imag, order='C')
# K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
bufR = numpy.empty((mydf.blockdim * nao**2))
bufI = numpy.empty((mydf.blockdim * nao**2))
max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
def make_kpt(ki, kj, swap_2e):
kpti = kpts[ki]
kptj = kpts_band[kj]
for LpqR, LpqI, sign in mydf.sr_loop((kpti, kptj), max_memory, False):
nrow = LpqR.shape[0]
pLqR = numpy.ndarray((nao, nrow, nao), buffer=bufR)
pLqI = numpy.ndarray((nao, nrow, nao), buffer=bufI)
tmpR = numpy.ndarray((nao, nrow * nao), buffer=LpqR)
tmpI = numpy.ndarray((nao, nrow * nao), buffer=LpqI)
pLqR[:] = LpqR.reshape(-1, nao, nao).transpose(1, 0, 2)
pLqI[:] = LpqI.reshape(-1, nao, nao).transpose(1, 0, 2)
for i in range(nset):
zdotNN(dmsR[i, ki], dmsI[i, ki], pLqR.reshape(nao, -1),
pLqI.reshape(nao, -1), 1, tmpR, tmpI)
zdotCN(
pLqR.reshape(-1, nao).T,
pLqI.reshape(-1, nao).T, tmpR.reshape(-1, nao),
tmpI.reshape(-1, nao), sign, vkR[i, kj], vkI[i, kj], 1)
if swap_2e:
tmpR = tmpR.reshape(nao * nrow, nao)
tmpI = tmpI.reshape(nao * nrow, nao)
for i in range(nset):
zdotNN(pLqR.reshape(-1, nao), pLqI.reshape(-1, nao),
dmsR[i, kj], dmsI[i, kj], 1, tmpR, tmpI)
zdotNC(tmpR.reshape(nao, -1), tmpI.reshape(nao, -1),
pLqR.reshape(nao, -1).T,
pLqI.reshape(nao, -1).T, sign, vkR[i, ki],
vkI[i, ki], 1)
if kpts_band is kpts: # normal k-points HF/DFT
for ki in range(nkpts):
for kj in range(ki):
make_kpt(ki, kj, True)
make_kpt(ki, ki, False)
t1 = log.timer_debug1('get_k_kpts: make_kpt ki>=kj (%d,*)' % ki,
*t1)
else:
for ki in range(nkpts):
for kj in range(nband):
make_kpt(ki, kj, False)
t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)
if (gamma_point(kpts) and gamma_point(kpts_band)
and not numpy.iscomplexobj(dm_kpts)):
vk_kpts = vkR
else:
vk_kpts = vkR + vkI * 1j
vk_kpts *= 1. / nkpts
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_j_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.zeros((1, 3)),
kpts_band=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = (logger.process_clock(), logger.perf_counter())
if mydf._cderi is None or not mydf.has_kpts(kpts_band):
if mydf._cderi is not None:
log.warn(
'DF integrals for band k-points were not found %s. '
'DF integrals will be rebuilt to include band k-points.',
mydf._cderi)
mydf.build(kpts_band=kpts_band)
t1 = log.timer_debug1('Init get_j_kpts', *t1)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
if mydf.auxcell is None:
# If mydf._cderi is the file that generated from another calculation,
# guess naux based on the contents of the integral file.
naux = mydf.get_naoaux()
else:
naux = mydf.auxcell.nao_nr()
nao_pair = nao * (nao + 1) // 2
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
j_real = gamma_point(kpts_band) and not numpy.iscomplexobj(dms)
dmsR = dms.real.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2)
dmsI = dms.imag.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2)
rhoR = numpy.zeros((nset, naux))
rhoI = numpy.zeros((nset, naux))
max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
for k, kpt in enumerate(kpts):
kptii = numpy.asarray((kpt, kpt))
p1 = 0
for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, False):
p0, p1 = p1, p1 + LpqR.shape[0]
#:Lpq = (LpqR + LpqI*1j).reshape(-1,nao,nao)
#:rhoR[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).real
#:rhoI[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).imag
rhoR[:,
p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsR[:, k])
rhoI[:,
p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsI[:, k])
if LpqI is not None:
rhoR[:, p0:p1] -= sign * numpy.einsum('Lp,xp->xL', LpqI,
dmsI[:, k])
rhoI[:, p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqI,
dmsR[:, k])
LpqR = LpqI = None
t1 = log.timer_debug1('get_j pass 1', *t1)
weight = 1. / nkpts
rhoR *= weight
rhoI *= weight
vjR = numpy.zeros((nset, nband, nao_pair))
vjI = numpy.zeros((nset, nband, nao_pair))
for k, kpt in enumerate(kpts_band):
kptii = numpy.asarray((kpt, kpt))
p1 = 0
for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, True):
p0, p1 = p1, p1 + LpqR.shape[0]
#:Lpq = (LpqR + LpqI*1j)#.reshape(-1,nao,nao)
#:vjR[:,k] += numpy.dot(rho[:,p0:p1], Lpq).real
#:vjI[:,k] += numpy.dot(rho[:,p0:p1], Lpq).imag
vjR[:, k] += numpy.dot(rhoR[:, p0:p1], LpqR)
if not j_real:
vjI[:, k] += numpy.dot(rhoI[:, p0:p1], LpqR)
if LpqI is not None:
vjR[:, k] -= numpy.dot(rhoI[:, p0:p1], LpqI)
vjI[:, k] += numpy.dot(rhoR[:, p0:p1], LpqI)
LpqR = LpqI = None
t1 = log.timer_debug1('get_j pass 2', *t1)
if j_real:
vj_kpts = vjR
else:
vj_kpts = vjR + vjI * 1j
vj_kpts = lib.unpack_tril(vj_kpts.reshape(-1, nao_pair))
vj_kpts = vj_kpts.reshape(nset, nband, nao, nao)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)