def pbc_h1():
# get the hamiltonian for pbc system
# setup the environment
pos = torch.tensor([0.0, 0.0, 0.0], dtype=dtype)
atomz = 2
atom = "He"
bases = loadbasis("%d:3-21G" % atomz, dtype=dtype)
atombases = [AtomCGTOBasis(atomz=atomz, bases=bases, pos=pos)]
kpts = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.2, 0.15]], dtype=dtype)
# setup the lattice
a = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) * 3
latt = Lattice(a)
# setup the auxbasis
auxbases = loadbasis("%d:def2-sv(p)-jkfit" % atomz, dtype=dtype)
atomauxbases = [AtomCGTOBasis(atomz=atomz, bases=auxbases, pos=pos)]
df = DensityFitInfo(method="gdf", auxbases=atomauxbases)
# build the hamiltonian
h = HamiltonCGTO_PBC(atombases, latt=latt, df=df, kpts=kpts,
lattsum_opt={"precision": 1e-8})
h.build()
# pyscf system build
mol = pyscf.pbc.gto.C(atom="%s 0 0 0" % atom, a=a.numpy(), basis="3-21G", unit="Bohr", spin=0)
df = pyscf.pbc.df.GDF(mol, kpts=kpts.detach().numpy())
df.auxbasis = "def2-svp-jkfit"
df.build()
return h, df
def _renormalize_auxbases(
auxbases: List[AtomCGTOBasis]) -> List[AtomCGTOBasis]:
# density basis renormalization, following pyscf here:
# https://github.com/pyscf/pyscf/blob/7be5e015b2b40181755c71d888449db936604660/pyscf/pbc/df/df.py#L95
# this renormalization makes the integral of auxbases (not auxbases * auxbases)
# to be 1
res: List[AtomCGTOBasis] = []
# libcint multiply np.sqrt(4*pi) to the basis
half_sph_norm = 0.5 / np.sqrt(np.pi)
for atb in auxbases: # atb is AtomCGTOBasis
bases: List[CGTOBasis] = []
for bas in atb.bases: # bas is CGTOBasis
assert bas.normalized
int1 = gaussian_int(bas.angmom * 2 + 2, bas.alphas)
s = torch.sum(bas.coeffs * int1)
coeffs = bas.coeffs * (half_sph_norm / s)
b2 = CGTOBasis(angmom=bas.angmom,
coeffs=coeffs,
alphas=bas.alphas,
normalized=True)
bases.append(b2)
atb2 = AtomCGTOBasis(atomz=atb.atomz, bases=bases, pos=atb.pos)
res.append(atb2)
return res
def _create_compensating_bases(self, atombases: List[AtomCGTOBasis], eta: float) -> List[AtomCGTOBasis]:
# create the list of atom bases containing the compensating basis with
# given `eta` as the exponentials
# see make_modchg_basis in
# https://github.com/pyscf/pyscf/blob/c9aa2be600d75a97410c3203abf35046af8ca615/pyscf/pbc/df/df.py#L116
# pre-calculate the norms up to angmom 6
half_sph_norm = 0.5 / np.sqrt(np.pi)
norms = [half_sph_norm / gaussian_int(2 * angmom + 2, eta) for angmom in range(7)]
norms_t = [torch.tensor([nrm], dtype=self.dtype, device=self.device) for nrm in norms]
res: List[AtomCGTOBasis] = []
alphas = torch.tensor([eta], dtype=self.dtype, device=self.device)
for atb in atombases:
# TODO: use reduced bases to optimize the integration time
# angmoms = set(bas.angmom for bas in atb.bases)
# bases = [
# CGTOBasis(angmom=angmom, alphas=alphas, coeffs=norms[angmom], normalized=True) \
# for angmom in angmoms
# ]
bases: List[CGTOBasis] = []
for bas in atb.bases:
# calculate the integral of the basis
int1 = gaussian_int(bas.angmom * 2 + 2, bas.alphas)
s = torch.sum(bas.coeffs * int1) / half_sph_norm
# set the coefficients of the compensating basis to have
# the same integral
coeffs = s * norms_t[bas.angmom]
b2 = CGTOBasis(angmom=bas.angmom, alphas=alphas,
coeffs=coeffs,
normalized=True)
bases.append(b2)
res.append(AtomCGTOBasis(atomz=0, bases=bases, pos=atb.pos))
return res
def __init__(
self,
soldesc: Union[str, Tuple[AtomZsType, AtomPosType]],
alattice: torch.Tensor,
basis: Union[str, List[CGTOBasis], List[str],
List[List[CGTOBasis]]],
*,
grid: Union[int, str] = "sg3",
spin: Optional[ZType] = None,
lattsum_opt: Optional[Union[PBCIntOption, Dict]] = None,
dtype: torch.dtype = torch.float64,
device: torch.device = torch.device('cpu'),
):
self._dtype = dtype
self._device = device
self._grid_inp = grid
self._grid: Optional[BaseGrid] = None
charge = 0 # we can't have charged solids for now
# get the AtomCGTOBasis & the hamiltonian
# atomzs: (natoms,) dtype: torch.int or dtype for floating point
# atompos: (natoms, ndim)
atomzs, atompos = parse_moldesc(soldesc, dtype, device)
allbases = _parse_basis(atomzs, basis) # list of list of CGTOBasis
atombases = [
AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas, atpos) in zip(atomzs, allbases, atompos)
]
self._atombases = atombases
self._atompos = atompos # (natoms, ndim)
self._atomzs = atomzs # (natoms,) int-type
nelecs_tot: torch.Tensor = torch.sum(atomzs)
# get the number of electrons and spin and orbital weights
nelecs, spin, frac_mode = _get_nelecs_spin(nelecs_tot, spin, charge)
assert not frac_mode, "Fractional Z mode for pbc is not supported"
_orb_weights, _orb_weights_u, _orb_weights_d = _get_orb_weights(
nelecs, spin, frac_mode, dtype, device)
# initialize cache
self._cache = Cache()
# save the system's properties
self._spin = spin
self._charge = charge
self._numel = nelecs
self._orb_weights = _orb_weights
self._orb_weights_u = _orb_weights_u
self._orb_weights_d = _orb_weights_d
self._lattice = Lattice(alattice)
self._lattsum_opt = PBCIntOption.get_default(lattsum_opt)
def _create_fake_nucl_bases(self, alpha: float,
chargemult: int) -> List[AtomCGTOBasis]:
# create a list of basis (of s-type) at every nuclei positions
res: List[AtomCGTOBasis] = []
alphas = torch.tensor([alpha], dtype=self.dtype, device=self.device)
# normalizing so the integral of the cgto is 1
# 0.5 / np.sqrt(np.pi) * 2 / scipy.special.gamma(1.5) * alphas ** 1.5
norm_coeff = 0.6366197723675814 * alphas**1.5
for atb in self._atombases:
# put the charge in the coefficients
coeffs = atb.atomz * norm_coeff
basis = CGTOBasis(angmom=0,
alphas=alphas,
coeffs=coeffs,
normalized=True)
res.append(AtomCGTOBasis(atomz=0, bases=[basis], pos=atb.pos))
return res
def densityfit(self,
method: Optional[str] = None,
auxbasis: Optional[BasisInpType] = None) -> BaseSystem:
"""
Indicate that the system's Hamiltonian uses density fit for its integral.
Arguments
---------
method: Optional[str]
Density fitting method. Available methods in this class are:
* "coulomb": Minimizing the Coulomb inner product, i.e. min <p-p_fit|r_12|p-p_fit>
Ref: Eichkorn, et al. Chem. Phys. Lett. 240 (1995) 283-290.
(default)
* "overlap": Minimizing the overlap inner product, i.e. min <p-p_fit|p-p_fit>
auxbasis: Optional[BasisInpType]
Auxiliary basis for the density fit. If not specified, then it uses
"cc-pvtz-jkfit".
"""
if method is None:
method = "coulomb"
if auxbasis is None:
# TODO: choose the auxbasis properly
auxbasis = "cc-pvtz-jkfit"
# get the auxiliary basis
assert auxbasis is not None
auxbasis_lst = _parse_basis(self._atomzs_int, auxbasis)
atomauxbases = [
AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas,
atpos) in zip(self._atomzs, auxbasis_lst, self._atompos)
]
# change the hamiltonian to have density fit
df = DensityFitInfo(method=method, auxbases=atomauxbases)
self._hamilton = HamiltonCGTO(self._atombases,
df=df,
efield=self._efield,
cache=self._cache.add_prefix("hamilton"))
return self
def densityfit(self,
method: Optional[str] = None,
auxbasis: Optional[BasisInpType] = None) -> BaseSystem:
"""
Indicate that the system's Hamiltonian uses density fit for its integral.
Arguments
---------
method: Optional[str]
Density fitting method. Available methods in this class are:
* ``"gdf"``: Density fit with gdf compensating charge to perform
the lattice sum. Ref https://doi.org/10.1063/1.4998644 (default)
auxbasis: Optional[BasisInpType]
Auxiliary basis for the density fit. If not specified, then it uses
``"cc-pvtz-jkfit"``.
"""
if method is None:
method = "gdf"
if auxbasis is None:
# TODO: choose the auxbasis properly
auxbasis = "cc-pvtz-jkfit"
# get the auxiliary basis
assert auxbasis is not None
auxbasis_lst = _parse_basis(self._atomzs, auxbasis)
atomauxbases = [
AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas,
atpos) in zip(self._atomzs, auxbasis_lst, self._atompos)
]
# change the hamiltonian to have density fit
df = DensityFitInfo(method=method, auxbases=atomauxbases)
self._hamilton = HamiltonCGTO_PBC(
self._atombases,
df=df,
latt=self._lattice,
lattsum_opt=self._lattsum_opt,
cache=self._cache.add_prefix("hamilton"))
return self
def get_uncontracted_wrapper(self) -> Tuple[LibcintWrapper, torch.Tensor]:
# returns the uncontracted LibcintWrapper as well as the mapping from
# uncontracted atomic orbital (relative index) to the relative index
# of the atomic orbital
new_atombases = []
for atombasis in self.atombases:
atomz = atombasis.atomz
pos = atombasis.pos
new_bases = []
for shell in atombasis.bases:
angmom = shell.angmom
alphas = shell.alphas
coeffs = shell.coeffs
normalized = shell.normalized
new_bases.extend([
CGTOBasis(angmom,
alpha[None],
coeff[None],
normalized=normalized)
for (alpha, coeff) in zip(alphas, coeffs)
])
new_atombases.append(
AtomCGTOBasis(atomz=atomz, bases=new_bases, pos=pos))
uncontr_wrapper = LibcintWrapper(new_atombases,
spherical=self.spherical)
# get the mapping uncontracted ao to the contracted ao
uao2ao: List[int] = []
idx_ao = 0
# iterate over shells
for i in range(len(self)):
nao = self._nao_at_shell(i)
uao2ao += list(range(idx_ao,
idx_ao + nao)) * self.ngauss_at_shell[i]
idx_ao += nao
uao2ao_res = torch.tensor(uao2ao, dtype=torch.long, device=self.device)
return uncontr_wrapper, uao2ao_res
def __init__(self,
moldesc: Union[str, Tuple[AtomZsType, AtomPosType]],
basis: BasisInpType,
*,
orthogonalize_basis: bool = True,
ao_parameterizer: str = "qr",
grid: Union[int, str] = "sg3",
spin: Optional[ZType] = None,
charge: ZType = 0,
orb_weights: Optional[SpinParam[torch.Tensor]] = None,
efield: Union[torch.Tensor, Tuple[torch.Tensor, ...], None] = None,
vext: Optional[torch.Tensor] = None,
dtype: torch.dtype = torch.float64,
device: torch.device = torch.device('cpu'),
):
self._dtype = dtype
self._device = device
self._grid_inp = grid
self._basis_inp = basis
self._grid: Optional[BaseGrid] = None
self._vext = vext
# make efield a tuple
self._efield = _normalize_efield(efield)
self._preproc_efield = _preprocess_efield(self._efield)
# initialize cache
self._cache = Cache()
# get the AtomCGTOBasis & the hamiltonian
# atomzs: (natoms,) dtype: torch.int or dtype for floating point
# atompos: (natoms, ndim)
atomzs, atompos = parse_moldesc(
moldesc, dtype=dtype, device=device)
atomzs_int = torch.round(atomzs).to(torch.int) if atomzs.is_floating_point() else atomzs
allbases = _parse_basis(atomzs_int, basis) # list of list of CGTOBasis
atombases = [AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas, atpos) in zip(atomzs, allbases, atompos)]
self._atombases = atombases
self._hamilton = HamiltonCGTO(atombases, efield=self._preproc_efield,
vext=self._vext,
cache=self._cache.add_prefix("hamilton"),
orthozer=orthogonalize_basis,
aoparamzer=ao_parameterizer)
self._orthogonalize_basis = orthogonalize_basis
self._aoparamzer = ao_parameterizer
self._atompos = atompos # (natoms, ndim)
self._atomzs = atomzs # (natoms,) int-type or dtype if floating point
self._atomzs_int = atomzs_int # (natoms,) int-type rounded from atomzs
nelecs_tot: torch.Tensor = torch.sum(atomzs)
# orb_weights is not specified, so determine it from spin and charge
if orb_weights is None:
# get the number of electrons and spin
nelecs, spin, frac_mode = _get_nelecs_spin(nelecs_tot, spin, charge)
_orb_weights, _orb_weights_u, _orb_weights_d = _get_orb_weights(
nelecs, spin, frac_mode, dtype, device)
# save the system's properties
self._spin = spin
self._charge = charge
self._numel = nelecs
self._orb_weights = _orb_weights
self._orb_weights_u = _orb_weights_u
self._orb_weights_d = _orb_weights_d
# orb_weights is specified, so calculate the spin and charge from it
else:
if not isinstance(orb_weights, SpinParam):
raise TypeError("Specifying orb_weights must be in SpinParam type")
assert orb_weights.u.ndim == 1
assert orb_weights.d.ndim == 1
assert len(orb_weights.u) == len(orb_weights.d)
# check if it is decreasing
orb_u_dec = torch.all(orb_weights.u[:-1] - orb_weights.u[1:] > -1e-4)
orb_d_dec = torch.all(orb_weights.d[:-1] - orb_weights.d[1:] > -1e-4)
if not (orb_u_dec and orb_d_dec):
# if not decreasing, the variational might give the wrong results
warnings.warn("The orbitals should be ordered in a non-increasing manner. "
"Otherwise, some calculations might be wrong.")
utot = orb_weights.u.sum()
dtot = orb_weights.d.sum()
self._numel = utot + dtot
self._spin = utot - dtot
self._charge = nelecs_tot - self._numel
self._orb_weights_u = orb_weights.u
self._orb_weights_d = orb_weights.d
self._orb_weights = orb_weights.u + orb_weights.d
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)
def __init__(
self,
moldesc: Union[str, Tuple[AtomZsType, AtomPosType]],
basis: BasisInpType,
*,
grid: Union[int, str] = "sg3",
spin: Optional[ZType] = None,
charge: ZType = 0,
orb_weights: Optional[SpinParam[torch.Tensor]] = None,
efield: Optional[torch.Tensor] = None,
dtype: torch.dtype = torch.float64,
device: torch.device = torch.device('cpu'),
):
self._dtype = dtype
self._device = device
self._grid_inp = grid
self._basis_inp = basis
self._grid: Optional[BaseGrid] = None
self._efield = efield
# initialize cache
self._cache = Cache()
# get the AtomCGTOBasis & the hamiltonian
# atomzs: (natoms,) dtype: torch.int or dtype for floating point
# atompos: (natoms, ndim)
atomzs, atompos = _parse_moldesc(moldesc, dtype, device)
atomzs_int = torch.round(atomzs).to(
torch.int) if atomzs.is_floating_point() else atomzs
allbases = _parse_basis(atomzs_int, basis) # list of list of CGTOBasis
atombases = [
AtomCGTOBasis(atomz=atz, bases=bas, pos=atpos)
for (atz, bas, atpos) in zip(atomzs, allbases, atompos)
]
self._atombases = atombases
self._hamilton = HamiltonCGTO(atombases,
efield=efield,
cache=self._cache.add_prefix("hamilton"))
self._atompos = atompos # (natoms, ndim)
self._atomzs = atomzs # (natoms,) int-type or dtype if floating point
self._atomzs_int = atomzs_int # (natoms,) int-type rounded from atomzs
nelecs_tot: torch.Tensor = torch.sum(atomzs)
# orb_weights is not specified, so determine it from spin and charge
if orb_weights is None:
# get the number of electrons and spin
nelecs, spin, frac_mode = _get_nelecs_spin(nelecs_tot, spin,
charge)
_orb_weights, _orb_weights_u, _orb_weights_d = _get_orb_weights(
nelecs, spin, frac_mode, dtype, device)
# save the system's properties
self._spin = spin
self._charge = charge
self._numel = nelecs
self._orb_weights = _orb_weights
self._orb_weights_u = _orb_weights_u
self._orb_weights_d = _orb_weights_d
# orb_weights is specified, so calculate the spin and charge from it
else:
assert isinstance(orb_weights, SpinParam)
assert orb_weights.u.ndim == 1
assert orb_weights.d.ndim == 1
assert len(orb_weights.u) == len(orb_weights.d)
utot = orb_weights.u.sum()
dtot = orb_weights.d.sum()
self._numel = utot + dtot
self._spin = utot - dtot
self._charge = nelecs_tot - self._numel
self._orb_weights_u = orb_weights.u
self._orb_weights_d = orb_weights.d
self._orb_weights = orb_weights.u + orb_weights.d