def mpi_aoR_loop(self, gs=None, kpts=None, kpts_band=None):
cell = self.cell
if kpts is None: kpts = self.kpts
kpts = numpy.asarray(kpts)
if gs is None:
gs = numpy.asarray(self.gs)
else:
gs = numpy.asarray(gs)
if any(gs != self.gs):
self.non0tab = None
self.gs = gs
ngrids = numpy.prod(gs*2+1)
ni = self._numint
coords = cell.gen_uniform_grids(gs)
if self.non0tab is None:
self.non0tab = ni.make_mask(cell, coords)
if kpts_band is None:
aoR = ni.eval_ao(cell, coords, kpts, non0tab=self.non0tab)
for k in mpi.static_partition(range(len(kpts))):
yield k, aoR[k]
else:
aoR = ni.eval_ao(cell, coords, kpts_band, non0tab=self.non0tab)
if kpts_band.ndim == 1:
if rank == 0:
yield 0, aoR
else:
return
else:
for k in mpi.static_partition(range(len(kpts_band))):
yield k, aoR[k]
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.
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 = 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_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
mesh = mydf.mesh
SI = cell.get_SI()
Gv = cell.get_Gv(mesh)
vpplocG = pseudo.get_vlocG(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', SI, vpplocG)
ngrids = len(vpplocG)
nao = cell.nao_nr()
nkpts = len(kpts_lst)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, mesh).real
vpp = numpy.zeros((nkpts, nao, nao), dtype=dtype)
for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts_lst):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
vpp[k] += lib.dot(ao.T.conj() * vpplocR[p0:p1], ao)
ao = ao_ks = None
vpp = mpi.reduce(lib.asarray(vpp))
# 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, ngrids), 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) * (ngrids / 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, ngrids),
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)
if p1 > 0:
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. / ngrids**2)
vppnl = []
for kpt in mpi.static_partition(kpts_lst):
vppnl.append(vppnl_by_k(kpt))
vppnl = mpi.gather(lib.asarray(vppnl))
if rank == 0:
for k in range(nkpts):
if dtype == numpy.float64:
vpp[k] += vppnl[k].real
else:
vpp[k] += vppnl[k]
if kpts is None or numpy.shape(kpts) == (3, ):
vpp = vpp[0]
return vpp
