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 _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:
# performing 3-centre integrals with libcint
drv = CGTO().GTOnr3c_drv
fill = CGTO().GTOnr3c_fill_s1
# TODO: create optimizer without the 3rd index like in
# https://github.com/pyscf/pyscf/blob/e833b9a4fd5fb24a061721e5807e92c44bb66d06/pyscf/gto/moleintor.py#L538
outsh = self.outshape
out = np.empty((*outsh[:-3], outsh[-1], outsh[-2], outsh[-3]), dtype=np.float64)
drv(self.op, fill,
out.ctypes.data_as(ctypes.c_void_p),
int2ctypes(self.ncomp),
(ctypes.c_int * len(self.shls_slice))(*self.shls_slice),
np2ctypes(self.wrapper0.full_shell_to_aoloc),
self.optimizer,
np2ctypes(self.atm), int2ctypes(self.atm.shape[0]),
np2ctypes(self.bas), int2ctypes(self.bas.shape[0]),
np2ctypes(self.env))
out = np.swapaxes(out, -3, -1)
return self._to_tensor(out)
def _int4c(self) -> torch.Tensor:
# performing 4-centre integrals with libcint
symm = self.int_nmgr.get_intgl_symmetry(self.wrapper_uniqueness)
outshape = symm.get_reduced_shape(self.outshape)
out = np.empty(outshape, dtype=np.float64)
drv = CGTO().GTOnr2e_fill_drv
fill = getattr(CGTO(), "GTOnr2e_fill_%s" % symm.code)
prescreen = ctypes.POINTER(ctypes.c_void_p)()
drv(self.op, fill, prescreen,
out.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(self.ncomp),
(ctypes.c_int * 8)(*self.shls_slice),
np2ctypes(self.wrapper0.full_shell_to_aoloc),
self.optimizer,
np2ctypes(self.atm), int2ctypes(self.atm.shape[0]),
np2ctypes(self.bas), int2ctypes(self.bas.shape[0]),
np2ctypes(self.env))
out = symm.reconstruct_array(out, self.outshape)
return self._to_tensor(out)
def gto_evaluator(wrapper: LibcintWrapper, shortname: str, rgrid: torch.Tensor,
to_transpose: bool):
# NOTE: this function do not propagate gradient and should only be used
# in this file only
# rgrid: (ngrid, ndim)
# returns: (*, nao, ngrid) if not to_transpose else (*, ngrid, nao)
ngrid = rgrid.shape[0]
nshells = len(wrapper)
nao = wrapper.nao()
opname = _get_evalgto_opname(shortname, wrapper.spherical)
outshape = _get_evalgto_compshape(shortname) + (nao, ngrid)
out = np.empty(outshape, dtype=np.float64)
non0tab = np.ones(((ngrid + BLKSIZE - 1) // BLKSIZE, nshells),
dtype=np.int8)
# TODO: check if we need to transpose it first?
rgrid = rgrid.contiguous()
coords = np.asarray(rgrid, dtype=np.float64, order='F')
ao_loc = np.asarray(wrapper.full_shell_to_aoloc, dtype=np.int32)
c_shls = (ctypes.c_int * 2)(*wrapper.shell_idxs)
c_ngrid = ctypes.c_int(ngrid)
# evaluate the orbital
operator = getattr(CGTO(), opname)
operator.restype = ctypes.c_double
atm, bas, env = wrapper.atm_bas_env
operator(c_ngrid, c_shls, np2ctypes(ao_loc), np2ctypes(out),
np2ctypes(coords), np2ctypes(non0tab), np2ctypes(atm),
int2ctypes(atm.shape[0]), np2ctypes(bas),
int2ctypes(bas.shape[0]), np2ctypes(env))
if to_transpose:
out = np.ascontiguousarray(np.moveaxis(out, -1, -2))
out_tensor = torch.as_tensor(out,
dtype=wrapper.dtype,
device=wrapper.device)
return out_tensor
def _int2c(self) -> torch.Tensor:
# performing 2-centre integrals with libcint
drv = CGTO().GTOint2c
outshape = self.outshape
out = np.empty((*outshape[:-2], outshape[-1], outshape[-2]), dtype=np.float64)
drv(self.op,
out.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(self.ncomp),
ctypes.c_int(0), # do not assume hermitian
(ctypes.c_int * len(self.shls_slice))(*self.shls_slice),
np2ctypes(self.wrapper0.full_shell_to_aoloc),
self.optimizer,
np2ctypes(self.atm), int2ctypes(self.atm.shape[0]),
np2ctypes(self.bas), int2ctypes(self.bas.shape[0]),
np2ctypes(self.env))
out = np.swapaxes(out, -2, -1)
# TODO: check if we need to do the lines below for 3rd order grad and higher
# if out.ndim > 2:
# out = np.moveaxis(out, -3, 0)
return self._to_tensor(out)
def __del__(self):
try:
CGTO.CINTdel_optimizer(ctypes.byref(self))
except AttributeError:
pass
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)