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 _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 _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 test_fft(self):
n = 31
a = numpy.random.random([2,n,n,n])
ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel()
v = tools.fft(a, [n,n,n]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
a = numpy.random.random([2,n,n,8])
ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel()
v = tools.fft(a, [n,n,8]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
a = numpy.random.random([2,8,n,8])
ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel()
v = tools.fft(a, [8,n,8]).ravel()
self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None,
compact=getattr(__config__, 'pbc_df_mo_pairs_compact', False)):
'''Calculate forward (G|ij) FFT of all MO pairs.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G : (ngrids, nmoi*nmoj) ndarray
The FFT of the real-space MO pairs.
'''
if kpts is None: kpts = numpy.zeros((2,3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.mesh)
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
ngrids = len(coords)
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T,aoi.T), order='C')
else:
moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C')
moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C')
mo_pairs_G = numpy.empty((nmoi,nmoj,ngrids), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.mesh)
mo_pairs_G = mo_pairs_G.reshape(-1,ngrids).T
return mo_pairs_G
if gamma_point(kpts): # gamma point, real
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
mo = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T), order='C')
npair = nmoi*(nmoi+1)//2
mo_pairs_G = numpy.empty((npair,ngrids), dtype=numpy.complex128)
ij = 0
for i in range(nmoi):
mo_pairs_G[ij:ij+i+1] = tools.fft(mo[i].conj() * mo[:i+1], mydf.mesh)
ij += i + 1
mo_pairs_G = mo_pairs_G.T
else:
mo_pairs_G = trans(ao, ao)
elif is_zero(kpts[0]-kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
mo_pairs_G = trans(ao, ao)
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts)
fac = numpy.exp(-1j * numpy.dot(coords, q))
mo_pairs_G = trans(aoi, aoj, fac)
return mo_pairs_G
def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), compact=False):
'''Calculate forward (G|ij) FFT of all MO pairs.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G : (ngs, nmoi*nmoj) ndarray
The FFT of the real-space MO pairs.
'''
if kpts is None: kpts = mydf.kpts
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = pdft.gen_grid.gen_uniform_grids(cell, mydf.gs)
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
ngs = len(coords)
def trans(aoiR, aojR, fac=1):
if id(aoiR) == id(aojR) and ao2mo.incore.iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moiR = mojR = numpy.asarray(lib.dot(mo_coeffs[0].T,aoiR.T), order='C')
else:
moiR = numpy.asarray(lib.dot(mo_coeffs[0].T, aoiR.T), order='C')
mojR = numpy.asarray(lib.dot(mo_coeffs[1].T, aojR.T), order='C')
mo_pairs_G = numpy.empty((nmoi,nmoj,ngs), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moiR[i].conj() * mojR, mydf.gs)
mo_pairs_G = mo_pairs_G.reshape(-1,ngs).T
return mo_pairs_G
if abs(kpts).sum() < 1e-9: # gamma point, real
aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
if compact and ao2mo.incore.iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moR = numpy.asarray(lib.dot(mo_coeffs[0].T, aoR.T), order='C')
npair = nmoi*(nmoi+1)//2
mo_pairs_G = numpy.empty((npair,ngs), dtype=numpy.complex128)
ij = 0
for i in range(nmoi):
mo_pairs_G[ij:ij+i+1] = tools.fft(moR[i].conj() * moR[:i+1], mydf.gs)
ij += i + 1
mo_pairs_G = mo_pairs_G.T
else:
mo_pairs_G = trans(aoR, aoR)
elif abs(kpts[0]-kpts[1]).sum() < 1e-9:
aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
mo_pairs_G = trans(aoR, aoR)
else:
aoiR, aojR = mydf._numint.eval_ao(cell, coords, kpts)
q = kpts[1] - kpts[0]
fac = numpy.exp(-1j * numpy.dot(coords, q))
mo_pairs_G = trans(aoiR, aojR, fac)
return mo_pairs_G
def test_ft_ao(self):
coords = pdft.gen_grid.gen_uniform_grids(cell)
aoR = pdft.numint.eval_ao(cell, coords)
ngs, nao = aoR.shape
ref = numpy.asarray([tools.fft(aoR[:,i], cell.gs) for i in range(nao)])
ref = ref.T * (cell.vol/ngs)
dat = ft_ao.ft_ao(cell, cell.Gv)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,0]-dat[:,0]) , 8.4358614794095722e-11, 9)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,1]-dat[:,1]) , 0.0041669297531642616 , 4)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:]-dat[:,2:]), 5.8677286005879366e-14, 9)
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 trans(aoiR, aojR, fac=1):
if id(aoiR) == id(aojR):
aoiR = aojR = numpy.asarray(aoiR.T, order='C')
else:
aoiR = numpy.asarray(aoiR.T, order='C')
aojR = numpy.asarray(aojR.T, order='C')
ao_pairs_G = numpy.empty((nao,nao,ngs), dtype=numpy.complex128)
for i in range(nao):
ao_pairs_G[i] = tools.fft(fac * aoiR[i].conj() * aojR, mydf.gs)
ao_pairs_G = ao_pairs_G.reshape(-1,ngs).T
return ao_pairs_G
def trans(aoiR, aojR, fac=1):
if id(aoiR) == id(aojR) and ao2mo.incore.iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moiR = mojR = numpy.asarray(lib.dot(mo_coeffs[0].T,aoiR.T), order='C')
else:
moiR = numpy.asarray(lib.dot(mo_coeffs[0].T, aoiR.T), order='C')
mojR = numpy.asarray(lib.dot(mo_coeffs[1].T, aojR.T), order='C')
mo_pairs_G = numpy.empty((nmoi,nmoj,ngs), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moiR[i].conj() * mojR, mydf.gs)
mo_pairs_G = mo_pairs_G.reshape(-1,ngs).T
return mo_pairs_G
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T,aoi.T), order='C')
else:
moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C')
moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C')
mo_pairs_G = numpy.empty((nmoi,nmoj,ngrids), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.mesh)
mo_pairs_G = mo_pairs_G.reshape(-1,ngrids).T
return mo_pairs_G
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 trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj):
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
ni = aoi.shape[0]
nj = aoj.shape[0]
ao_pairs_G = numpy.empty((ni,nj,ngrids), dtype=numpy.complex128)
for i in range(ni):
ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.mesh)
ao_pairs_G = ao_pairs_G.reshape(-1,ngrids).T
return ao_pairs_G
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_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), compact=False):
'''Calculate forward (G|ij) FFT of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngs, nao*(nao+1)/2); otherwise the
shape is (ngs, nao*nao)
'''
if kpts is None: kpts = mydf.kpts
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = pdft.gen_grid.gen_uniform_grids(cell, mydf.gs)
nao = cell.nao_nr()
ngs = len(coords)
def trans(aoiR, aojR, fac=1):
if id(aoiR) == id(aojR):
aoiR = aojR = numpy.asarray(aoiR.T, order='C')
else:
aoiR = numpy.asarray(aoiR.T, order='C')
aojR = numpy.asarray(aojR.T, order='C')
ao_pairs_G = numpy.empty((nao,nao,ngs), dtype=numpy.complex128)
for i in range(nao):
ao_pairs_G[i] = tools.fft(fac * aoiR[i].conj() * aojR, mydf.gs)
ao_pairs_G = ao_pairs_G.reshape(-1,ngs).T
return ao_pairs_G
if compact and abs(kpts).sum() < 1e-9: # gamma point
aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
aoR = numpy.asarray(aoR.T, order='C')
npair = nao*(nao+1)//2
ao_pairs_G = numpy.empty((npair,ngs), dtype=numpy.complex128)
ij = 0
for i in range(nao):
ao_pairs_G[ij:ij+i+1] = tools.fft(aoR[i] * aoR[:i+1], mydf.gs)
ij += i + 1
ao_pairs_G = ao_pairs_G.T
elif abs(kpts[0]-kpts[1]).sum() < 1e-9:
aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao_pairs_G = trans(aoR, aoR)
else:
aoiR, aojR = mydf._numint.eval_ao(cell, coords, kpts[:2])
q = kpts[1] - kpts[0]
fac = numpy.exp(-1j * numpy.dot(coords, q))
ao_pairs_G = trans(aoiR, aojR, fac)
return ao_pairs_G
def get_ao_pairs_G(cell, kpt=np.zeros(3)):
'''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.
Args:
cell : instance of :class:`Cell`
Returns:
ao_pairs_G, ao_pairs_invG : (ngs, nao*(nao+1)/2) ndarray
The FFTs of the real-space AO pairs.
'''
coords = gen_uniform_grids(cell)
aoR = eval_ao(cell, coords, kpt) # shape = (coords, nao)
ngs, nao = aoR.shape
gamma_point = abs(kpt).sum() < 1e-9
if gamma_point:
npair = nao*(nao+1)//2
ao_pairs_G = np.empty([ngs, npair], np.complex128)
ij = 0
for i in range(nao):
for j in range(i+1):
ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
ao_pairs_G[:,ij] = tools.fft(ao_ij_R, cell.gs)
#ao_pairs_invG[:,ij] = ngs*tools.ifft(ao_ij_R, cell.gs)
ij += 1
ao_pairs_invG = ao_pairs_G.conj()
else:
ao_pairs_G = np.zeros([ngs, nao,nao], np.complex128)
for i in range(nao):
for j in range(nao):
ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
ao_pairs_G[:,i,j] = tools.fft(ao_ij_R, cell.gs)
ao_pairs_invG = ao_pairs_G.transpose(0,2,1).conj().reshape(-1,nao**2)
ao_pairs_G = ao_pairs_G.reshape(-1,nao**2)
return ao_pairs_G, ao_pairs_invG
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 test_ft_aoao(self):
coords = pdft.gen_grid.gen_uniform_grids(cell)
aoR = pdft.numint.eval_ao(cell, coords)
ngs, nao = aoR.shape
ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.gs)
for i in range(nao) for j in range(nao)])
ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngs)
dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi')
self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 1.869103994619606e-06 , 7)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373 , 4)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 5.4648896424693173e-14, 9)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 4.0352047774658308e-11, 9)
self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 4.0352047774658308e-11, 9)
idx = numpy.tril_indices(nao)
ref = dat[:,idx[0],idx[1]]
dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2')
self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9)
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 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 _fft_quad_integrals(mydf, dm, efg_nuc):
# Use FFTDF to compute the integrals of quadrupole operator
# (3 \vec{r} \vec{r} - r^2) / r^5
cell = mydf.cell
if cell.dimension != 3:
raise NotImplementedError
mesh = mydf.mesh
kpts = mydf.kpts
kpts_lst = numpy.reshape(kpts, (-1,3))
nkpts = len(kpts_lst)
nao = cell.nao_nr()
dm_kpts = dm.reshape((nkpts,nao,nao), order='C')
ni = mydf._numint
hermi = 1
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi)
ngrids = numpy.prod(mesh)
rhoR = numpy.zeros(ngrids)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_lst):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
rhoR[p0:p1] += make_rho(0, ao_ks, mask, 'LDA')
ao = ao_ks = None
rhoG = tools.fft(rhoR, mesh)
Gv = cell.get_Gv(mesh)
coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv)
GG = numpy.einsum('gx,gy->gxy', Gv, Gv)
absG2 = numpy.einsum('gxx->g', GG)
# Corresponding to FC term, that makes the tensor traceless
idx = numpy.arange(3)
GG[:,idx,idx] -= 1./3 * absG2[:,None]
vG = 1./ngrids * numpy.einsum('g,g,gxy->gxy', rhoG, coulG, GG)
SI = cell.get_SI(Gv)
efg_e = lib.einsum('zg,gxy->zxy', SI[efg_nuc], vG.conj()).real
return efg_e
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
low_dim_ft_type = mydf.low_dim_ft_type
coords = cell.gen_uniform_grids(mesh)
ngrids = 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 / 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,
low_dim_ft_type=low_dim_ft_type)
else:
coulG = tools.get_coulG(cell,
kpt2 - kpt1,
True,
mydf,
mesh,
low_dim_ft_type=low_dim_ft_type)
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_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_3d(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), q=None, shls_slice=None,
compact=getattr(__config__, 'pbc_df_ao_pairs_compact', False)):
'''Calculate forward (G|ij) FFT of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the
shape is (ngrids, nao*nao)
'''
if kpts is None: kpts = numpy.zeros((2,3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
if shls_slice is None:
i0, i1 = j0, j1 = (0, cell.nao_nr())
else:
ish0, ish1, jsh0, jsh1 = shls_slice
ao_loc = cell.ao_loc_nr()
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
j0 = ao_loc[jsh0]
j1 = ao_loc[jsh1]
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj):
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
ni = aoi.shape[0]
nj = aoj.shape[0]
ao_pairs_G = numpy.empty((ni,nj,ngrids), dtype=numpy.complex128)
for i in range(ni):
ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.mesh)
ao_pairs_G = ao_pairs_G.reshape(-1,ngrids).T
return ao_pairs_G
if compact and gamma_point(kpts): # gamma point
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao = numpy.asarray(ao.T, order='C')
npair = i1*(i1+1)//2 - i0*(i0+1)//2
ao_pairs_G = numpy.empty((npair,ngrids), dtype=numpy.complex128)
ij = 0
for i in range(i0, i1):
ao_pairs_G[ij:ij+i+1] = tools.fft(ao[i] * ao[:i+1], mydf.mesh)
ij += i + 1
ao_pairs_G = ao_pairs_G.T
elif is_zero(kpts[0]-kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao_pairs_G = trans(ao[:,i0:i1], ao[:,j0:j1])
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts[:2])
fac = numpy.exp(-1j * numpy.dot(coords, q))
ao_pairs_G = trans(aoi[:,i0:i1], aoj[:,j0:j1], fac)
return ao_pairs_G
def get_k_kpts_occ(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)
]
# occ
if mo_coeff is None:
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
# occ
if gamma_point(kpts_band) and gamma_point(kpts):
for i in range(nset):
occ = mo_occ[i][0]
kiv = numpy.zeros(
(nset, nband, mo_coeff[i][0][:, occ > 0].shape[1], nao),
dtype=dms.dtype)
else:
kiv = [[] * nset]
for i in range(nset):
for k1 in range(nband):
occ = mo_occ[i][k1]
kiv[i].append(
numpy.zeros((mo_coeff[i][k1][:, occ > 0].shape[1], nao),
dtype=numpy.complex128))
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
# occ
ao3T_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 ao3T_buf[k1] is None:
if mo_coeff is not None:
ao3T = []
for i, dm in enumerate(dms):
occ = mo_occ[i][k1]
mo_scaled = mo_coeff[i][k1][:, occ > 0]
ao3T.append(lib.dot(mo_scaled.T, ao1T))
ao3T_buf[k1] = ao3T
else:
ao3T = ao3T_buf[k1]
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 i in range(nset):
for p0, p1 in lib.prange(0, ao3T[i].shape[0], blksize):
rho1 = numpy.einsum('ig,jg->ijg',
ao3T[i][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):
kiv[i][k1] += weight * lib.dot(
vR_dm[i, 0:mo_coeff[i][k1][:, mo_occ[i][k1] > 0].shape[1]],
ao1T.T)
kiv = mpi.reduce(lib.asarray(kiv))
if rank == 0:
for i in range(nset):
for k1 in range(nband):
kij = lib.einsum('ui,ju->ij',
mo_coeff[i][k1][:, mo_occ[i][k1] > 0],
kiv[i][k1])
kr = scipy.linalg.solve(kij.conj(), kiv[i][k1])
vk_kpts[i, k1] = lib.dot(kiv[i][k1].T.conj(), kr)
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)
cell.a = numpy.diag([L, L, L])
cell.mesh = numpy.array([n, n, n])
cell.atom = '''C 1.3 .2 .3
C .1 .1 1.1
'''
cell.basis = 'ccpvdz'
#cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
#cell.basis = {'C': [[0, (2.4, 1)]]}
cell.unit = 'B'
#cell.verbose = 4
cell.build(0, 0)
#cell.nimgs = (2,2,2)
ao2 = ft_aopair(cell, cell.Gv)
nao = cell.nao_nr()
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = cell.pbc_eval_gto('GTOval', coords)
aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR)
ngrids = aoR.shape[0]
for i in range(nao):
for j in range(nao):
ao2ref = tools.fft(aoR2[:, i, j], cell.mesh) * cell.vol / ngrids
print(i, j, numpy.linalg.norm(ao2ref - ao2[:, i, j]))
aoG = ft_ao(cell, cell.Gv)
for i in range(nao):
aoref = tools.fft(aoR[:, i], cell.mesh) * cell.vol / ngrids
print(i, numpy.linalg.norm(aoref - aoG[:, i]))
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)
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_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)
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_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)
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 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 _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_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)
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)
cell.gs = numpy.array([n,n,n])
cell.atom = '''C 1.3 .2 .3
C .1 .1 1.1
'''
cell.basis = 'ccpvdz'
#cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
#cell.basis = {'C': [[0, (2.4, 1)]]}
cell.unit = 'B'
#cell.verbose = 4
cell.build(0,0)
#cell.nimgs = (2,2,2)
ao2 = ft_aopair(cell, cell.Gv)
nao = cell.nao_nr()
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords)
aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR)
ngs = aoR.shape[0]
for i in range(nao):
for j in range(nao):
ao2ref = tools.fft(aoR2[:,i,j], cell.gs) * cell.vol/ngs
print i, j, numpy.linalg.norm(ao2ref - ao2[:,i,j])
aoG = ft_ao(cell, cell.Gv)
for i in range(nao):
aoref = tools.fft(aoR[:,i], cell.gs) * cell.vol/ngs
print i, numpy.linalg.norm(aoref - aoG[:,i])
cell.mesh = numpy.array([n,n,n])
cell.atom = '''C 1.3 .2 .3
C .1 .1 1.1
'''
cell.basis = 'ccpvdz'
#cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
#cell.basis = {'C': [[0, (2.4, 1)]]}
cell.unit = 'B'
#cell.verbose = 4
cell.build(0,0)
#cell.nimgs = (2,2,2)
ao2 = ft_aopair(cell, cell.Gv)
nao = cell.nao_nr()
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = cell.pbc_eval_gto('GTOval', coords)
aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR)
ngrids = aoR.shape[0]
for i in range(nao):
for j in range(nao):
ao2ref = tools.fft(aoR2[:,i,j], cell.mesh) * cell.vol/ngrids
print(i, j, numpy.linalg.norm(ao2ref - ao2[:,i,j]))
aoG = ft_ao(cell, cell.Gv)
for i in range(nao):
aoref = tools.fft(aoR[:,i], cell.mesh) * cell.vol/ngrids
print(i, numpy.linalg.norm(aoref - aoG[:,i]))
def test_ft_aoao(self):
#coords = pdft.gen_grid.gen_uniform_grids(cell)
#aoR = pdft.numint.eval_ao(cell, coords)
#ngrids, nao = aoR.shape
#ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.mesh)
# for i in range(nao) for j in range(nao)])
#ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngrids)
#dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi')
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 0, 5)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373, 4)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 0, 9)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 0, 9)
#self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 0, 9)
#idx = numpy.tril_indices(nao)
#ref = dat[:,idx[0],idx[1]]
#dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2')
#self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9)
coords = pdft.gen_grid.gen_uniform_grids(cell1)
Gv, Gvbase, kws = cell1.get_Gv_weights(cell1.mesh)
b = cell1.reciprocal_vectors()
gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s1',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
1.5666516306798806 + 1.953555017583245j, 9)
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s2',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
-0.85276967757297917 + 1.0378751267506394j, 9)
dat = ft_ao.ft_aopair(cell1,
cell1.Gv,
aosym='s1hermi',
b=b,
gxyz=gxyz,
Gvbase=Gvbase)
self.assertAlmostEqual(finger(dat),
1.5666516306798806 + 1.953555017583245j, 9)
aoR = pdft.numint.eval_ao(cell1, coords)
ngrids, nao = aoR.shape
aoaoR = numpy.einsum('pi,pj->ijp', aoR, aoR)
ref = tools.fft(aoaoR.reshape(nao * nao, -1), cell1.mesh)
ref = ref.reshape(nao, nao, -1).transpose(2, 0,
1) * (cell1.vol / ngrids)
self.assertAlmostEqual(numpy.linalg.norm(ref[:, 0, 0] - dat[:, 0, 0]),
0, 7)
self.assertAlmostEqual(numpy.linalg.norm(ref[:, 1, 1] - dat[:, 1, 1]),
0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 2:, 2:] - dat[:, 2:, 2:]), 0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 0, 2:] - dat[:, 0, 2:]), 0, 7)
self.assertAlmostEqual(
numpy.linalg.norm(ref[:, 2:, 0] - dat[:, 2:, 0]), 0, 7)
idx = numpy.tril_indices(nao)
ref = dat[:, idx[0], idx[1]]
dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s2')
self.assertAlmostEqual(abs(dat - ref).sum(), 0, 9)
def get_ao_pairs_G(mydf,
kpts=numpy.zeros((2, 3)),
q=None,
shls_slice=None,
compact=False):
'''Calculate forward (G|ij) FFT of all AO pairs.
Returns:
ao_pairs_G : 2D complex array
For gamma point, the shape is (ngs, nao*(nao+1)/2); otherwise the
shape is (ngs, nao*nao)
'''
if kpts is None: kpts = numpy.zeros((2, 3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.gs)
ngs = len(coords)
if shls_slice is None:
i0, i1 = j0, j1 = (0, cell.nao_nr())
else:
ish0, ish1, jsh0, jsh1 = shls_slice
ao_loc = cell.ao_loc_nr()
i0 = ao_loc[ish0]
i1 = ao_loc[ish1]
j0 = ao_loc[jsh0]
j1 = ao_loc[jsh1]
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj):
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
ni = aoi.shape[0]
nj = aoj.shape[0]
ao_pairs_G = numpy.empty((ni, nj, ngs), dtype=numpy.complex128)
for i in range(ni):
ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.gs)
ao_pairs_G = ao_pairs_G.reshape(-1, ngs).T
return ao_pairs_G
if compact and gamma_point(kpts): # gamma point
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao = numpy.asarray(ao.T, order='C')
npair = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2
ao_pairs_G = numpy.empty((npair, ngs), dtype=numpy.complex128)
ij = 0
for i in range(i0, i1):
ao_pairs_G[ij:ij + i + 1] = tools.fft(ao[i] * ao[:i + 1], mydf.gs)
ij += i + 1
ao_pairs_G = ao_pairs_G.T
elif is_zero(kpts[0] - kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
ao_pairs_G = trans(ao[:, i0:i1], ao[:, j0:j1])
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts[:2])
fac = numpy.exp(-1j * numpy.dot(coords, q))
ao_pairs_G = trans(aoi[:, i0:i1], aoj[:, j0:j1], fac)
return ao_pairs_G
cell.a = numpy.diag([L, L, L])
cell.gs = numpy.array([n, n, n])
cell.atom = '''C 1.3 .2 .3
C .1 .1 1.1
'''
cell.basis = 'ccpvdz'
#cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
#cell.basis = {'C': [[0, (2.4, 1)]]}
cell.unit = 'B'
#cell.verbose = 4
cell.build(0, 0)
#cell.nimgs = (2,2,2)
ao2 = ft_aopair(cell, cell.Gv)
nao = cell.nao_nr()
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords)
aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR)
ngs = aoR.shape[0]
for i in range(nao):
for j in range(nao):
ao2ref = tools.fft(aoR2[:, i, j], cell.gs) * cell.vol / ngs
print(i, j, numpy.linalg.norm(ao2ref - ao2[:, i, j]))
aoG = ft_ao(cell, cell.Gv)
for i in range(nao):
aoref = tools.fft(aoR[:, i], cell.gs) * cell.vol / ngs
print(i, numpy.linalg.norm(aoref - aoG[:, i]))
def get_mo_pairs_G(mydf,
mo_coeffs,
kpts=numpy.zeros((2, 3)),
q=None,
compact=False):
'''Calculate forward (G|ij) FFT of all MO pairs.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G : (ngs, nmoi*nmoj) ndarray
The FFT of the real-space MO pairs.
'''
if kpts is None: kpts = numpy.zeros((2, 3))
cell = mydf.cell
kpts = numpy.asarray(kpts)
coords = cell.gen_uniform_grids(mydf.gs)
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
ngs = len(coords)
def trans(aoi, aoj, fac=1):
if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T),
order='C')
else:
moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C')
moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C')
mo_pairs_G = numpy.empty((nmoi, nmoj, ngs), dtype=numpy.complex128)
for i in range(nmoi):
mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.gs)
mo_pairs_G = mo_pairs_G.reshape(-1, ngs).T
return mo_pairs_G
if gamma_point(kpts): # gamma point, real
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
mo = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T), order='C')
npair = nmoi * (nmoi + 1) // 2
mo_pairs_G = numpy.empty((npair, ngs), dtype=numpy.complex128)
ij = 0
for i in range(nmoi):
mo_pairs_G[ij:ij + i + 1] = tools.fft(
mo[i].conj() * mo[:i + 1], mydf.gs)
ij += i + 1
mo_pairs_G = mo_pairs_G.T
else:
mo_pairs_G = trans(ao, ao)
elif is_zero(kpts[0] - kpts[1]):
ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0]
mo_pairs_G = trans(ao, ao)
else:
if q is None:
q = kpts[1] - kpts[0]
aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts)
fac = numpy.exp(-1j * numpy.dot(coords, q))
mo_pairs_G = trans(aoi, aoj, fac)
return mo_pairs_G
def _make_fftdf_eris(mycc, eris):
mydf = mycc._scf.with_df
mo_coeff = eris.mo_coeff
kpts = mycc.kpts
logger = Logger(mycc.stdout, mycc.verbose)
cell = mydf.cell
gvec = cell.reciprocal_vectors()
nao = cell.nao_nr()
coords = cell.gen_uniform_grids(mydf.mesh)
ngrids = len(coords)
nkpts = len(kpts)
nocc, nmo = mycc.nocc, mycc.nmo
nvir = nmo - nocc
cput1 = cput0 = (time.clock(), time.time())
ijG = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128)
iaG = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128)
abG = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128)
ijR = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128)
iaR = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128)
aiR = ctf.zeros([nkpts,nkpts,nvir,nocc,ngrids], dtype=np.complex128)
abR = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128)
jobs = []
for ki in range(nkpts):
for kj in range(ki,nkpts):
jobs.append([ki,kj])
tasks = list(static_partition(jobs))
ntasks = max(comm.allgather(len(tasks)))
idx_ooG = np.arange(nocc*nocc*ngrids)
idx_ovG = np.arange(nocc*nvir*ngrids)
idx_vvG = np.arange(nvir*nvir*ngrids)
for itask in range(ntasks):
if itask >= len(tasks):
ijR.write([], [])
iaR.write([], [])
aiR.write([], [])
abR.write([], [])
ijR.write([], [])
iaR.write([], [])
aiR.write([], [])
abR.write([], [])
ijG.write([], [])
iaG.write([], [])
abG.write([], [])
ijG.write([], [])
iaG.write([], [])
abG.write([], [])
continue
ki, kj = tasks[itask]
kpti, kptj = kpts[ki], kpts[kj]
ao_kpti = mydf._numint.eval_ao(cell, coords, kpti)[0]
ao_kptj = mydf._numint.eval_ao(cell, coords, kptj)[0]
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
wcoulG = coulG * (cell.vol/ngrids)
fac = np.exp(-1j * np.dot(coords, q))
mo_kpti = np.dot(ao_kpti, mo_coeff[ki]).T
mo_kptj = np.dot(ao_kptj, mo_coeff[kj]).T
mo_pairs = np.einsum('ig,jg->ijg', mo_kpti.conj(), mo_kptj)
mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids)*fac, mydf.mesh)
off = ki * nkpts + kj
ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel())
iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel())
aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel())
abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel())
off = kj * nkpts + ki
mo_pairs = mo_pairs.transpose(1,0,2).conj()
ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel())
iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel())
aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel())
abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel())
mo_pairs = None
mo_pairs_G*= wcoulG
v = tools.ifft(mo_pairs_G, mydf.mesh)
v *= fac.conj()
v = v.reshape(nmo,nmo,ngrids)
off = ki * nkpts + kj
ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel())
iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel())
abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel())
off = kj * nkpts + ki
v = v.transpose(1,0,2).conj()
ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel())
iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel())
abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel())
cput1 = logger.timer("Generating ijG", *cput1)
sym1 = ["+-+", [kpts,]*3, None, gvec]
sym2 = ["+--", [kpts,]*3, None, gvec]
ooG = tensor(ijG, sym1, verbose=mycc.SYMVERBOSE)
ovG = tensor(iaG, sym1, verbose=mycc.SYMVERBOSE)
vvG = tensor(abG, sym1, verbose=mycc.SYMVERBOSE)
ooR = tensor(ijR, sym2, verbose=mycc.SYMVERBOSE)
ovR = tensor(iaR, sym2, verbose=mycc.SYMVERBOSE)
voR = tensor(aiR, sym2, verbose=mycc.SYMVERBOSE)
vvR = tensor(abR, sym2, verbose=mycc.SYMVERBOSE)
eris.oooo = lib.einsum('ijg,klg->ijkl', ooG, ooR)/ nkpts
eris.ooov = lib.einsum('ijg,kag->ijka', ooG, ovR)/ nkpts
eris.oovv = lib.einsum('ijg,abg->ijab', ooG, vvR)/ nkpts
ooG = ooR = ijG = ijR = None
eris.ovvo = lib.einsum('iag,bjg->iabj', ovG, voR)/ nkpts
eris.ovov = lib.einsum('iag,jbg->iajb', ovG, ovR)/ nkpts
ovR = iaR = voR = aiR = None
eris.ovvv = lib.einsum('iag,bcg->iabc', ovG, vvR)/ nkpts
ovG = iaG = None
eris.vvvv = lib.einsum('abg,cdg->abcd', vvG, vvR)/ nkpts
cput1 = logger.timer("ijG to eri", *cput1)
