#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
r'''
Integral transformation with FFT
(ij|kl) = \int dr1 dr2 i*(r1) j(r1) v(r12) k*(r2) l(r2)
= (ij|G) v(G) (G|kl)
i*(r) j(r) = 1/N \sum_G e^{iGr} (G|ij)
= 1/N \sum_G e^{-iGr} (ij|G)
"forward" FFT:
(G|ij) = \sum_r e^{-iGr} i*(r) j(r) = fft[ i*(r) j(r) ]
"inverse" FFT:
(ij|G) = \sum_r e^{iGr} i*(r) j(r) = N * ifft[ i*(r) j(r) ]
= conj[ \sum_r e^{-iGr} j*(r) i(r) ]
'''
import numpy
from pyscf import lib
from pyscf import ao2mo
from pyscf.ao2mo.incore import iden_coeffs
from pyscf.pbc import tools
from pyscf.pbc.lib import kpts_helper
from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point, unique
from pyscf import __config__
[docs]
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
[docs]
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
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
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 _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
[docs]
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
[docs]
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
[docs]
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None):
cell = mydf.cell
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
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 _format_kpts(kpts):
if kpts is None:
kptijkl = numpy.zeros((4,3))
else:
kpts = numpy.asarray(kpts)
if kpts.size == 3:
kptijkl = numpy.vstack([kpts]*4).reshape(4,3)
else:
kptijkl = kpts.reshape(4,3)
return kptijkl
def _iskconserv(cell, kpts):
dk = kpts[1] - kpts[0] + kpts[3] - kpts[2]
if abs(dk).sum() < 1e-9:
return True
else:
s = 1./(2*numpy.pi)*numpy.dot(dk, cell.lattice_vectors().T)
s_int = s.round(0)
return abs(s - s_int).sum() < 1e-9
if __name__ == '__main__':
import pyscf.pbc.gto as pgto
from pyscf.pbc import df
L = 5.
n = 11
cell = pgto.Cell()
cell.a = numpy.diag([L,L,L])
cell.mesh = numpy.array([n,n,n])
cell.atom = '''He 3. 2. 3.
He 1. 1. 1.'''
#cell.basis = {'He': [[0, (1.0, 1.0)]]}
#cell.basis = '631g'
#cell.basis = {'He': [[0, (2.4, 1)], [1, (1.1, 1)]]}
cell.basis = 'ccpvdz'
cell.verbose = 0
cell.build(0,0)
nao = cell.nao_nr()
numpy.random.seed(1)
kpts = numpy.random.random((4,3))
kpts[3] = -numpy.einsum('ij->j', kpts[:3])
with_df = df.FFTDF(cell)
with_df.kpts = kpts
mo =(numpy.random.random((nao,nao)) +
numpy.random.random((nao,nao))*1j)
eri = with_df.get_eri(kpts).reshape((nao,)*4)
eri0 = numpy.einsum('pjkl,pi->ijkl', eri , mo.conj())
eri0 = numpy.einsum('ipkl,pj->ijkl', eri0, mo )
eri0 = numpy.einsum('ijpl,pk->ijkl', eri0, mo.conj())
eri0 = numpy.einsum('ijkp,pl->ijkl', eri0, mo ).reshape(nao**2,-1)
eri1 = with_df.ao2mo(mo, kpts)
print(abs(eri1-eri0.reshape(eri1.shape)).sum())
