def __init__(self,
atombases: List[AtomCGTOBasis],
spherical: bool = True,
df: Optional[DensityFitInfo] = None,
efield: Optional[torch.Tensor] = None,
cache: Optional[Cache] = None) -> None:
self.atombases = atombases
self.spherical = spherical
self.libcint_wrapper = intor.LibcintWrapper(atombases, spherical)
self.dtype = self.libcint_wrapper.dtype
self.device = self.libcint_wrapper.device
if df is None:
self._df: Optional[DFMol] = None
else:
self._df = DFMol(df, wrapper=self.libcint_wrapper)
self._efield = efield
if efield is not None:
assert efield.ndim == 1 and efield.numel() == 3
self.is_grid_set = False
self.is_ao_set = False
self.is_grad_ao_set = False
self.is_lapl_ao_set = False
self.xc: Optional[BaseXC] = None
self.xcfamily = 1
self.is_built = False
# initialize cache
self._cache = cache if cache is not None else Cache.get_dummy()
self._cache.add_cacheable_params(
["overlap", "kinetic", "nuclattr", "efield"])
if self._df is None:
self._cache.add_cacheable_params(["elrep"])
def build(self) -> BaseDF:
self._is_built = True
# construct the matrix used to calculate the electron repulsion for
# density fitting method
method = self.dfinfo.method
auxbasiswrapper = intor.LibcintWrapper(self.dfinfo.auxbases,
spherical=self.wrapper.spherical)
basisw, auxbw = intor.LibcintWrapper.concatenate(self.wrapper, auxbasiswrapper)
if method == "coulomb":
logger.log("Calculating the 2e2c integrals")
j2c = intor.coul2c(auxbw) # (nxao, nxao)
logger.log("Calculating the 2e3c integrals")
j3c = intor.coul3c(basisw, other1=basisw,
other2=auxbw) # (nao, nao, nxao)
elif method == "overlap":
j2c = intor.overlap(auxbw) # (nxao, nxao)
# TODO: implement overlap3c
raise NotImplementedError(
"Density fitting with overlap minimization is not implemented")
self._j2c = j2c # (nxao, nxao)
self._j3c = j3c # (nao, nao, nxao)
logger.log("Precompute matrix for density fittings")
self._inv_j2c = torch.inverse(j2c)
# if the memory is too big, then don't precompute elmat
if get_memory(j3c) > config.THRESHOLD_MEMORY:
self._precompute_elmat = False
else:
self._precompute_elmat = True
self._el_mat = torch.matmul(j3c, self._inv_j2c) # (nao, nao, nxao)
logger.log("Density fitting done")
return self
def build(self) -> BaseDF:
self._is_built = True
# construct the matrix used to calculate the electron repulsion for
# density fitting method
method = self.dfinfo.method
auxbasiswrapper = intor.LibcintWrapper(
self.dfinfo.auxbases, spherical=self.wrapper.spherical)
basisw, auxbw = intor.LibcintWrapper.concatenate(
self.wrapper, auxbasiswrapper)
if method == "coulomb":
j2c = intor.coul2c(auxbw) # (nxao, nxao)
j3c = intor.coul3c(basisw, other1=basisw,
other2=auxbw) # (nao, nao, nxao)
elif method == "overlap":
j2c = intor.overlap(auxbw) # (nxao, nxao)
# TODO: implement overlap3c
raise NotImplementedError(
"Density fitting with overlap minimization is not implemented")
self._j2c = j2c # (nxao, nxao)
self._j3c = j3c # (nao, nao, nxao)
self._inv_j2c = torch.inverse(j2c)
self._el_mat = torch.matmul(j3c, self._inv_j2c) # (nao, nao, nxao)
return self
def __init__(
self,
atombases: List[AtomCGTOBasis],
latt: Lattice,
*,
kpts: Optional[torch.Tensor] = None,
wkpts: Optional[
torch.Tensor] = None, # weights of k-points to get the density
spherical: bool = True,
df: Optional[DensityFitInfo] = None,
lattsum_opt: Optional[Union[intor.PBCIntOption, Dict]] = None,
cache: Optional[Cache] = None) -> None:
self._atombases = atombases
self._spherical = spherical
self._lattice = latt
# alpha for the compensating charge
# TODO: calculate eta properly or put it in lattsum_opt
self._eta = 0.2
self._eta = 0.46213127322256375 # temporary to follow pyscf.df
# lattice sum integral options
self._lattsum_opt = intor.PBCIntOption.get_default(lattsum_opt)
self._basiswrapper = intor.LibcintWrapper(atombases,
spherical=spherical,
lattice=latt)
self.dtype = self._basiswrapper.dtype
self.cdtype = get_complex_dtype(self.dtype)
self.device = self._basiswrapper.device
# set the default k-points and their weights
self._kpts = kpts if kpts is not None else \
torch.zeros((1, 3), dtype=self.dtype, device=self.device)
nkpts = self._kpts.shape[0]
# default weights are just 1/nkpts (nkpts,)
self._wkpts = wkpts if wkpts is not None else \
torch.ones((nkpts,), dtype=self.dtype, device=self.device) / nkpts
assert self._wkpts.shape[0] == self._kpts.shape[0]
assert self._wkpts.ndim == 1
assert self._kpts.ndim == 2
# initialize cache
self._cache = cache if cache is not None else Cache.get_dummy()
self._cache.add_cacheable_params(["overlap", "kinetic", "nuclattr"])
if df is None:
self._df: Optional[BaseDF] = None
else:
self._df = DFPBC(dfinfo=df,
wrapper=self._basiswrapper,
kpts=self._kpts,
wkpts=self._wkpts,
eta=self._eta,
lattsum_opt=self._lattsum_opt,
cache=self._cache.add_prefix("df"))
self._is_built = False
def _calc_nucl_attr(self) -> torch.Tensor:
# calculate the nuclear attraction matrix
# this follows the equation (31) in Sun, et al., J. Chem. Phys. 147 (2017)
# construct the fake nuclei atombases for nuclei
# (in this case, we assume each nucleus is a very sharp s-type orbital)
nucl_atbases = self._create_fake_nucl_bases(alpha=1e16, chargemult=1)
# add a compensating charge
cnucl_atbases = self._create_fake_nucl_bases(alpha=self._eta,
chargemult=-1)
# real charge + compensating charge
nucl_atbases_all = nucl_atbases + cnucl_atbases
nucl_wrapper = intor.LibcintWrapper(nucl_atbases_all,
spherical=self._spherical,
lattice=self._lattice)
cnucl_wrapper = intor.LibcintWrapper(cnucl_atbases,
spherical=self._spherical,
lattice=self._lattice)
natoms = nucl_wrapper.nao() // 2
# construct the k-points ij
# duplicating kpts to have shape of (nkpts, 2, ndim)
kpts_ij = self._kpts.unsqueeze(-2) * torch.ones(
(2, 1), dtype=self.dtype, device=self.device)
############# 1st part of nuclear attraction: short range #############
# get the 1st part of the nuclear attraction: the charge and compensating charge
# nuc1: (nkpts, nao, nao, 2 * natoms)
# nuc1 is not hermitian
basiswrapper1, nucl_wrapper1 = intor.LibcintWrapper.concatenate(
self._basiswrapper, nucl_wrapper)
nuc1_c = intor.pbc_coul3c(basiswrapper1,
other1=basiswrapper1,
other2=nucl_wrapper1,
kpts_ij=kpts_ij,
options=self._lattsum_opt)
nuc1 = -nuc1_c[..., :natoms] + nuc1_c[..., natoms:]
nuc1 = torch.sum(nuc1, dim=-1) # (nkpts, nao, nao)
# add vbar for 3 dimensional cell
# vbar is the interaction between the background charge and the
# compensating function.
# https://github.com/pyscf/pyscf/blob/c9aa2be600d75a97410c3203abf35046af8ca615/pyscf/pbc/df/aft.py#L239
nucbar = sum([-atb.atomz / self._eta for atb in self._atombases])
nuc1_b = -nucbar * np.pi / self._lattice.volume() * self._olp_mat
nuc1 = nuc1 + nuc1_b
############# 2nd part of nuclear attraction: long range #############
# get the 2nd part from the Fourier Transform
# get the G-points, choosing min because the two FTs are multiplied
gcut = get_gcut(self._lattsum_opt.precision,
wrappers=[cnucl_wrapper, self._basiswrapper],
reduce="min")
# gvgrids: (ngv, ndim), gvweights: (ngv,)
gvgrids, gvweights = self._lattice.get_gvgrids(gcut)
# the compensating charge's Fourier Transform
# TODO: split gvgrids and gvweights to reduce the memory usage
cnucl_ft = intor.eval_gto_ft(cnucl_wrapper, gvgrids) # (natoms, ngv)
# overlap integral of the electron basis' Fourier Transform
cbas_ft = intor.pbcft_overlap(
self._basiswrapper,
gvgrid=-gvgrids,
kpts=self._kpts,
options=self._lattsum_opt) # (nkpts, nao, nao, ngv)
# coulomb kernel Fourier Transform
coul_ft = unweighted_coul_ft(gvgrids) * gvweights # (ngv,)
coul_ft = coul_ft.to(cbas_ft.dtype) # cast to complex
# optimized by opt_einsum
# nuc2 = -torch.einsum("tg,kabg,g->kab", cnucl_ft, cbas_ft, coul_ft)
nuc2_temp = torch.einsum("g,tg->g", coul_ft, cnucl_ft)
nuc2 = -torch.einsum("g,kabg->kab", nuc2_temp,
cbas_ft) # (nkpts, nao, nao)
# print((nuc2 - nuc2.conj().transpose(-2, -1)).abs().max()) # check hermitian-ness
# get the total contribution from the short range and long range
nuc = nuc1 + nuc2
# symmetrize for more stable numerical calculation
nuc = (nuc + nuc.conj().transpose(-2, -1)) * 0.5
return nuc
def build(self) -> BaseDF:
self._is_built = True
df = self._dfinfo
# calculate the matrices required to calculate the electron repulsion operator
# i.e. the 3-centre 2-electron integrals (short + long range) and j3c @ (j2c^-1)
method = df.method.lower()
df_auxbases = _renormalize_auxbases(df.auxbases)
aux_comp_bases = self._create_compensating_bases(df_auxbases,
eta=self._eta)
fuse_aux_bases = df_auxbases + aux_comp_bases
fuse_aux_wrapper = intor.LibcintWrapper(
fuse_aux_bases,
spherical=self._wrapper.spherical,
lattice=self._lattice)
aux_comp_wrapper = intor.LibcintWrapper(
aux_comp_bases,
spherical=self._wrapper.spherical,
lattice=self._lattice)
aux_wrapper = intor.LibcintWrapper(df_auxbases,
spherical=self._wrapper.spherical,
lattice=self._lattice)
nxcao = aux_comp_wrapper.nao(
) # number of aux compensating basis wrapper
nxao = fuse_aux_wrapper.nao() - nxcao # number of aux basis wrapper
assert nxcao == nxao
# only gaussian density fitting is implemented at the moment
if method != "gdf":
raise NotImplementedError(
"Density fitting %s is not implemented (only gdf)" % df.method)
# get the k-points needed for the integrations
nkpts = self._kpts.shape[0]
kpts_ij = _combine_kpts_to_kpts_ij(self._kpts) # (nkpts_ij, 2, ndim)
kpts_reduce = _reduce_kpts_ij(kpts_ij) # (nkpts_ij, ndim)
nkpts_ij = kpts_ij.shape[0]
kpts_j = kpts_ij[..., 1, :] # (nkpts_ij, ndim)
def _calc_integrals():
######################## short-range integrals ########################
############# 3-centre 2-electron integral #############
_basisw, _fusew = intor.LibcintWrapper.concatenate(
self._wrapper, fuse_aux_wrapper)
# (nkpts_ij, nao, nao, nxao+nxcao)
j3c_short_f = intor.pbc_coul3c(_basisw,
other2=_fusew,
kpts_ij=kpts_ij,
options=self._lattsum_opt)
j3c_short = j3c_short_f[..., :nxao] - j3c_short_f[
..., nxao:] # (nkpts_ij, nao, nao, nxao)
############# 2-centre 2-electron integrals #############
# (nkpts_unique, nxao+nxcao, nxao+nxcao)
j2c_short_f = intor.pbc_coul2c(fuse_aux_wrapper,
kpts=kpts_reduce,
options=self._lattsum_opt)
# j2c_short: (nkpts_unique, nxao, nxao)
j2c_short = j2c_short_f[..., :nxao, :nxao] + j2c_short_f[..., nxao:, nxao:] \
- j2c_short_f[..., :nxao, nxao:] - j2c_short_f[..., nxao:, :nxao]
######################## long-range integrals ########################
# only use the compensating wrapper as the gcut
gcut = get_gcut(self._lattsum_opt.precision, [aux_comp_wrapper])
# gvgrids: (ngv, ndim), gvweights: (ngv,)
gvgrids, gvweights = self._lattice.get_gvgrids(gcut)
ngv = gvgrids.shape[0]
gvk = gvgrids.unsqueeze(-2) + kpts_reduce # (ngv, nkpts_ij, ndim)
gvk = gvk.view(-1, gvk.shape[-1]) # (ngv * nkpts_ij, ndim)
# get the fourier transform variables
# TODO: iterate over ngv axis
# ft of the compensating basis
comp_ft = intor.eval_gto_ft(aux_comp_wrapper,
gvk) # (nxcao, ngv * nkpts_ij)
comp_ft = comp_ft.view(-1, ngv, nkpts_ij) # (nxcao, ngv, nkpts_ij)
# ft of the auxiliary basis
auxb_ft_c = intor.eval_gto_ft(aux_wrapper,
gvk) # (nxao, ngv * nkpts_ij)
auxb_ft_c = auxb_ft_c.view(-1, ngv,
nkpts_ij) # (nxao, ngv, nkpts_ij)
auxb_ft = auxb_ft_c - comp_ft # (nxao, ngv, nkpts_ij)
# ft of the overlap integral of the basis (nkpts_ij, nao, nao, ngv)
aoao_ft = self._get_pbc_overlap_with_kpts_ij(
gvgrids, kpts_reduce, kpts_j)
# ft of the coulomb kernel
coul_ft = unweighted_coul_ft(gvk) # (ngv * nkpts_ij,)
coul_ft = coul_ft.to(comp_ft.dtype).view(
ngv, nkpts_ij) * gvweights.unsqueeze(-1) # (ngv, nkpts_ij)
# 1: (nkpts_ij, nxao, nxao)
pattern = "gi,xgi,ygi->ixy"
j2c_long = torch.einsum(pattern, coul_ft, comp_ft.conj(), auxb_ft)
# 2: (nkpts_ij, nxao, nxao)
j2c_long += torch.einsum(pattern, coul_ft, auxb_ft.conj(), comp_ft)
# 3: (nkpts_ij, nxao, nxao)
j2c_long += torch.einsum(pattern, coul_ft, comp_ft.conj(), comp_ft)
# calculate the j3c long-range
patternj3 = "gi,xgi,iyzg->iyzx"
# (nkpts_ij, nao, nao, nxao)
j3c_long = torch.einsum(patternj3, coul_ft, comp_ft.conj(),
aoao_ft)
# get the average potential
auxbar_f = self._auxbar(
kpts_reduce, fuse_aux_wrapper) # (nkpts_ij, nxao + nxcao)
auxbar = auxbar_f[:, :nxao] - auxbar_f[:,
nxao:] # (nkpts_ij, nxao)
auxbar = auxbar.reshape(nkpts, nkpts,
auxbar.shape[-1]) # (nkpts, nkpts, nxao)
olp_mat = intor.pbc_overlap(
self._wrapper, kpts=self._kpts,
options=self._lattsum_opt) # (nkpts, nao, nao)
j3c_bar = auxbar[:, :, None, None, :] * olp_mat[
..., None] # (nkpts, nkpts, nao, nao, nxao)
j3c_bar = j3c_bar.reshape(
-1, *j3c_bar.shape[2:]) # (nkpts_ij, nao, nao, nxao)
######################## combining integrals ########################
j2c = j2c_short + j2c_long # (nkpts_ij, nxao, nxao)
j3c = j3c_short + j3c_long - j3c_bar # (nkpts_ij, nao, nao, nxao)
el_mat = torch.einsum("kxy,kaby->kabx", torch.inverse(j2c),
j3c) # (nkpts_ij, nao, nao, nxao)
return j2c, j3c, el_mat
with self._cache.open():
# check the signature
self._cache.check_signature({
"dfinfo":
self._dfinfo,
"kpts":
self._kpts.detach(),
"wkpts":
self._wkpts.detach(),
"atombases":
self._wrapper.atombases,
"alattice":
self._lattice.lattice_vectors().detach(),
})
j2c, j3c, el_mat = self._cache.cache_multi(
["j2c", "j3c", "el_mat"], _calc_integrals)
self._j2c = j2c
self._j3c = j3c
self._el_mat = el_mat
return self