def test_coulG(self):
numpy.random.seed(19)
kpt = numpy.random.random(3)
ase_atom = bulk('C', 'diamond', a=3.5668)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom)
cell.a = ase_atom.cell + numpy.random.random((3, 3)).T
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.gs = [5, 4, 3]
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
coulG = tools.get_coulG(cell, kpt)
self.assertAlmostEqual(finger(coulG), 62.75448804333378, 9)
cell.a = numpy.eye(3)
cell.unit = 'B'
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]))
self.assertAlmostEqual(finger(coulG), 4.6737453679713905, 9)
coulG = tools.get_coulG(cell,
numpy.array([0, numpy.pi, 0]),
wrap_around=False)
self.assertAlmostEqual(finger(coulG), 4.5757877990664744, 9)
coulG = tools.get_coulG(cell, exx='ewald')
self.assertAlmostEqual(finger(coulG), 4.888843468914021, 9)
def test_coulG(self):
numpy.random.seed(19)
kpt = numpy.random.random(3)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = 'C 0., 0., 0.; C 0.8917, 0.8917, 0.8917'
cell.a = numpy.array(((0. , 1.7834, 1.7834),
(1.7834, 0. , 1.7834),
(1.7834, 1.7834, 0. ),)) + numpy.random.random((3,3)).T
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.mesh = [11,9,7]
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
coulG = tools.get_coulG(cell, kpt)
self.assertAlmostEqual(lib.finger(coulG), 62.75448804333378, 9)
cell.a = numpy.eye(3)
cell.unit = 'B'
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]))
self.assertAlmostEqual(lib.finger(coulG), 4.6737453679713905, 9)
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]),
wrap_around=False)
self.assertAlmostEqual(lib.finger(coulG), 4.5757877990664744, 9)
coulG = tools.get_coulG(cell, exx='ewald')
self.assertAlmostEqual(lib.finger(coulG), 4.888843468914021, 9)
def test_coulG(self):
numpy.random.seed(19)
kpt = numpy.random.random(3)
ase_atom = bulk('C', 'diamond', a=3.5668)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom)
cell.a = ase_atom.cell + numpy.random.random((3,3)).T
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.gs = [5,4,3]
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
coulG = tools.get_coulG(cell, kpt)
self.assertAlmostEqual(finger(coulG), 62.75448804333378, 9)
cell.a = numpy.eye(3)
cell.unit = 'B'
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]))
self.assertAlmostEqual(finger(coulG), 4.6737453679713905, 9)
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]),
wrap_around=False)
self.assertAlmostEqual(finger(coulG), 4.5757877990664744, 9)
coulG = tools.get_coulG(cell, exx='ewald')
self.assertAlmostEqual(finger(coulG), 4.888843468914021, 9)
def test_coulG(self):
numpy.random.seed(19)
kpt = numpy.random.random(3)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = 'C 0., 0., 0.; C 0.8917, 0.8917, 0.8917'
cell.a = numpy.array((
(0., 1.7834, 1.7834),
(1.7834, 0., 1.7834),
(1.7834, 1.7834, 0.),
)) + numpy.random.random((3, 3)).T
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.gs = [5, 4, 3]
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
coulG = tools.get_coulG(cell, kpt)
self.assertAlmostEqual(finger(coulG), 62.75448804333378, 9)
cell.a = numpy.eye(3)
cell.unit = 'B'
coulG = tools.get_coulG(cell, numpy.array([0, numpy.pi, 0]))
self.assertAlmostEqual(finger(coulG), 4.6737453679713905, 9)
coulG = tools.get_coulG(cell,
numpy.array([0, numpy.pi, 0]),
wrap_around=False)
self.assertAlmostEqual(finger(coulG), 4.5757877990664744, 9)
coulG = tools.get_coulG(cell, exx='ewald')
self.assertAlmostEqual(finger(coulG), 4.888843468914021, 9)
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 get_vkR_(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2):
'''Get the real-space 2-index "exchange" potential V_{i,k1; j,k2}(r).
Kwargs:
kpts : (nkpts, 3) ndarray
The sampled k-points; may be required for G=0 correction.
Returns:
vR : (ngs, nao, nao) ndarray
The real-space "exchange" potential at every grid point, for all
AO pairs.
Note:
This is essentially a density-fitting or resolution-of-the-identity.
The returned object is of size ngs*nao**2 and could be precomputed and
saved in vhfopt.
'''
coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell)
ngs, nao = aoR_k1.shape
coulG = tools.get_coulG(cell, kpt1-kpt2, exx=True, mf=mf)
vR = np.zeros((ngs,nao,nao), dtype=np.complex128)
for i in range(nao):
for j in range(nao):
rhoR = aoR_k1[:,i] * aoR_k2[:,j].conj()
rhoG = tools.fftk(rhoR, cell.gs, coords, kpt1-kpt2)
vG = coulG*rhoG
vR[:,i,j] = tools.ifftk(vG, cell.gs, coords, kpt1-kpt2)
return vR
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 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 assemble_eri(cell,
orb_pair_invG1,
orb_pair_G2,
q=None,
verbose=logger.INFO):
'''Assemble 4-index electron repulsion integrals.
Returns:
(nmo1*nmo2, nmo3*nmo4) ndarray
'''
log = logger.Logger
if isinstance(verbose, logger.Logger):
log = verbose
else:
log = logger.Logger(cell.stdout, verbose)
log.debug('Performing periodic ERI assembly of (%i, %i) ij,kl pairs',
orb_pair_invG1.shape[1], orb_pair_G2.shape[1])
if q is None:
q = np.zeros(3)
coulqG = tools.get_coulG(cell, -1.0 * q)
ngrids = orb_pair_invG1.shape[0]
Jorb_pair_G2 = np.einsum('g,gn->gn', coulqG,
orb_pair_G2) * (cell.vol / ngrids**2)
eri = np.dot(orb_pair_invG1.T, Jorb_pair_G2)
return eri
def generate_orbital_products(self, Q, X, Xaoik, Xaolj):
kpts = self.kpts
cell = self.cell
ngs = self.ngs
for K in range(self.part.nkk):
k1 = K + self.part.kk0
k2 = self.QKToK2[Q][k1]
if self.part.ij0 > self.nmo_pk[k1] * self.nmo_pk[k2]:
continue
i0 = self.part.ij0 // self.nmo_pk[k2]
iN = self.part.ijN // self.nmo_pk[k2]
if self.part.ijN % self.nmo_pk[k2] != 0:
iN += 1
if iN > self.nmo_pk[k1]:
iN = self.nmo_pk[k1]
pij = self.part.ij0 % self.nmo_pk[k2]
n_ = min(self.part.ijN,
self.nmo_pk[k1] * self.nmo_pk[k2]) - self.part.ij0
X_t = X[k1][:, i0:iN].copy()
Xaoik[K, :, 0:n_] = self.df.get_mo_pairs_G(
(X_t, X[k2].copy()), (kpts[k1], kpts[k2]),
(kpts[k2] - kpts[k1]),
compact=False)[:, pij:pij + n_]
Xaolj[K, :, :] = Xaoik[K, :, :]
coulG = tools.get_coulG(cell,
kpts[k2] - kpts[k1],
mesh=self.df.mesh)
Xaoik[K, :, :] *= (coulG * cell.vol / ngs**2).reshape(-1, 1)
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 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_j_kpts(mydf,
dm_kpts,
hermi=1,
kpts=numpy.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
rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
rhoG = rhoG[:, 0]
coulG = tools.get_coulG(cell,
mesh=cell.mesh,
low_dim_ft_type=mydf.low_dim_ft_type)
ngrids = coulG.size
vG = numpy.einsum('ng,g->ng', rhoG.reshape(-1, ngrids), coulG)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
vj_kpts = _get_j_pass2(mydf, vG, kpts_band)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def gen_orbital_products(cell, mydf, X, nmo_pk, ngs, part, kpts, nmo_max):
# left (includes v(G)) and right pair densities
nkpts = len(kpts)
try:
Xaoik = numpy.zeros((part.nkk, ngs, part.nij), dtype=numpy.complex128)
Xaolj = numpy.zeros((part.nkk, ngs, part.nij), dtype=numpy.complex128)
except:
mem = part.nkk * ngs * part.nij * 16 * 2 / 1024.0**3
print(" # Problems allocating memory. Trying to allocate: {}"
" GB.".format(mem))
sys.exit()
for k in range(part.nkk):
k1 = part.n2k1[k]
k2 = part.n2k2[k]
i0 = part.ij0 // nmo_pk[k2]
iN = part.ijN // nmo_pk[k2]
if part.ijN % nmo_pk[k2] != 0:
iN += 1
if iN > nmo_pk[k1]:
iN = nmo_pk[k1]
pij = part.ij0 % nmo_pk[k2]
n_ = min(part.ijN, nmo_pk[k1] * nmo_pk[k2]) - part.ij0
X_t = X[k1][:, i0:iN].copy()
Xaoik[k, :, 0:n_] = mydf.get_mo_pairs_G((X_t, X[k2]),
(kpts[k1], kpts[k2]),
kpts[k2] - kpts[k1],
compact=False)[:, pij:pij + n_]
X_t = None
Xaolj[k, :, :] = Xaoik[k, :, :]
coulG = tools.get_coulG(cell, kpts[k2] - kpts[k1], mesh=mydf.mesh)
Xaoik[k, :, :] *= (coulG * cell.vol / ngs**2).reshape(-1, 1)
t1 = time.clock()
return Xaoik, Xaolj
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 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 weighted_coulG(self, kpt=numpy.zeros(3), exx=False, gs=None):
cell = self.cell
if gs is None:
gs = self.gs
Gv, Gvbase, kws = cell.get_Gv_weights(gs)
coulG = tools.get_coulG(cell, kpt, exx, self, gs, Gv)
coulG *= kws
return coulG
def weighted_coulG(mydf, kpt=numpy.zeros(3), exx=False, mesh=None):
cell = mydf.cell
if mesh is None:
mesh = mydf.mesh
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
coulG = tools.get_coulG(cell, kpt, exx, mydf, mesh, Gv)
coulG *= kws
return coulG
def weighted_coulG(self, kpt=numpy.zeros(3), exx=False, gs=None):
cell = self.cell
if gs is None:
gs = self.gs
Gv, Gvbase, kws = cell.get_Gv_weights(gs)
coulG = tools.get_coulG(cell, kpt, exx, self, gs, Gv)
coulG *= kws
return coulG
def weighted_coulG(mydf, kpt=numpy.zeros(3), exx=False, mesh=None):
cell = mydf.cell
if mesh is None:
mesh = mydf.mesh
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
coulG = tools.get_coulG(cell, kpt, exx, mydf, mesh, Gv)
coulG *= kws
return coulG
def _aft_quad_integrals(mydf, dm, efg_nuc):
# Use AFTDF 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
log = lib.logger.new_logger(mydf)
t0 = t1 = (lib.logger.process_clock(), lib.logger.perf_counter())
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')
ngrids = numpy.prod(mesh)
rhoG = numpy.zeros(ngrids, dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts_lst,
max_memory=max_memory,
aosym='s1'):
rhoG[p0:p1] = numpy.einsum('kgpq,kqp->g',
aoaoks.reshape(nkpts, p1 - p0, nao, nao),
dm_kpts)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]' % (p0, p1), *t1)
t0 = log.timer_debug1('contracting Vnuc', *t0)
rhoG *= 1. / nkpts
Gv, Gvbase, kws = cell.get_Gv_weights(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. / cell.vol * numpy.einsum('g,g,gxy->gxy', rhoG, coulG, GG)
if mydf.eta == 0:
SI = cell.get_SI(Gv)
efg_e = numpy.einsum('zg,gxy->zxy', SI[efg_nuc], vG.conj()).real
else:
nuccell = aft._compensate_nuccell(mydf)
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
efg_e = _int_nuc_vloc(mydf, nuccell, kpts_lst, dm_kpts)
t0 = log.timer_debug1('vnuc pass1: analytic int', *t0)
aoaux = df.ft_ao.ft_ao(nuccell, Gv)
efg_e += numpy.einsum('gz,gxy->zxy', aoaux[:, efg_nuc], vG.conj()).real
return efg_e
def get_eri(mydf, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)):
cell = mydf.cell
nao = cell.nao_nr()
kptijkl = _format_kpts(kpts)
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'fft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros((nao,nao,nao,nao))
kpti, kptj, kptk, kptl = kptijkl
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coords = cell.gen_uniform_grids(mydf.mesh)
max_memory = mydf.max_memory - lib.current_memory()[0]
####################
# gamma point, the integral is real and with s4 symmetry
if gamma_point(kptijkl):
#:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=compact)
#:ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1)
#:eri = lib.dot(ao_pairs_G.T, ao_pairs_G, cell.vol/ngrids**2)
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
ao = numpy.asarray(ao.T, order='C')
eri = _contract_compact(mydf, (ao,ao), coulG, max_memory=max_memory)
if not compact:
eri = ao2mo.restore(1, eri, nao).reshape(nao**2,nao**2)
return eri
####################
# aosym = s1, complex integrals
else:
#:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=False)
#:# ao_pairs_invG = rho_kl(-(G+k_ij)) = conj(rho_lk(G+k_ij)).swap(r,s)
#:#=get_ao_pairs_G(mydf, [kptl,kptk], q, compact=False).transpose(0,2,1).conj()
#:ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], q, compact=False).conj()
#:ao_pairs_G *= coulG.reshape(-1,1)
#:eri = lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol/ngrids**2)
if is_zero(kpti-kptl) and is_zero(kptj-kptk):
if is_zero(kpti-kptj):
aoi = mydf._numint.eval_ao(cell, coords, kpti)[0]
aoi = aoj = numpy.asarray(aoi.T, order='C')
else:
aoi, aoj = mydf._numint.eval_ao(cell, coords, kptijkl[:2])
aoi = numpy.asarray(aoi.T, order='C')
aoj = numpy.asarray(aoj.T, order='C')
aos = (aoi, aoj, aoj, aoi)
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
aos = [numpy.asarray(x.T, order='C') for x in aos]
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - aos[0].nbytes*4*1e-6
eri = _contract_plain(mydf, aos, coulG, fac, max_memory=max_memory)
return eri
def make_kpt(kpt): # kpt = kptj - kpti
# search for all possible ki and kj that has ki-kj+kpt=0
kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
nkptj = len(kptj_idx)
log.debug1('kpt = %s', kpt)
log.debug1('kpti_idx = %s', kpti_idx)
log.debug1('kptj_idx = %s', kptj_idx)
kk_todo[kpti_idx, kptj_idx] = False
if swap_2e and abs(kpt).sum() > 1e-9:
kk_todo[kptj_idx, kpti_idx] = False
mydf.exxdiv = exxdiv
vkcoulG = tools.get_coulG(cell, kpt, True, mydf, mydf.gs) / cell.vol
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(cell, mydf.gs, kpt, kpts[kptj_idx],
max_memory=max_memory):
ki = kpti_idx[k]
kj = kptj_idx[k]
coulG = numpy.sqrt(vkcoulG[p0:p1])
# case 1: k_pq = (pi|iq)
pqkR *= coulG
pqkI *= coulG
rsk = (pqkR.reshape(nao, nao, -1).transpose(1, 0, 2) -
pqkI.reshape(nao, nao, -1).transpose(1, 0, 2) * 1j)
qpk = rsk.conj()
for i in range(nset):
qsk = lib.dot(dms[i, kj],
rsk.reshape(nao, -1)).reshape(nao, nao, -1)
#:vk_kpts[i,ki] += numpy.einsum('qpk,qsk->ps', qpk, qsk)
vk_kpts[i, ki] += lib.dot(
qpk.transpose(1, 0, 2).reshape(nao, -1),
qsk.transpose(1, 0, 2).reshape(nao, -1).T)
qsk = None
rsk = qpk = None
# case 2: k_pq = (iq|pi)
if swap_2e and abs(kpt).sum() > 1e-9:
srk = pqkR - pqkI * 1j
pqk = srk.reshape(nao, nao, -1).conj()
for i in range(nset):
prk = lib.dot(dms[i, ki].T,
srk.reshape(nao, -1)).reshape(nao, nao, -1)
#:vk_kpts[i,kj] += numpy.einsum('prk,pqk->rq', prk, pqk)
vk_kpts[i, kj] += lib.dot(
prk.transpose(1, 0, 2).reshape(nao, -1),
pqk.transpose(1, 0, 2).reshape(nao, -1).T)
prk = None
srk = pqk = None
pqkR = pqkI = coulG = None
return None
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 loop(self):
coulG = tools.get_coulG(self.cell, numpy.zeros(3), gs=mydf.gs)
ngs = len(coulG)
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact=True)
ao_pairs_G *= numpy.sqrt(coulG * (cell.vol / ngs**2)).reshape(-1, 1)
Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1]))
for p0, p1 in lib.prange(0, ngs, self.blockdim):
Lpq[:p1 - p0] = ao_pairs_G[p0:p1].real
yield Lpq[:p1 - p0]
Lpq[:p1 - p0] = ao_pairs_G[p0:p1].imag
yield Lpq[:p1 - p0]
def weighted_coulG_SR(self, kpt=np.zeros(3), exx=False, mesh=None):
cell = self.cell
if mesh is None:
mesh = self.mesh
if cell.omega != 0:
raise RuntimeError('RangeSeparationJKBuilder cannot be used to evaluate '
'the long-range HF exchange in RSH functional')
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
coulG = pbctools.get_coulG(cell, kpt, False, self, mesh, Gv,
omega=-self.omega)
coulG *= kws
return coulG
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 loop(self):
coulG = tools.get_coulG(self.cell, numpy.zeros(3), gs=mydf.gs)
ngs = len(coulG)
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact=True)
ao_pairs_G *= numpy.sqrt(coulG * (cell.vol / ngs ** 2)).reshape(-1, 1)
Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1]))
for p0, p1 in lib.prange(0, ngs, self.blockdim):
Lpq[: p1 - p0] = ao_pairs_G[p0:p1].real
yield Lpq[: p1 - p0]
Lpq[: p1 - p0] = ao_pairs_G[p0:p1].imag
yield Lpq[: p1 - p0]
def _aft_quad_integrals(mydf, dm, efg_nuc):
# Use AFTDF 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
log = lib.logger.new_logger(mydf)
t0 = t1 = (time.clock(), time.time())
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')
ngrids = numpy.prod(mesh)
rhoG = numpy.zeros(ngrids, dtype=numpy.complex128)
kpt_allow = numpy.zeros(3)
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_lst,
max_memory=max_memory, aosym='s1'):
rhoG[p0:p1] = numpy.einsum('kgpq,kqp->g', aoaoks.reshape(nkpts,p1-p0,nao,nao),
dm_kpts)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]'%(p0, p1), *t1)
t0 = log.timer_debug1('contracting Vnuc', *t0)
rhoG *= 1./nkpts
Gv, Gvbase, kws = cell.get_Gv_weights(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./cell.vol * numpy.einsum('g,g,gxy->gxy', rhoG, coulG, GG)
if mydf.eta == 0:
SI = cell.get_SI(Gv)
efg_e = numpy.einsum('zg,gxy->zxy', SI[efg_nuc], vG.conj()).real
else:
nuccell = aft._compensate_nuccell(mydf)
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
efg_e = _int_nuc_vloc(mydf, nuccell, kpts_lst, dm_kpts)
t0 = log.timer_debug1('vnuc pass1: analytic int', *t0)
aoaux = df.ft_ao.ft_ao(nuccell, Gv)
efg_e += numpy.einsum('gz,gxy->zxy', aoaux[:,efg_nuc], vG.conj()).real
return efg_e
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 general(mydf, mo_coeffs, kpts=None,
compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)):
'''General MO integral transformation'''
from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
warn_pbc2d_eri(mydf)
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs,) * 4
mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs]
if not _iskconserv(cell, kptijkl):
lib.logger.warn(cell, 'fft_ao2mo: momentum conservation not found in '
'the given k-points %s', kptijkl)
return numpy.zeros([mo.shape[1] for mo in mo_coeffs])
allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
q = kptj - kpti
coulG = tools.get_coulG(cell, q, mesh=mydf.mesh)
coords = cell.gen_uniform_grids(mydf.mesh)
max_memory = mydf.max_memory - lib.current_memory()[0]
if gamma_point(kptijkl) and allreal:
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
if ((iden_coeffs(mo_coeffs[0], mo_coeffs[1]) and
iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
iden_coeffs(mo_coeffs[0], mo_coeffs[3]))):
moiT = mojT = numpy.asarray(lib.dot(mo_coeffs[0].T,ao.T), order='C')
ao = None
max_memory = max_memory - moiT.nbytes*1e-6
eri = _contract_compact(mydf, (moiT,mojT), coulG, max_memory=max_memory)
if not compact:
nmo = moiT.shape[0]
eri = ao2mo.restore(1, eri, nmo).reshape(nmo**2,nmo**2)
else:
mos = [numpy.asarray(lib.dot(c.T, ao.T), order='C') for c in mo_coeffs]
ao = None
fac = numpy.array(1.)
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory).real
return eri
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
mos = [numpy.asarray(lib.dot(c.T, aos[i].T), order='C')
for i,c in enumerate(mo_coeffs)]
aos = None
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - sum([x.nbytes for x in mos])*1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory)
return eri
def general(mydf, mo_coeffs, kpts=None, compact=False):
'''General MO integral transformation'''
cell = mydf.cell
kptijkl = _format_kpts(kpts)
kpti, kptj, kptk, kptl = kptijkl
if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
mo_coeffs = (mo_coeffs, ) * 4
mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs]
allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs)
q = kptj - kpti
coulG = tools.get_coulG(cell, q, gs=mydf.gs)
coords = cell.gen_uniform_grids(mydf.gs)
max_memory = mydf.max_memory - lib.current_memory()[0]
if gamma_point(kptijkl) and allreal:
ao = mydf._numint.eval_ao(cell, coords, kpti)[0]
if ((iden_coeffs(mo_coeffs[0], mo_coeffs[2])
and iden_coeffs(mo_coeffs[1], mo_coeffs[3]))):
moiT = mojT = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T),
order='C')
ao = None
max_memory = max_memory - moiT.nbytes * 1e-6
eri = _contract_compact(mydf, (moiT, mojT),
coulG,
max_memory=max_memory)
if not compact:
nmo = moiT.shape[0]
eri = ao2mo.restore(1, eri, nmo).reshape(nmo**2, nmo**2)
else:
mos = [
numpy.asarray(lib.dot(c.T, ao.T), order='C') for c in mo_coeffs
]
ao = None
fac = numpy.array(1.)
max_memory = max_memory - sum([x.nbytes for x in mos]) * 1e-6
eri = _contract_plain(mydf, mos, coulG, fac,
max_memory=max_memory).real
return eri
else:
aos = mydf._numint.eval_ao(cell, coords, kptijkl)
mos = [
numpy.asarray(lib.dot(c.T, aos[i].T), order='C')
for i, c in enumerate(mo_coeffs)
]
aos = None
fac = numpy.exp(-1j * numpy.dot(coords, q))
max_memory = max_memory - sum([x.nbytes for x in mos]) * 1e-6
eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory)
return eri
def make_kpt(kpt): # kpt = kptj - kpti
# search for all possible ki and kj that has ki-kj+kpt=0
kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
nkptj = len(kptj_idx)
log.debug1('kpt = %s', kpt)
log.debug1('kpti_idx = %s', kpti_idx)
log.debug1('kptj_idx = %s', kptj_idx)
kk_todo[kpti_idx,kptj_idx] = False
if swap_2e and abs(kpt).sum() > 1e-9:
kk_todo[kptj_idx,kpti_idx] = False
mydf.exxdiv = exxdiv
vkcoulG = tools.get_coulG(cell, kpt, True, mydf, mydf.gs) / cell.vol
# <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(cell, mydf.gs, kpt, kpts[kptj_idx],
max_memory=max_memory):
ki = kpti_idx[k]
kj = kptj_idx[k]
coulG = numpy.sqrt(vkcoulG[p0:p1])
# case 1: k_pq = (pi|iq)
pqkR *= coulG
pqkI *= coulG
rsk =(pqkR.reshape(nao,nao,-1).transpose(1,0,2) -
pqkI.reshape(nao,nao,-1).transpose(1,0,2)*1j)
qpk = rsk.conj()
for i in range(nset):
qsk = lib.dot(dms[i,kj], rsk.reshape(nao,-1)).reshape(nao,nao,-1)
#:vk_kpts[i,ki] += numpy.einsum('qpk,qsk->ps', qpk, qsk)
vk_kpts[i,ki] += lib.dot(qpk.transpose(1,0,2).reshape(nao,-1),
qsk.transpose(1,0,2).reshape(nao,-1).T)
qsk = None
rsk = qpk = None
# case 2: k_pq = (iq|pi)
if swap_2e and abs(kpt).sum() > 1e-9:
srk = pqkR - pqkI*1j
pqk = srk.reshape(nao,nao,-1).conj()
for i in range(nset):
prk = lib.dot(dms[i,ki].T, srk.reshape(nao,-1)).reshape(nao,nao,-1)
#:vk_kpts[i,kj] += numpy.einsum('prk,pqk->rq', prk, pqk)
vk_kpts[i,kj] += lib.dot(prk.transpose(1,0,2).reshape(nao,-1),
pqk.transpose(1,0,2).reshape(nao,-1).T)
prk = None
srk = pqk = None
pqkR = pqkI = coulG = None
return None
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 weighted_coulG(mydf, omega, kpt=np.zeros(3), exx=False, mesh=None):
cell = mydf.cell
if cell.omega != 0:
raise RuntimeError('RSGDF cannot be used '
'to evaluate the long-range HF exchange in RSH '
'functional')
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
if abs(omega) < 1.e-10:
omega_ = None
else:
omega_ = omega
coulG = pbctools.get_coulG(cell, kpt, False, mydf, mesh, Gv, omega=omega_)
coulG *= kws
return coulG
def get_eri(mydf, kpts=None, compact=False):
cell = mydf.cell
if kpts is None:
kptijkl = numpy.zeros((4, 3))
elif numpy.shape(kpts) == (3, ):
kptijkl = numpy.vstack([kpts] * 4)
else:
kptijkl = numpy.reshape(kpts, (4, 3))
kpti, kptj, kptk, kptl = kptijkl
nao = cell.nao_nr()
nao_pair = nao * (nao + 1) // 2
coulG = tools.get_coulG(cell, kptj - kpti, gs=mydf.gs)
ngs = len(coulG)
####################
# gamma point, the integral is real and with s4 symmetry
if abs(kptijkl).sum() < 1e-9:
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact)
ao_pairs_G *= numpy.sqrt(coulG).reshape(-1, 1)
aoijR = ao_pairs_G.real.copy()
aoijI = ao_pairs_G.imag.copy()
ao_pairs_G = None
eri = lib.dot(aoijR.T, aoijR, cell.vol / ngs**2)
eri = lib.dot(aoijI.T, aoijI, cell.vol / ngs**2, eri, 1)
return eri
####################
# (kpt) i == j == k == l != 0
# (kpt) i == l && j == k && i != j && j != k =>
#
# complex integrals, N^4 elements
elif (abs(kpti - kptl).sum() < 1e-9) and (abs(kptj - kptk).sum() < 1e-9):
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], False)
ao_pairs_G *= numpy.sqrt(coulG).reshape(-1, 1)
ao_pairs_invG = ao_pairs_G.T.reshape(nao, nao, -1).transpose(1, 0,
2).conj()
ao_pairs_invG = ao_pairs_invG.reshape(-1, ngs)
return lib.dot(ao_pairs_G.T, ao_pairs_invG.T, cell.vol / ngs**2)
####################
# aosym = s1, complex integrals
#
else:
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], False)
# ao_pairs_invG = rho_rs(-G+k_rs) = conj(rho_sr(G+k_sr)).swap(r,s)
ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], False).conj()
ao_pairs_G *= coulG.reshape(-1, 1)
return lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol / ngs**2)
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_(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 assemble_eri(cell, orb_pair_invG1, orb_pair_G2, q=None):
'''Assemble 4-index electron repulsion integrals.
Returns:
(nmo1*nmo2, nmo3*nmo4) ndarray
'''
if q is None:
q = np.zeros(3)
coulqG = tools.get_coulG(cell, -1.0*q)
ngs = orb_pair_invG1.shape[0]
Jorb_pair_G2 = np.einsum('g,gn->gn',coulqG,orb_pair_G2)*(cell.vol/ngs**2)
eri = np.dot(orb_pair_invG1.T, Jorb_pair_G2)
return eri
def assemble_eri(cell, orb_pair_invG1, orb_pair_G2, q=None):
'''Assemble 4-index electron repulsion integrals.
Returns:
(nmo1*nmo2, nmo3*nmo4) ndarray
'''
if q is None:
q = np.zeros(3)
coulqG = tools.get_coulG(cell, -1.0*q)
ngrids = orb_pair_invG1.shape[0]
Jorb_pair_G2 = np.einsum('g,gn->gn',coulqG,orb_pair_G2)*(cell.vol/ngrids**2)
eri = np.dot(orb_pair_invG1.T, Jorb_pair_G2)
return eri
def get_eri(mydf, kpts=None, compact=False):
cell = mydf.cell
if kpts is None:
kptijkl = numpy.zeros((4,3))
elif numpy.shape(kpts) == (3,):
kptijkl = numpy.vstack([kpts]*4)
else:
kptijkl = numpy.reshape(kpts, (4,3))
kpti, kptj, kptk, kptl = kptijkl
nao = cell.nao_nr()
nao_pair = nao * (nao+1) // 2
coulG = tools.get_coulG(cell, kptj-kpti, gs=mydf.gs)
ngs = len(coulG)
####################
# gamma point, the integral is real and with s4 symmetry
if abs(kptijkl).sum() < 1e-9:
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact)
ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1)
aoijR = ao_pairs_G.real.copy()
aoijI = ao_pairs_G.imag.copy()
ao_pairs_G = None
eri = lib.dot(aoijR.T, aoijR, cell.vol/ngs**2)
eri = lib.dot(aoijI.T, aoijI, cell.vol/ngs**2, eri, 1)
return eri
####################
# (kpt) i == j == k == l != 0
# (kpt) i == l && j == k && i != j && j != k =>
#
# complex integrals, N^4 elements
elif (abs(kpti-kptl).sum() < 1e-9) and (abs(kptj-kptk).sum() < 1e-9):
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], False)
ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1)
ao_pairs_invG = ao_pairs_G.T.reshape(nao,nao,-1).transpose(1,0,2).conj()
ao_pairs_invG = ao_pairs_invG.reshape(-1,ngs)
return lib.dot(ao_pairs_G.T, ao_pairs_invG.T, cell.vol/ngs**2)
####################
# aosym = s1, complex integrals
#
else:
ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], False)
# ao_pairs_invG = rho_rs(-G+k_rs) = conj(rho_sr(G+k_sr)).swap(r,s)
ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], False).conj()
ao_pairs_G *= coulG.reshape(-1,1)
return lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol/ngs**2)
def test_coulG_ws(self):
ase_atom = bulk('C', 'diamond', a=3.5668)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom)
cell.a = ase_atom.cell
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.gs = [5] * 3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
mf = khf.KRHF(cell, exxdiv='vcut_ws')
mf.kpts = cell.make_kpts([2, 2, 2])
coulG = tools.get_coulG(cell, mf.kpts[2], True, mf)
self.assertAlmostEqual(finger(coulG), 1.3245117871351604, 9)
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 test_coulG_ws(self):
ase_atom = bulk('C', 'diamond', a=3.5668)
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom)
cell.a = ase_atom.cell
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.gs = [5]*3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
mf = khf.KRHF(cell, exxdiv='vcut_ws')
mf.kpts = cell.make_kpts([2,2,2])
coulG = tools.get_coulG(cell, mf.kpts[2], True, mf)
self.assertAlmostEqual(finger(coulG), 1.3245117871351604, 9)
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_coulG_ws(self):
cell = pbcgto.Cell()
cell.unit = 'A'
cell.atom = 'C 0., 0., 0.; C 0.8917, 0.8917, 0.8917'
cell.a = '''0. 1.7834 1.7834
1.7834 0. 1.7834
1.7834 1.7834 0. '''
cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade'
cell.mesh = [11]*3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()
mf = khf.KRHF(cell, exxdiv='vcut_ws')
mf.kpts = cell.make_kpts([2,2,2])
coulG = tools.get_coulG(cell, mf.kpts[2], True, mf, gs=[5,5,5])
self.assertAlmostEqual(lib.finger(coulG), 1.3245365170998518+0j, 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 get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2):
coords = gen_grid.gen_uniform_grids(cell)
ngs, nao = aoR_k1.shape
expmikr = np.exp(-1j*np.dot(kpt1-kpt2,coords.T))
coulG = tools.get_coulG(cell, kpt1-kpt2, exx=True, mf=mf)
def prod(ij):
i, j = divmod(ij, nao)
rhoR = aoR_k1[:,i] * aoR_k2[:,j].conj()
rhoG = tools.fftk(rhoR, cell.gs, expmikr)
vG = coulG*rhoG
vR = tools.ifftk(vG, cell.gs, expmikr.conj())
return vR
if aoR_k1.dtype == np.double and aoR_k2.dtype == np.double:
vR = numpy.asarray([prod(ij).real for ij in range(nao**2)])
else:
vR = numpy.asarray([prod(ij) for ij in range(nao**2)])
return vR.reshape(nao,nao,-1).transpose(2,0,1)
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 get_nuc(mydf, kpts=None):
log = logger.Logger(mydf.stdout, mydf.verbose)
t1 = t0 = (time.clock(), time.time())
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
nkpts = len(kpts_lst)
nao = cell.nao_nr()
charge = -cell.atom_charges()
Gv = cell.get_Gv(mydf.gs)
SI = cell.get_SI(Gv)
rhoG = numpy.dot(charge, SI)
kpt_allow = numpy.zeros(3)
real = gamma_point(kpts_lst)
coulG = tools.get_coulG(cell, kpt_allow, gs=mydf.gs, Gv=Gv) / cell.vol
vneG = rhoG * coulG
if real:
vne = numpy.zeros((nkpts,nao**2))
else:
vne = numpy.zeros((nkpts,nao**2), dtype=numpy.complex128)
max_memory = mydf.max_memory - lib.current_memory()[0]
for k, pqkR, pqkI, p0, p1 \
in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts_lst,
max_memory=max_memory):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
vG = vneG[p0:p1]
if not real:
vne[k] += numpy.einsum('k,xk->x', vG.real, pqkI) * 1j
vne[k] += numpy.einsum('k,xk->x', vG.imag, pqkR) *-1j
vne[k] += numpy.einsum('k,xk->x', vG.real, pqkR)
vne[k] += numpy.einsum('k,xk->x', vG.imag, pqkI)
vne = vne.reshape(-1,nao,nao)
t1 = log.timer_debug1('contracting Vnuc', *t1)
if kpts is None or numpy.shape(kpts) == (3,):
vne = vne[0]
return vne
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 assemble_eri(cell, orb_pair_invG1, orb_pair_G2, q=None, verbose=logger.INFO):
'''Assemble 4-index electron repulsion integrals.
Returns:
(nmo1*nmo2, nmo3*nmo4) ndarray
'''
log = logger.Logger
if isinstance(verbose, logger.Logger):
log = verbose
else:
log = logger.Logger(cell.stdout, verbose)
log.debug('Performing periodic ERI assembly of (%i, %i) ij,kl pairs',
orb_pair_invG1.shape[1], orb_pair_G2.shape[1])
if q is None:
q = np.zeros(3)
coulqG = tools.get_coulG(cell, -1.0*q)
ngrids = orb_pair_invG1.shape[0]
Jorb_pair_G2 = np.einsum('g,gn->gn',coulqG,orb_pair_G2)*(cell.vol/ngrids**2)
eri = np.dot(orb_pair_invG1.T, Jorb_pair_G2)
return eri
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 _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 _ewald_exxdiv_for_G0(mf, dm, kpt, vk):
cell = mf.cell
gs = (0, 0, 0)
Gv = np.zeros((1, 3))
ovlp = mf.get_ovlp(cell, kpt)
coulGk = tools.get_coulG(cell, kpt - kpt, True, mf, gs, Gv)[0]
logger.debug(
mf, "Total energy shift = -1/2 * Nelec*madelung/cell.vol = %.12g", coulGk / cell.vol * cell.nelectron * -0.5
)
if isinstance(dm, np.ndarray) and dm.ndim == 2:
vk += coulGk / cell.vol * reduce(np.dot, (ovlp, dm, ovlp))
nelec = np.einsum("ij,ij", ovlp, dm)
else:
nelec = 0
for k, dmi in enumerate(dm):
vk[k] += coulGk / cell.vol * reduce(np.dot, (ovlp, dmi, ovlp))
nelec += np.einsum("ij,ij", ovlp, dmi)
# if abs(nelec - cell.nelectron) > .1 and abs(nelec) > .1:
# logger.debug(mf, 'Tr(dm,S) = %g', nelec)
# sys.stderr.write('Warning: The input dm is not SCF density matrix. '
# 'The Ewald treatment on G=0 term for K matrix '
# 'might not be well defined. Set mf.exxdiv=None '
# 'to switch off the Ewald term.\n')
return vk
def get_vkR(mydf, cell, aoR_k1, aoR_k2, kpt1, kpt2, coords, gs, exxdiv):
"""Get the real-space 2-index "exchange" potential V_{i,k1; j,k2}(r)
where {i,k1} = exp^{i k1 r) |i> , {j,k2} = exp^{-i k2 r) <j|
Kwargs:
kpt1, kpt2 : (3,) ndarray
The sampled k-points; may be required for G=0 correction.
Returns:
vR : (ngs, nao, nao) ndarray
The real-space "exchange" potential at every grid point, for all
AO pairs.
Note:
This is essentially a density-fitting or resolution-of-the-identity.
The returned object is of size ngs*nao**2
"""
ngs, nao = aoR_k1.shape
expmikr = np.exp(-1j * np.dot(kpt1 - kpt2, coords.T))
mydf.exxdiv = exxdiv
coulG = tools.get_coulG(cell, kpt1 - kpt2, True, mydf, gs)
aoR_k1 = np.asarray(aoR_k1.T, order="C")
aoR_k2 = np.asarray(aoR_k2.T, order="C")
vR = np.empty((nao, nao, ngs), dtype=np.complex128)
for i in range(nao):
rhoR = aoR_k1 * aoR_k2[i].conj()
rhoG = tools.fftk(rhoR, gs, expmikr)
vG = rhoG * coulG
vR[:, i] = tools.ifftk(vG, gs, expmikr.conj())
vR = vR.transpose(2, 0, 1)
if aoR_k1.dtype == np.double and aoR_k2.dtype == np.double:
return vR.real
else:
return vR