def eval_ao(cell, coords, kpt=None, deriv=0, relativity=0, bastart=0,
bascount=None, non0tab=None, verbose=None):
'''Collocate AO crystal orbitals (opt. gradients) on the real-space grid.
Args:
cell : instance of :class:`Cell`
coords : (nx*ny*nz, 3) ndarray
The real-space grid point coordinates.
Kwargs:
kpt : (3,) ndarray
The k-point corresponding to the crystal AO.
deriv : int
AO derivative order. It affects the shape of the return array.
If deriv=0, the returned AO values are stored in a (N,nao) array.
Otherwise the AO values are stored in an array of shape (M,N,nao).
Here N is the number of grids, nao is the number of AO functions,
M is the size associated to the derivative deriv.
Returns:
aoR : ([4,] nx*ny*nz, nao=cell.nao_nr()) ndarray
The value of the AO crystal orbitals on the real-space grid by default.
If deriv=1, also contains the value of the orbitals gradient in the
x, y, and z directions. It can be either complex or float array,
depending on the kpt argument. If kpt is not given (gamma point),
aoR is a float array.
See Also:
pyscf.dft.numint.eval_ao
'''
aoR = 0
for L in tools.get_lattice_Ls(cell, cell.nimgs):
if kpt is None:
aoR += pyscf.dft.numint.eval_ao(cell, coords-L, deriv, relativity,
bastart, bascount,
non0tab, verbose)
else:
factor = numpy.exp(1j*numpy.dot(kpt,L))
aoR += pyscf.dft.numint.eval_ao(cell, coords-L, deriv, relativity,
bastart, bascount,
non0tab, verbose) * factor
if cell.ke_cutoff is not None:
ke = 0.5*numpy.einsum('gi,gi->g', cell.Gv, cell.Gv)
ke_mask = ke < cell.ke_cutoff
aoG = numpy.zeros_like(aoR)
for i in range(cell.nao_nr()):
if deriv == 1:
for c in range(4):
aoG[c][ke_mask, i] = tools.fft(aoR[c][:,i], cell.gs)[ke_mask]
aoR[c][:,i] = tools.ifft(aoG[c][:,i], cell.gs)
else:
aoG[ke_mask, i] = tools.fft(aoR[:,i], cell.gs)[ke_mask]
aoR[:,i] = tools.ifft(aoG[:,i], cell.gs)
return numpy.asarray(aoR)
def get_nuc(mydf, kpts):
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: # gamma_point
dtype = numpy.float64
else:
dtype = numpy.complex128
gs = mydf.gs
charge = -cell.atom_charges()
Gv = cell.get_Gv(gs)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell, gs=gs, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, mydf.gs).real
vne = [lib.dot(aoR.T.conj()*vneR, aoR)
for k, aoR in mydf.mpi_aoR_loop(gs, kpts_lst)]
vne = mpi.gather(lib.asarray(vne, dtype=dtype))
if rank == 0:
if kpts is None or numpy.shape(kpts) == (3,):
vne = vne[0]
return vne
def test_pnucp(self):
cell1 = gto.Cell()
cell1.atom = '''
He 1.3 .2 .3
He .1 .1 1.1 '''
cell1.basis = {'He': [[0, [0.8, 1]],
[1, [0.6, 1]]
]}
cell1.mesh = [15]*3
cell1.a = numpy.array(([2.0, .9, 0. ],
[0.1, 1.9, 0.4],
[0.8, 0 , 2.1]))
cell1.build()
charge = -cell1.atom_charges()
Gv = cell1.get_Gv(cell1.mesh)
SI = cell1.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell1, mesh=cell1.mesh, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, cell1.mesh).real
coords = cell1.gen_uniform_grids(cell1.mesh)
aoR = dft.numint.eval_ao(cell1, coords, deriv=1)
ngrids, nao = aoR.shape[1:]
vne_ref = numpy.einsum('p,xpi,xpj->ij', vneR, aoR[1:4], aoR[1:4])
mydf = df.AFTDF(cell1)
dat = sfx2c1e.get_pnucp(mydf)
self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7)
mydf.eta = 0
dat = sfx2c1e.get_pnucp(mydf)
self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7)
def get_nuc(mydf, kpts):
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: # gamma_point
dtype = numpy.float64
else:
dtype = numpy.complex128
mesh = mydf.mesh
charge = -cell.atom_charges()
Gv = cell.get_Gv(mesh)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, mydf.mesh).real
vne = [0] * len(kpts_lst)
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):
vne[k] += lib.dot(ao.T.conj() * vneR[p0:p1], ao)
ao = ao_ks = None
vne = mpi.reduce(lib.asarray(vne))
if rank == 0:
if kpts is None or numpy.shape(kpts) == (3, ):
vne = vne[0]
return vne
def _contract_compact(mydf, mos, coulG, max_memory):
cell = mydf.cell
moiT, mokT = mos
nmoi, ngrids = moiT.shape
nmok = mokT.shape[0]
wcoulG = coulG * (cell.vol/ngrids)
def fill_orbital_pair(moT, i0, i1, buf):
npair = i1*(i1+1)//2 - i0*(i0+1)//2
out = numpy.ndarray((npair,ngrids), dtype=buf.dtype, buffer=buf)
ij = 0
for i in range(i0, i1):
numpy.einsum('p,jp->jp', moT[i], moT[:i+1], out=out[ij:ij+i+1])
ij += i + 1
return out
eri = numpy.empty((nmoi*(nmoi+1)//2,nmok*(nmok+1)//2))
blksize = int(min(max(nmoi*(nmoi+1)//2, nmok*(nmok+1)//2),
(max_memory*1e6/8 - eri.size)/2/ngrids+1))
buf = numpy.empty((blksize,ngrids))
for p0, p1 in lib.prange_tril(0, nmoi, blksize):
mo_pairs_G = tools.fft(fill_orbital_pair(moiT, p0, p1, buf), mydf.mesh)
mo_pairs_G*= wcoulG
v = tools.ifft(mo_pairs_G, mydf.mesh)
vR = numpy.asarray(v.real, order='C')
for q0, q1 in lib.prange_tril(0, nmok, blksize):
mo_pairs = numpy.asarray(fill_orbital_pair(mokT, q0, q1, buf), order='C')
eri[p0*(p0+1)//2:p1*(p1+1)//2,
q0*(q0+1)//2:q1*(q1+1)//2] = lib.ddot(vR, mo_pairs.T)
v = None
return eri
def test_pnucp(self):
cell1 = gto.Cell()
cell1.atom = '''
He 1.3 .2 .3
He .1 .1 1.1 '''
cell1.basis = {'He': [[0, [0.8, 1]], [1, [0.6, 1]]]}
cell1.mesh = [15] * 3
cell1.a = numpy.array(([2.0, .9, 0.], [0.1, 1.9, 0.4], [0.8, 0, 2.1]))
cell1.build()
charge = -cell1.atom_charges()
Gv = cell1.get_Gv(cell1.mesh)
SI = cell1.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell1, mesh=cell1.mesh, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, cell1.mesh).real
coords = cell1.gen_uniform_grids(cell1.mesh)
aoR = dft.numint.eval_ao(cell1, coords, deriv=1)
ngrids, nao = aoR.shape[1:]
vne_ref = numpy.einsum('p,xpi,xpj->ij', vneR, aoR[1:4], aoR[1:4])
mydf = df.AFTDF(cell1)
dat = sfx2c1e.get_pnucp(mydf)
self.assertAlmostEqual(abs(dat - vne_ref).max(), 0, 7)
mydf.eta = 0
dat = sfx2c1e.get_pnucp(mydf)
self.assertAlmostEqual(abs(dat - vne_ref).max(), 0, 7)
def _contract_plain(mydf, mos, coulG, phase, max_memory):
cell = mydf.cell
moiT, mojT, mokT, molT = mos
nmoi, nmoj, nmok, nmol = [x.shape[0] for x in mos]
ngrids = moiT.shape[1]
wcoulG = coulG * (cell.vol/ngrids)
dtype = numpy.result_type(phase, *mos)
eri = numpy.empty((nmoi*nmoj,nmok*nmol), dtype=dtype)
blksize = int(min(max(nmoi,nmok), (max_memory*1e6/16 - eri.size)/2/ngrids/max(nmoj,nmol)+1))
assert blksize > 0
buf0 = numpy.empty((blksize,max(nmoj,nmol),ngrids), dtype=dtype)
buf1 = numpy.ndarray((blksize,nmoj,ngrids), dtype=dtype, buffer=buf0)
buf2 = numpy.ndarray((blksize,nmol,ngrids), dtype=dtype, buffer=buf0)
for p0, p1 in lib.prange(0, nmoi, blksize):
mo_pairs = numpy.einsum('ig,jg->ijg', moiT[p0:p1].conj()*phase,
mojT, out=buf1[:p1-p0])
mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
mo_pairs = None
mo_pairs_G*= wcoulG
v = tools.ifft(mo_pairs_G, mydf.mesh)
mo_pairs_G = None
v *= phase.conj()
if dtype == numpy.double:
v = numpy.asarray(v.real, order='C')
for q0, q1 in lib.prange(0, nmok, blksize):
mo_pairs = numpy.einsum('ig,jg->ijg', mokT[q0:q1].conj(),
molT, out=buf2[:q1-q0])
eri[p0*nmoj:p1*nmoj,q0*nmol:q1*nmol] = lib.dot(v, mo_pairs.reshape(-1,ngrids).T)
v = None
return eri
def get_nuc(mydf, kpts=None):
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
cell = mydf.cell
mesh = mydf.mesh
charge = -cell.atom_charges()
Gv = cell.get_Gv(mesh)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, mesh).real
vne = [0] * len(kpts_lst)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_lst):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
vne[k] += lib.dot(ao.T.conj() * vneR[p0:p1], ao)
ao = ao_ks = None
if kpts is None or numpy.shape(kpts) == (3, ):
vne = vne[0]
return numpy.asarray(vne)
def get_nuc(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
gs = mydf.gs
charge = -cell.atom_charges()
Gv = cell.get_Gv(gs)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell, gs=gs, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, mydf.gs).real
vne = [
lib.dot(aoR.T.conj() * vneR, aoR)
for k, aoR in mydf.aoR_loop(gs, kpts_lst)
]
if kpts is None or numpy.shape(kpts) == (3, ):
vne = vne[0]
return numpy.asarray(vne)
def test_ifft(self):
n = 31
a = numpy.random.random([2,n,n,n])
ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel()
v = tools.ifft(a, [n,n,n]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
a = numpy.random.random([2,n,n,8])
ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel()
v = tools.ifft(a, [n,n,8]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
a = numpy.random.random([2,8,n,8])
ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel()
v = tools.ifft(a, [8,n,8]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
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
low_dim_ft_type = mydf.low_dim_ft_type
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, low_dim_ft_type=low_dim_ft_type)
ngrids = len(coulG)
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
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):
vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i, p0:p1],
mask, 'LDA')
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_hcore(cell, kpts):
'''Part of the nuclear gradients of core Hamiltonian'''
h1 = np.asarray(cell.pbc_intor('int1e_ipkin', kpts=kpts))
dtype = h1.dtype
if cell._pseudo:
SI = cell.get_SI()
nao = cell.nao_nr()
Gv = cell.Gv
natom = cell.natm
coords = cell.get_uniform_grids()
ngrids, nkpts = len(coords), len(kpts)
vlocG = get_vlocG(cell)
vpplocG = -np.einsum('ij,ij->j', SI, vlocG)
vpplocG[0] = np.sum(get_alphas(cell))
vpplocR = tools.ifft(vpplocG, cell.mesh).real
fakemol = _make_fakemol()
ptr = mole.PTR_ENV_START
for kn, kpt in enumerate(kpts):
aos = eval_ao_kpts(cell, coords, kpt, deriv=1)[0]
vloc = np.einsum('agi,g,gj->aij', aos[1:].conj(), vpplocR, aos[0])
expir = np.exp(-1j * np.dot(coords, kpt))
aokG = np.asarray([
tools.fftk(np.asarray(ao.T, order='C'), cell.mesh, expir).T
for ao in aos
])
Gk = Gv + kpt
G_rad = lib.norm(Gk, axis=1)
vnl = np.zeros(vloc.shape, dtype=np.complex128)
for ia in range(natom):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0, mole.ANG_OF] = l
fakemol._env[ptr + 3] = .5 * rl**2
fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25
pYlm_part = fakemol.eval_gto('GTOval', Gk)
pYlm = np.empty((nl, l * 2 + 1, ngrids))
for k in range(nl):
qkl = _qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,agp->anmp', SPG_lmi, aokG)
tmp = np.einsum('ij,ajmp->aimp', hl, SPG_lm_aoG[1:])
vnl += np.einsum('aimp,imq->apq', tmp.conj(),
SPG_lm_aoG[0])
vnl *= (1. / ngrids**2)
if dtype == np.float64:
h1[kn, :] += vloc.real + vnl.real
else:
h1[kn, :] += vloc + vnl
else:
raise NotImplementedError
return h1
def get_pp_o1(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = pseudo.get_vlocG(cell)
vpplocG = -np.sum(SI * vlocG, axis=0)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.gs)
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
# vppnonloc evaluated in reciprocal space
aokG = np.empty(aoR.shape, np.complex128)
for i in range(nao):
aokG[:,i] = tools.fftk(aoR[:,i], cell.gs, coords, kpt)
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1,pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1,pyscf.gto.BAS_SLOTS), dtype=np.int32)
fakemol._env = np.zeros(5)
fakemol._bas[0,pyscf.gto.NPRIM_OF ] = 1
fakemol._bas[0,pyscf.gto.NCTR_OF ] = 1
fakemol._bas[0,pyscf.gto.PTR_EXP ] = 3
fakemol._bas[0,pyscf.gto.PTR_COEFF] = 4
Gv = np.asarray(cell.Gv+kpt)
G_rad = pyscf.lib.norm(Gv, axis=1)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
for ia in range(cell.natm):
pp = cell._pseudo[cell.atom_symbol(ia)]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0,pyscf.gto.ANG_OF] = l
fakemol._env[3] = .5*rl**2
fakemol._env[4] = rl**(l+1.5)*np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl,l*2+1,ngs))
for k in range(nl):
qkl = pseudo.pp._qli(G_rad*rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1./ngs**2)
return vpploc + vppnl
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./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 get_vjR(cell, dm, aoR):
coulG = tools.get_coulG(cell)
rhoR = numint.eval_rho(cell, aoR, dm)
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG*rhoG
vR = tools.ifft(vG, cell.gs)
if rhoR.dtype == np.double:
vR = vR.real
return vR
def get_vjR(cell, dm, aoR):
coulG = tools.get_coulG(cell)
rhoR = numint.eval_rho(cell, aoR, dm)
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG * rhoG
vR = tools.ifft(vG, cell.gs)
if rhoR.dtype == np.double:
vR = vR.real
return vR
def test_eval_rhoG_nonorth_gga(self):
mydf = multigrid.MultiGridFFTDF(cell_nonorth)
rhoG = multigrid._eval_rhoG(mydf, dm, hermi=1, kpts=kpts, deriv=1,
rhog_high_order=True)
mydf = df.FFTDF(cell_nonorth)
ni = dft.numint.KNumInt()
ao_kpts = ni.eval_ao(cell_nonorth, mydf.grids.coords, kpts, deriv=1)
ref = ni.eval_rho(cell_nonorth, ao_kpts, dm, xctype='GGA')
rhoR = tools.ifft(rhoG[0], cell_nonorth.mesh).real
rhoR *= numpy.prod(cell_nonorth.mesh)/cell_nonorth.vol
self.assertAlmostEqual(abs(rhoR-ref).max(), 0, 5)
def hcore_deriv(atm_id):
shl0, shl1, p0, p1 = aoslices[atm_id]
symb = cell.atom_symbol(atm_id)
fakemol = _make_fakemol()
vloc_g = 1j * np.einsum('ga,g->ag', Gv, SI[atm_id] * vlocG[atm_id])
nkpts, nao = h1.shape[0], h1.shape[2]
hcore = np.zeros([3, nkpts, nao, nao], dtype=h1.dtype)
for kn, kpt in enumerate(kpts):
ao = eval_ao_kpts(cell, coords, kpt)[0]
rho = np.einsum('gi,gj->gij', ao.conj(), ao)
for ax in range(3):
vloc_R = tools.ifft(vloc_g[ax], mesh).real
vloc = np.einsum('gij,g->ij', rho, vloc_R)
hcore[ax, kn] += vloc
rho = None
aokG = tools.fftk(np.asarray(ao.T, order='C'), mesh,
np.exp(-1j * np.dot(coords, kpt))).T
ao = None
Gk = Gv + kpt
G_rad = lib.norm(Gk, axis=1)
if symb not in cell._pseudo: continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0, mole.ANG_OF] = l
fakemol._env[ptr + 3] = .5 * rl**2
fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25
pYlm_part = fakemol.eval_gto('GTOval', Gk)
pYlm = np.empty((nl, l * 2 + 1, ngrids))
for k in range(nl):
qkl = _qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
SPG_lmi = np.einsum('g,nmg->nmg', SI[atm_id].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
SPG_lmi_G = 1j * np.einsum('nmg, ga->anmg', SPG_lmi, Gv)
SPG_lm_G_aoG = np.einsum('anmg, gp->anmp', SPG_lmi_G, aokG)
tmp_1 = np.einsum('ij,ajmp->aimp', hl, SPG_lm_G_aoG)
tmp_2 = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl = np.einsum(
'imp,aimq->apq', SPG_lm_aoG.conj(), tmp_1) + np.einsum(
'aimp,imq->apq', SPG_lm_G_aoG.conj(), tmp_2)
vppnl *= (1. / ngrids**2)
if dtype == np.float64:
hcore[:, kn] += vppnl.real
else:
hcore[:, kn] += vppnl
hcore[:, kn, p0:p1] -= h1[kn, :, p0:p1]
hcore[:, kn, :, p0:p1] -= h1[kn, :, p0:p1].transpose(0, 2,
1).conj()
return hcore
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 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
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. / nkpts
rhoG = tools.fft(rhoR[i], gs)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, gs).real
kpts_band, single_kpt_band = _format_kpts_band(kpts_band, kpts)
nband = len(kpts_band)
vj_kpts = []
weight = cell.vol / ngs
if gamma_point(kpts_band):
vj_kpts = np.empty((nset, nband, nao, nao))
else:
vj_kpts = np.empty((nset, nband, nao, nao), dtype=np.complex128)
for k, aoR in mydf.aoR_loop(gs, kpts_band):
for i in range(nset):
vj_kpts[i, k] = weight * lib.dot(aoR.T.conj() * vR[i], aoR)
return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
def get_vjR_kpts(cell, dm_kpts, aoR_kpts):
nkpts = len(aoR_kpts)
coulG = tools.get_coulG(cell)
rhoR = 0
for k in range(nkpts):
rhoR += 1. / nkpts * numint.eval_rho(cell, aoR_kpts[k], dm_kpts[k])
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG * rhoG
vR = tools.ifft(vG, cell.gs)
if rhoR.dtype == np.double:
vR = vR.real
return vR
def get_vjR_kpts(cell, dm_kpts, aoR_kpts):
nkpts = len(aoR_kpts)
coulG = tools.get_coulG(cell)
rhoR = 0
for k in range(nkpts):
rhoR += 1./nkpts*numint.eval_rho(cell, aoR_kpts[k], dm_kpts[k])
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG*rhoG
vR = tools.ifft(vG, cell.gs)
if rhoR.dtype == np.double:
vR = vR.real
return vR
def test_eval_rhoG_orth_kpts(self):
numpy.random.seed(9)
dm = numpy.random.random(dm1.shape) + numpy.random.random(dm1.shape) * 1j
mydf = multigrid.MultiGridFFTDF(cell_orth)
rhoG = multigrid._eval_rhoG(mydf, dm, hermi=0, kpts=kpts, deriv=0,
rhog_high_order=True)
self.assertTrue(rhoG.dtype == numpy.complex128)
mydf = df.FFTDF(cell_orth)
ni = dft.numint.KNumInt()
ao_kpts = ni.eval_ao(cell_orth, mydf.grids.coords, kpts, deriv=0)
ref = ni.eval_rho(cell_orth, ao_kpts, dm, hermi=0, xctype='LDA')
rhoR = tools.ifft(rhoG[0], cell_orth.mesh).real
rhoR *= numpy.prod(cell_orth.mesh)/cell_orth.vol
self.assertAlmostEqual(abs(rhoR-ref).max(), 0, 7)
def get_vjR_(cell, aoR, dm):
'''Get the real-space Hartree potential of the given density matrix.
Returns:
vR : (ngs,) ndarray
The real-space Hartree potential at every grid point.
'''
coulG = tools.get_coulG(cell)
rhoR = pyscf.pbc.dft.numint.eval_rho(cell, aoR, dm)
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG*rhoG
vR = tools.ifft(vG, cell.gs)
return vR
def get_nuc(cell, kpt=np.zeros(3)):
'''Get the bare periodic nuc-el AO matrix, with G=0 removed.
See Martin (12.16)-(12.21).
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
chargs = cell.atom_charges()
SI = cell.get_SI()
coulG = tools.get_coulG(cell)
vneG = -np.dot(chargs, SI) * coulG
vneR = tools.ifft(vneG, cell.gs).real
vne = np.dot(aoR.T.conj(), vneR.reshape(-1, 1) * aoR)
return vne
def test_eval_rhoG_nonorth_gga(self):
mydf = multigrid.MultiGridFFTDF(cell_nonorth)
rhoG = multigrid._eval_rhoG(mydf,
dm,
hermi=1,
kpts=kpts,
deriv=1,
rhog_high_order=True)
mydf = df.FFTDF(cell_nonorth)
ni = dft.numint.KNumInt()
ao_kpts = ni.eval_ao(cell_nonorth, mydf.grids.coords, kpts, deriv=1)
ref = ni.eval_rho(cell_nonorth, ao_kpts, dm, xctype='GGA')
rhoR = tools.ifft(rhoG[0], cell_nonorth.mesh).real
rhoR *= numpy.prod(cell_nonorth.mesh) / cell_nonorth.vol
self.assertAlmostEqual(abs(rhoR - ref).max(), 0, 7)
def get_nuc(cell, kpt=np.zeros(3)):
'''Get the bare periodic nuc-el AO matrix, with G=0 removed.
See Martin (12.16)-(12.21).
'''
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords, kpt)
chargs = [cell.atom_charge(i) for i in range(cell.natm)]
SI = cell.get_SI()
coulG = tools.get_coulG(cell)
vneG = -np.dot(chargs,SI) * coulG
vneR = tools.ifft(vneG, cell.gs)
vne = np.dot(aoR.T.conj(), vneR.reshape(-1,1)*aoR)
return vne
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)),
kpts_band=None):
mydf = _sync_mydf(mydf)
cell = mydf.cell
mesh = mydf.mesh
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)
vR = rhoR = numpy.zeros((nset,ngrids))
for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
for i in range(nset):
rhoR[i,p0:p1] += numint.eval_rho(cell, ao, dms[i,k])
ao = ao_ks = None
rhoR = mpi.allreduce(rhoR)
for i in range(nset):
rhoR[i] *= 1./nkpts
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
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 = numpy.zeros((nset,nband,nao,nao))
else:
vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts_band):
ao_ks = ao_ks_etc[0]
for k, ao in enumerate(ao_ks):
for i in range(nset):
vj_kpts[i,k] += lib.dot(ao.T.conj()*vR[i,p0:p1], ao)
vj_kpts = mpi.reduce(vj_kpts)
if gamma_point(kpts_band):
vj_kpts = vj_kpts.real
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = pseudo.get_vlocG(cell)
vlocG[:, 0] = 0
vpplocG = -np.sum(SI * vlocG, axis=0)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.mesh)
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR)
# vppnonloc evaluated in reciprocal space
aokG = np.empty(aoR.shape, np.complex128)
expmikr = np.exp(-1j * np.dot(coords, kpt))
for i in range(nao):
aokG[:, i] = tools.fftk(aoR[:, i], cell.mesh, expmikr)
ngrids = len(aokG)
vppnl = np.zeros((nao, nao), dtype=np.complex128)
hs, projGs = pseudo.get_projG(cell, kpt)
for ia, [h_ia, projG_ia] in enumerate(zip(hs, projGs)):
for l, h in enumerate(h_ia):
nl = h.shape[0]
for m in range(-l, l + 1):
SPG_lm_aoG = np.zeros((nl, nao), dtype=np.complex128)
for i in range(nl):
SPG_lmi = SI[ia, :] * projG_ia[l][m][i]
SPG_lm_aoG[i, :] = np.einsum('g,gp->p', SPG_lmi.conj(),
aokG)
for i in range(nl):
for j in range(nl):
# Note: There is no (-1)^l here.
vppnl += h[i, j] * np.einsum('p,q->pq',
SPG_lm_aoG[i, :].conj(),
SPG_lm_aoG[j, :])
vppnl *= (1. / ngrids**2)
ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt)
vpploc += 1. / cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp
return vpploc + vppnl
def get_vjR_(cell, dm_kpts, aoR_kpts):
'''Get the real-space Hartree potential of the k-point sampled density matrix.
Returns:
vR : (ngs,) ndarray
The real-space Hartree potential at every grid point.
'''
nkpts, ngs, nao = aoR_kpts.shape
coulG = tools.get_coulG(cell)
rhoR = np.zeros(ngs)
for k in range(nkpts):
rhoR += 1./nkpts*pyscf.pbc.dft.numint.eval_rho(cell, aoR_kpts[k,:,:], dm_kpts[k,:,:])
rhoG = tools.fft(rhoR, cell.gs)
vG = coulG*rhoG
vR = tools.ifft(vG, cell.gs)
return vR
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = pseudo.get_vlocG(cell)
vlocG[:,0] = 0
vpplocG = -np.sum(SI * vlocG, axis=0)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.mesh)
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
# vppnonloc evaluated in reciprocal space
aokG = np.empty(aoR.shape, np.complex128)
expmikr = np.exp(-1j*np.dot(coords,kpt))
for i in range(nao):
aokG[:,i] = tools.fftk(aoR[:,i], cell.mesh, expmikr)
ngrids = len(aokG)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
hs, projGs = pseudo.get_projG(cell, kpt)
for ia, [h_ia,projG_ia] in enumerate(zip(hs,projGs)):
for l, h in enumerate(h_ia):
nl = h.shape[0]
for m in range(-l,l+1):
SPG_lm_aoG = np.zeros((nl,nao), dtype=np.complex128)
for i in range(nl):
SPG_lmi = SI[ia,:] * projG_ia[l][m][i]
SPG_lm_aoG[i,:] = np.einsum('g,gp->p', SPG_lmi.conj(), aokG)
for i in range(nl):
for j in range(nl):
# Note: There is no (-1)^l here.
vppnl += h[i,j]*np.einsum('p,q->pq',
SPG_lm_aoG[i,:].conj(),
SPG_lm_aoG[j,:])
vppnl *= (1./ngrids**2)
ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt)
vpploc += 1./cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp
return vpploc + vppnl
def get_jmod_pw_poisson(cell, modcell, kpt=None):
modchg = np.asarray([cell.atom_charge(ia) for ia in range(cell.natm)])
rhok = 0
k2 = np.einsum('ij,ij->i', cell.Gv, cell.Gv)
for ib in range(modcell.nbas):
e = modcell.bas_exp(ib)[0]
r = modcell.bas_coord(ib)
si = np.exp(-1j*np.einsum('ij,j->i', cell.Gv, r))
rhok += modchg[ib] * si * np.exp(-k2/(4*e))
vk = rhok * tools.get_coulG(cell)
# weight = vol/N, 1/vol * weight = 1/N
# ifft has 1/N
vw = tools.ifft(vk, cell.gs)
coords = pdft.gen_grid.gen_uniform_grids(cell)
aoR = pdft.numint.eval_ao(cell, coords, None)
vj = np.dot(aoR.T.conj(), vw.reshape(-1,1) * aoR)
return vj
def test_eval_rhoG_orth_kpts(self):
numpy.random.seed(9)
dm = numpy.random.random(
dm1.shape) + numpy.random.random(dm1.shape) * 1j
mydf = multigrid.MultiGridFFTDF(cell_orth)
rhoG = multigrid._eval_rhoG(mydf,
dm,
hermi=0,
kpts=kpts,
deriv=0,
rhog_high_order=True)
self.assertTrue(rhoG.dtype == numpy.complex128)
mydf = df.FFTDF(cell_orth)
ni = dft.numint.KNumInt()
ao_kpts = ni.eval_ao(cell_orth, mydf.grids.coords, kpts, deriv=0)
ref = ni.eval_rho(cell_orth, ao_kpts, dm, hermi=0, xctype='LDA')
rhoR = tools.ifft(rhoG[0], cell_orth.mesh).real
rhoR *= numpy.prod(cell_orth.mesh) / cell_orth.vol
self.assertAlmostEqual(abs(rhoR - ref).max(), 0, 8)
def precompute_exx(self):
print "# Precomputing Wigner-Seitz EXX kernel"
from pyscf.pbc import gto as pbcgto
Nk = tools.get_monkhorst_pack_size(self.cell, self.kpts)
print "# Nk =", Nk
kcell = pbcgto.Cell()
kcell.atom = 'H 0. 0. 0.'
kcell.spin = 1
kcell.unit = 'B'
kcell.h = self.cell._h * Nk
Lc = 1.0/np.linalg.norm(np.linalg.inv(kcell.h.T), axis=0)
print "# Lc =", Lc
Rin = Lc.min() / 2.0
print "# Rin =", Rin
# ASE:
alpha = 5./Rin # sqrt(-ln eps) / Rc, eps ~ 10^{-11}
kcell.gs = np.array([2*int(L*alpha*3.0) for L in Lc])
# QE:
#alpha = 3./Rin * np.sqrt(0.5)
#kcell.gs = (4*alpha*np.linalg.norm(kcell.h,axis=0)).astype(int)
print "# kcell.gs FFT =", kcell.gs
kcell.build(False,False)
vR = tools.ifft( tools.get_coulG(kcell), kcell.gs )
kngs = len(vR)
print "# kcell kngs =", kngs
rs = pyscf.pbc.dft.gen_grid.gen_uniform_grids(kcell)
corners = np.dot(np.indices((2,2,2)).reshape((3,8)).T, kcell._h.T)
for i, rv in enumerate(rs):
# Minimum image convention to corners of kcell parallelepiped
r = np.linalg.norm(rv-corners, axis=1).min()
if np.isclose(r, 0.):
vR[i] = 2*alpha / np.sqrt(np.pi)
else:
vR[i] = scipy.special.erf(alpha*r) / r
vG = (kcell.vol/kngs) * tools.fft(vR, kcell.gs)
self.exx_alpha = alpha
self.exx_kcell = kcell
self.exx_q = kcell.Gv
self.exx_vq = vG
self.exx_built = True
print "# Finished precomputing"
def get_nuc(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
gs = mydf.gs
charge = -cell.atom_charges()
Gv = cell.get_Gv(gs)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
coulG = tools.get_coulG(cell, gs=gs, Gv=Gv)
vneG = rhoG * coulG
vneR = tools.ifft(vneG, mydf.gs).real
vne = [lib.dot(aoR.T.conj() * vneR, aoR) for k, aoR in mydf.aoR_loop(gs, kpts_lst)]
if kpts is None or numpy.shape(kpts) == (3,):
vne = vne[0]
return vne
def get_pp_loc_part2(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
G = lib.norm(cell.Gv, axis=1)
vlocG = np.zeros((cell.natm,len(G)))
for ia in range(cell.natm):
Zia = cell.atom_charge(ia)
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
vlocG[ia] = 0
continue
pp = cell._pseudo[symb]
rloc, nexp, cexp = pp[1:3+1]
G_red = G*rloc
cfacs = np.array(
[1*G_red**0,
3 - G_red**2,
15 - 10*G_red**2 + G_red**4,
105 - 105*G_red**2 + 21*G_red**4 - G_red**6])
with np.errstate(divide='ignore'):
# Note the signs -- potential here is positive
vlocG[ia,:] = (# 4*np.pi * Zia * np.exp(-0.5*G_red**2)/G**2
- (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G_red**2)*(
np.dot(cexp, cfacs[:nexp])) )
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
if aoR.dtype == np.double:
return vpploc.real
else:
return vpploc
def get_pp_loc_part2(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
G = lib.norm(cell.Gv, axis=1)
vlocG = np.zeros((cell.natm,len(G)))
for ia in range(cell.natm):
Zia = cell.atom_charge(ia)
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
vlocG[ia] = 0
continue
pp = cell._pseudo[symb]
rloc, nexp, cexp = pp[1:3+1]
G_red = G*rloc
cfacs = np.array(
[1*G_red**0,
3 - G_red**2,
15 - 10*G_red**2 + G_red**4,
105 - 105*G_red**2 + 21*G_red**4 - G_red**6])
with np.errstate(divide='ignore'):
# Note the signs -- potential here is positive
vlocG[ia,:] = (# 4*np.pi * Zia * np.exp(-0.5*G_red**2)/G**2
- (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G_red**2)*(
np.dot(cexp, cfacs[:nexp])) )
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
if aoR.dtype == np.double:
return vpploc.real
else:
return vpploc
def get_pp(mydf, kpts=None):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
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)
# 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.aoR_loop(gs, kpts_lst)
]
# 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)
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. / ngs**2)
for k, kpt in enumerate(kpts_lst):
vppnl = vppnl_by_k(kpt)
if gamma_point(kpt):
vpp[k] = vpp[k].real + vppnl.real
else:
vpp[k] += vppnl
if kpts is None or numpy.shape(kpts) == (3, ):
vpp = vpp[0]
return numpy.asarray(vpp)
def get_pp(mydf, kpts=None):
"""Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
"""
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
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.aoR_loop(gs, kpts_lst)]
# 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
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]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = numpy.asarray(hl)
fakemol._bas[0, gto.ANG_OF] = l
fakemol._env[ptr + 3] = 0.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)
pYlm = numpy.empty((nl, l * 2 + 1, ngs))
for k in range(nl):
qkl = pseudo.pp._qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
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)
return vppnl * (1.0 / ngs ** 2)
for k, kpt in enumerate(kpts_lst):
vppnl = vppnl_by_k(kpt)
if abs(kpt).sum() < 1e-9: # gamma_point
vpp[k] = vpp[k].real + vppnl.real
else:
vpp[k] += vppnl
if kpts is None or numpy.shape(kpts) == (3,):
vpp = vpp[0]
return vpp
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix
'''
import pyscf.dft
from pyscf.pbc import tools
from pyscf.pbc.dft import gen_grid
from pyscf.pbc.dft import numint
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = get_vlocG(cell)
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR)
# vppnonloc evaluated in reciprocal space
aokG = tools.fftk(np.asarray(aoR.T, order='C'), cell.gs,
np.exp(-1j * np.dot(coords, kpt))).T
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1, pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1, pyscf.gto.BAS_SLOTS), dtype=np.int32)
ptr = pyscf.gto.PTR_ENV_START
fakemol._env = np.zeros(ptr + 10)
fakemol._bas[0, pyscf.gto.NPRIM_OF] = 1
fakemol._bas[0, pyscf.gto.NCTR_OF] = 1
fakemol._bas[0, pyscf.gto.PTR_EXP] = ptr + 3
fakemol._bas[0, pyscf.gto.PTR_COEFF] = ptr + 4
Gv = np.asarray(cell.Gv + kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao, nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0, pyscf.gto.ANG_OF] = l
fakemol._env[ptr + 3] = .5 * rl**2
fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl, l * 2 + 1, ngs))
for k in range(nl):
qkl = _qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1. / ngs**2)
if aoR.dtype == np.double:
return vpploc.real + vppnl.real
else:
return vpploc + vppnl
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
cell = mydf.cell
if kpts is None:
kpts = mydf.kpts
nkpts = len(kpts)
if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3:
mo_coeff_kpts = [mo_coeff_kpts] * 4
else:
mo_coeff_kpts = list(mo_coeff_kpts)
mo_ids = [id(x) for x in mo_coeff_kpts]
moTs = []
coords = cell.gen_uniform_grids(mydf.mesh)
aos = mydf._numint.eval_ao(cell, coords, kpts)
for n, mo_id in enumerate(mo_ids):
if mo_id in mo_ids[:n]:
moTs.append(moTs[mo_ids[:n].index(mo_id)])
else:
moTs.append([lib.dot(mo.T, aos[k].T) for k,mo in enumerate(mo_coeff_kpts[n])])
# Shape of the orbitals can be different on different k-points. The
# orbital coefficients must be formatted (padded by zeros) so that the
# shape of the orbital coefficients are the same on all k-points. This can
# be achieved by calling pbc.mp.kmp2.padded_mo_coeff function
nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts]
eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol)
if gamma_point(kpts):
dtype = numpy.result_type(*mo_coeff_kpts)
else:
dtype = numpy.complex128
if out is None:
out = numpy.empty(eri_shape, dtype=dtype)
else:
assert(out.shape == eri_shape)
kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts])
kptis_lst = kptij_lst[:,0]
kptjs_lst = kptij_lst[:,1]
kpt_ji = kptjs_lst - kptis_lst
uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
ngrids = numpy.prod(mydf.mesh)
# To hold intermediates
fswap = lib.H5TmpFile()
kconserv = kpts_helper.get_kconserv(cell, kpts)
for uniq_id, kpt in enumerate(uniq_kpts):
q = uniq_kpts[uniq_id]
adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0]
ki = adapted_ji_idx[0] // nkpts
kj = adapted_ji_idx[0] % nkpts
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coulG *= (cell.vol/ngrids) * factor
phase = numpy.exp(-1j * numpy.dot(coords, q))
for kk in range(nkpts):
kl = kconserv[ki, kj, kk]
mokT = moTs[2][kk]
molT = moTs[3][kl]
mo_pairs = numpy.einsum('ig,g,jg->ijg', mokT.conj(), phase.conj(), molT)
v = tools.ifft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
v *= coulG
v = tools.fft(v.reshape(-1,ngrids), mydf.mesh)
v *= phase
fswap['zkl/'+str(kk)] = v
for ji_idx in adapted_ji_idx:
ki = ji_idx // nkpts
kj = ji_idx % nkpts
for kk in range(nkpts):
moiT = moTs[0][ki]
mojT = moTs[1][kj]
mo_pairs = numpy.einsum('ig,jg->ijg', moiT.conj(), mojT)
tmp = lib.dot(mo_pairs.reshape(-1,ngrids),
numpy.asarray(fswap['zkl/'+str(kk)]).T)
if dtype == numpy.double:
tmp = tmp.real
out[ki,kj,kk] = tmp.reshape(eri_shape[3:])
del(fswap['zkl'])
return out
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix
'''
from pyscf.pbc import tools
coords = cell.get_uniform_grids()
aoR = cell.pbc_eval_gto('GTOval', coords, kpt=kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = get_vlocG(cell)
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.mesh).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
# vppnonloc evaluated in reciprocal space
aokG = tools.fftk(np.asarray(aoR.T, order='C'),
cell.mesh, np.exp(-1j*np.dot(coords, kpt))).T
ngrids = len(aokG)
fakemol = mole.Mole()
fakemol._atm = np.zeros((1,mole.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1,mole.BAS_SLOTS), dtype=np.int32)
ptr = mole.PTR_ENV_START
fakemol._env = np.zeros(ptr+10)
fakemol._bas[0,mole.NPRIM_OF ] = 1
fakemol._bas[0,mole.NCTR_OF ] = 1
fakemol._bas[0,mole.PTR_EXP ] = ptr+3
fakemol._bas[0,mole.PTR_COEFF] = ptr+4
Gv = np.asarray(cell.Gv+kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0,mole.ANG_OF] = l
fakemol._env[ptr+3] = .5*rl**2
fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25
pYlm_part = fakemol.eval_gto('GTOval', Gv)
pYlm = np.empty((nl,l*2+1,ngrids))
for k in range(nl):
qkl = _qli(G_rad*rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1./ngrids**2)
if aoR.dtype == np.double:
return vpploc.real + vppnl.real
else:
return vpploc + vppnl
def get_k_e1_kpts(mydf,
dm_kpts,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Derivatives of exchange (K) AO matrix at sampled k-points.
'''
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_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)
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_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_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None):
'''Derivatives of Coulomb (J) AO matrix at sampled k-points.
'''
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=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_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)]
naoj = ao2_kpts[0].shape[0]
else:
naoj = nao
max_memory = mydf.max_memory - lib.current_memory()[0]
blksize = int(max(max_memory * 1e6 / 16 / 2 / ngs / nao, 1))
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
buf = np.empty((blksize, naoj, ngs),
dtype=np.result_type(ao1_dtype, ao2_dtype))
vR_dm = np.empty((nset, nao, ngs), dtype=vk_kpts.dtype)
ao_dms = np.empty((nset, naoj, ngs), dtype=np.result_type(dms, ao2_dtype))
for k2, ao2T in enumerate(ao2_kpts):
kpt2 = kpts[k2]
if mo_coeff is None or nset > 1:
for i in range(nset):
lib.dot(dms[i, k2], ao2T.conj(), c=ao_dms[i])
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,
out=buf[:p1 - p0])
vG = tools.fft(rho1.reshape(-1, ngs), gs)
vG *= coulG
vR = tools.ifft(vG, gs).reshape(p1 - p0, naoj, 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)