def _get_pbc_overlap_with_kpts_ij(self, gvgrids: torch.Tensor,
kpts_reduce: torch.Tensor,
kpts_j: torch.Tensor) -> torch.Tensor:
# compute the pbcft overlap integral for the basis by filling in
# the k-points one by one.
# this is to compute eq. (16) on Sun et al., https://doi.org/10.1063/1.4998644
# gvgrids: (ngv, ndim)
# kpts_reduce: (nkpts_ij, ndim)
# kpts_j: (nkpts_ij, ndim)
# returns: (nkpts_ij, nao, nao, ngv)
nkpts_ij = kpts_reduce.shape[0]
nao = self._wrapper.nao()
ngv = gvgrids.shape[0]
dctype = get_complex_dtype(self.dtype)
res = torch.empty((nkpts_ij, nao, nao, ngv),
dtype=dctype,
device=self.device)
for i in range(nkpts_ij):
kpt_ij = kpts_reduce[i] # (ndim,)
kpt_j = kpts_j[i:i + 1] # (1, ndim)
gvk = gvgrids + kpt_ij # (ngv, ndim)
aoao_ft_i = intor.pbcft_overlap(
self._wrapper,
gvgrid=gvk,
kpts=kpt_j,
options=self._lattsum_opt) # (1, nao, nao, ngv)
res[i] = aoao_ft_i[0]
return res
def _int2c(self) -> torch.Tensor:
# 2-centre integral
# this function works mostly in numpy
# no gradients propagated in this function (and it's OK)
# this function mostly replicate the `ft_aopair_kpts` function in pyscf
# https://github.com/pyscf/pyscf/blob/master/pyscf/pbc/df/ft_ao.py
# https://github.com/pyscf/pyscf/blob/c9aa2be600d75a97410c3203abf35046af8ca615/pyscf/pbc/df/ft_ao.py#L52
assert len(self.wrappers) == 2
# if the ls is too big, it might produce segfault
if (self.ls.shape[0] > 1e6):
warnings.warn("The number of neighbors in the integral is too many, "
"it might causes segfault")
# libpbc will do in-place shift of the basis of one of the wrappers, so
# we need to make a concatenated copy of the wrapper's atm_bas_env
atm, bas, env, ao_loc = _concat_atm_bas_env(self.wrappers[0], self.wrappers[1])
i0, i1 = self.wrappers[0].shell_idxs
j0, j1 = self.wrappers[1].shell_idxs
nshls0 = len(self.wrappers[0].parent)
shls_slice = (i0, i1, j0 + nshls0, j1 + nshls0)
# get the lattice translation vectors and the exponential factors
expkl = np.asarray(np.exp(1j * np.dot(self.kpts_inp_np, self.ls.T)), order='C')
# prepare the output
nGv = self.GvT.shape[-1]
nkpts = len(self.kpts_inp_np)
outshape = (nkpts,) + self.comp_shape + tuple(w.nao() for w in self.wrappers) + (nGv,)
out = np.empty(outshape, dtype=np.complex128)
# do the integration
cintor = getattr(CGTO, self.opname)
eval_gz = CPBC.GTO_Gv_general
fill = CPBC.PBC_ft_fill_ks1
drv = CPBC.PBC_ft_latsum_drv
p_gxyzT = c_null_ptr()
p_mesh = (ctypes.c_int * 3)(0, 0, 0)
p_b = (ctypes.c_double * 1)(0)
drv(cintor, eval_gz, fill,
np2ctypes(out), # ???
int2ctypes(nkpts),
int2ctypes(self.ncomp),
int2ctypes(len(self.ls)),
np2ctypes(self.ls),
np2ctypes(expkl),
(ctypes.c_int * len(shls_slice))(*shls_slice),
np2ctypes(ao_loc),
np2ctypes(self.GvT),
p_b, p_gxyzT, p_mesh,
int2ctypes(nGv),
np2ctypes(atm), int2ctypes(len(atm)),
np2ctypes(bas), int2ctypes(len(bas)),
np2ctypes(env))
out_tensor = torch.as_tensor(out, dtype=get_complex_dtype(self.dtype),
device=self.device)
return out_tensor
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 _int2c(self) -> torch.Tensor:
# 2-centre integral
# this function works mostly in numpy
# no gradients propagated in this function (and it's OK)
# this function mostly replicate the `intor_cross` function in pyscf
# https://github.com/pyscf/pyscf/blob/master/pyscf/pbc/gto/cell.py
# https://github.com/pyscf/pyscf/blob/f1321d5dd4fa103b5b04f10f31389c408949269d/pyscf/pbc/gto/cell.py#L345
assert len(self.wrappers) == 2
# libpbc will do in-place shift of the basis of one of the wrappers, so
# we need to make a concatenated copy of the wrapper's atm_bas_env
atm, bas, env, ao_loc = _concat_atm_bas_env(self.wrappers[0],
self.wrappers[1])
i0, i1 = self.wrappers[0].shell_idxs
j0, j1 = self.wrappers[1].shell_idxs
nshls0 = len(self.wrappers[0].parent)
shls_slice = (i0, i1, j0 + nshls0, j1 + nshls0)
# prepare the output
nkpts = len(self.kpts_inp_np)
outshape = (nkpts, ) + self.comp_shape + tuple(w.nao()
for w in self.wrappers)
out = np.empty(outshape, dtype=np.complex128)
# TODO: add symmetry here
fill = CPBC().PBCnr2c_fill_ks1
fintor = getattr(CGTO(), self.opname)
# TODO: use proper optimizers
cintopt = _get_intgl_optimizer(self.opname, atm, bas, env)
cpbcopt = c_null_ptr()
# get the lattice translation vectors and the exponential factors
expkl = np.asarray(np.exp(1j * np.dot(self.kpts_inp_np, self.ls.T)),
order='C')
# if the ls is too big, it might produce segfault
if (self.ls.shape[0] > 1e6):
warnings.warn(
"The number of neighbors in the integral is too many, "
"it might causes segfault")
# perform the integration
drv = CPBC().PBCnr2c_drv
drv(fintor, fill,
out.ctypes.data_as(ctypes.c_void_p), int2ctypes(nkpts),
int2ctypes(self.ncomp), int2ctypes(len(self.ls)),
np2ctypes(self.ls),
np2ctypes(expkl), (ctypes.c_int * len(shls_slice))(*shls_slice),
np2ctypes(ao_loc), cintopt, cpbcopt, np2ctypes(atm),
int2ctypes(atm.shape[0]), np2ctypes(bas), int2ctypes(bas.shape[0]),
np2ctypes(env), int2ctypes(env.size))
out_tensor = torch.as_tensor(out,
dtype=get_complex_dtype(self.dtype),
device=self.device)
return out_tensor
def _int3c(self) -> torch.Tensor:
# 3-centre integral
# this function works mostly in numpy
# no gradients propagated in this function (and it's OK)
# this function mostly replicate the `aux_e2` and `wrap_int3c` functions in pyscf
# https://github.com/pyscf/pyscf/blob/master/pyscf/pbc/df/incore.py
# https://github.com/pyscf/pyscf/blob/f1321d5dd4fa103b5b04f10f31389c408949269d/pyscf/pbc/df/incore.py#L46
assert len(self.wrappers) == 3
# libpbc will do in-place shift of the basis of one of the wrappers, so
# we need to make a concatenated copy of the wrapper's atm_bas_env
atm, bas, env, ao_loc = _concat_atm_bas_env(*self.wrappers)
i0, i1 = self.wrappers[0].shell_idxs
j0, j1 = self.wrappers[1].shell_idxs
k0, k1 = self.wrappers[2].shell_idxs
nshls0 = len(self.wrappers[0].parent)
nshls01 = len(self.wrappers[1].parent) + nshls0
shls_slice = (i0, i1, j0 + nshls0, j1 + nshls0, k0 + nshls01,
k1 + nshls01)
# kpts is actually kpts_ij in this function
nkpts_ij = len(self.kpts_inp_np)
outshape = (nkpts_ij, ) + self.comp_shape + tuple(
w.nao() for w in self.wrappers)
out = np.empty(outshape, dtype=np.complex128)
# get the unique k-points
kpts_i = self.kpts_inp_np[:, 0, :] # (nkpts, NDIM)
kpts_j = self.kpts_inp_np[:, 1, :]
kpts_stack = np.concatenate((kpts_i, kpts_j), axis=0)
kpt_diff_tol = self.options.kpt_diff_tol
_, kpts_idxs = np.unique(np.floor(kpts_stack / kpt_diff_tol) *
kpt_diff_tol,
axis=0,
return_index=True)
kpts = kpts_stack[kpts_idxs, :]
nkpts = len(kpts)
expkl = np.asarray(np.exp(1j * np.dot(kpts, self.ls.T)), order="C")
# get the kpts_ij_idxs
# TODO: check if it is the index inverse from unique
wherei = np.where(
np.abs(kpts_i.reshape(-1, 1, 3) -
kpts).sum(axis=2) < kpt_diff_tol)[1]
wherej = np.where(
np.abs(kpts_j.reshape(-1, 1, 3) -
kpts).sum(axis=2) < kpt_diff_tol)[1]
kpts_ij_idxs = np.asarray(wherei * nkpts + wherej, dtype=np.int32)
# prepare the optimizers
# TODO: use proper optimizers
# NOTE: using _get_intgl_optimizer in this case produce wrong results (I don't know why)
cintopt = c_null_ptr(
) # _get_intgl_optimizer(self.opname, atm, bas, env)
cpbcopt = c_null_ptr()
# do the integration
drv = CPBC().PBCnr3c_drv
fill = CPBC(
).PBCnr3c_fill_kks1 # TODO: optimize the kk-type and symmetry
fintor = getattr(CPBC(), self.opname)
drv(fintor, fill, np2ctypes(out), int2ctypes(nkpts_ij),
int2ctypes(nkpts), int2ctypes(self.ncomp),
int2ctypes(len(self.ls)), np2ctypes(self.ls), np2ctypes(expkl),
np2ctypes(kpts_ij_idxs),
(ctypes.c_int * len(shls_slice))(*shls_slice), np2ctypes(ao_loc),
cintopt, cpbcopt, np2ctypes(atm), int2ctypes(atm.shape[0]),
np2ctypes(bas), int2ctypes(bas.shape[0]), np2ctypes(env),
int2ctypes(env.size))
out_tensor = torch.as_tensor(out,
dtype=get_complex_dtype(self.dtype),
device=self.device)
return out_tensor
def gto_ft_evaluator(wrapper: LibcintWrapper,
Gvgrid: torch.Tensor) -> torch.Tensor:
# evaluate Fourier Transform of the basis which is defined as
# FT(f(r)) = integral(f(r) * exp(-ik.r) dr)
# NOTE: this function do not propagate gradient and should only be used
# in this file only
# this is mainly from PySCF
# https://github.com/pyscf/pyscf/blob/c9aa2be600d75a97410c3203abf35046af8ca615/pyscf/gto/ft_ao.py#L107
assert Gvgrid.ndim == 2
assert Gvgrid.shape[-1] == NDIM
# Gvgrid: (ngrid, ndim)
# returns: (nao, ngrid)
dtype = wrapper.dtype
device = wrapper.device
fill = CGTO.GTO_ft_fill_s1
if wrapper.spherical:
intor = CGTO.GTO_ft_ovlp_sph
else:
intor = CGTO.GTO_ft_ovlp_cart
fn = CGTO.GTO_ft_fill_drv
eval_gz = CGTO.GTO_Gv_general
p_gxyzT = c_null_ptr()
p_gs = (ctypes.c_int * 3)(0, 0, 0)
p_b = (ctypes.c_double * 1)(0)
# add another dummy basis to provide the multiplier
c = np.sqrt(4 * np.pi) # s-type normalization
ghost_basis = CGTOBasis(
angmom=0,
alphas=torch.tensor([0.], dtype=dtype, device=device),
coeffs=torch.tensor([c], dtype=dtype, device=device),
normalized=True,
)
ghost_atom_basis = AtomCGTOBasis(atomz=0,
bases=[ghost_basis],
pos=torch.tensor([0.0, 0.0, 0.0],
dtype=dtype,
device=device))
ghost_wrapper = LibcintWrapper([ghost_atom_basis],
spherical=wrapper.spherical,
lattice=wrapper.lattice)
wrapper, ghost_wrapper = LibcintWrapper.concatenate(wrapper, ghost_wrapper)
shls_slice = (*wrapper.shell_idxs, *ghost_wrapper.shell_idxs)
ao_loc = wrapper.full_shell_to_aoloc
atm, bas, env = wrapper.atm_bas_env
# prepare the Gvgrid
GvT = np.asarray(Gvgrid.detach().numpy().T, order="C")
nGv = Gvgrid.shape[0]
# prepare the output matrix
outshape = (wrapper.nao(), nGv)
out = np.zeros(outshape, dtype=np.complex128, order="C")
fn(intor, eval_gz, fill, np2ctypes(out), int2ctypes(1),
(ctypes.c_int * len(shls_slice))(*shls_slice), np2ctypes(ao_loc),
ctypes.c_double(0), np2ctypes(GvT), p_b, p_gxyzT, p_gs, int2ctypes(nGv),
np2ctypes(atm), int2ctypes(len(atm)), np2ctypes(bas),
int2ctypes(len(bas)), np2ctypes(env))
return torch.as_tensor(out, dtype=get_complex_dtype(dtype), device=device)