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 __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 __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 __init__(self,
dfinfo: DensityFitInfo,
wrapper: intor.LibcintWrapper,
kpts: torch.Tensor,
wkpts: torch.Tensor,
eta: float,
lattsum_opt: intor.PBCIntOption,
*,
cache: Optional[Cache] = None):
self._dfinfo = dfinfo
self._wrapper = wrapper
self._eta = eta
self._kpts = kpts
self._wkpts = wkpts # weights of each k-points
self._lattsum_opt = lattsum_opt
self.dtype = wrapper.dtype
self.device = wrapper.device
assert wrapper.lattice is not None
self._lattice = wrapper.lattice
self._is_built = False
# set up cache
self._cache = cache if cache is not None else Cache.get_dummy()
self._cache.add_cacheable_params(["j2c", "j3c", "el_mat"])
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
class Mol(BaseSystem):
"""
Describe the system of an isolated molecule.
Arguments
---------
* moldesc: str or 2-elements tuple
Description of the molecule system.
If string, it can be described like ``"H 1 0 0; H -1 0 0"``.
If tuple, the first element of the tuple is the Z number of the atoms while
the second element is the position of the atoms: ``(atomzs, atomposs)``.
* basis: str, CGTOBasis, list of str, or CGTOBasis
The string describing the gto basis. If it is a list, then it must have
the same length as the number of atoms.
* grid: int
Describe the grid.
If it is an integer, then it uses the default grid with specified level
of accuracy.
* spin: int, float, torch.Tensor, or None
The difference between spin-up and spin-down electrons.
It must be an integer or ``None``.
If ``None``, then it is ``num_electrons % 2``.
For floating point atomzs and/or charge, the ``spin`` must be specified.
* charge: int, float, or torch.Tensor
The charge of the molecule.
* orb_weights: SpinParam[torch.Tensor] or None
Specifiying the orbital occupancy (or weights) directly. If specified,
``spin`` and ``charge`` arguments are ignored.
* vext: tensor or None
The tensor describing the external potential given in the grid.
The grid position can be obtained by ``Mol().get_grid().get_rgrid()``.
* efield: tensor, tuple of tensor, or None
Uniform electric field of the system. If a tensor, then it is assumed
to be a constant electric field with the energy is
calculated based on potential at ``(0, 0, 0)`` is ``0``.
If a tuple of tensor, then the first element will have a shape of ``(ndim,)``
representing the constant electric field, second element is the gradient
of electric field with the last dimension is the direction of the electric
field, third element is the gradgrad of electric field, etc.
If ``None``, then the electric field is assumed to be ``0``.
* dtype: torch.dtype
The data type of tensors in this class.
* device: torch.device
The device on which the tensors in this class are stored.
* orthogonalize_basis: bool
(computational option)
If True, orthogonalize the basis in the hamiltonian calculation.
If False, then use the raw basis, this might not work with over-complete
basis.
* ao_parameterizer: str
(computational option)
Specifying the atomic orbital parameterizer.
"""
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 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._preproc_efield,
vext=self._vext,
cache=self._cache.add_prefix("hamilton"),
orthozer=self._orthogonalize_basis,
aoparamzer=self._aoparamzer)
return self
def get_hamiltonian(self) -> BaseHamilton:
"""
Returns the Hamiltonian that corresponds to the system, i.e.
:class:`~dqc.hamilton.HamiltonCGTO`
"""
return self._hamilton
def set_cache(self, fname: str, paramnames: Optional[List[str]] = None) -> BaseSystem:
"""
Setup the cache of some parameters specified by `paramnames` to be read/written
on a file.
If the file exists, then the parameters will not be recomputed, but just
loaded from the cache instead.
Cache is usually used for repeated calculations where the cached parameters
are not changed (e.g. running multiple systems with slightly different environment.)
Arguments
---------
fname: str
The file to store the cache.
paramnames: list of str or None
List of parameter names to be read/write from the cache.
"""
all_paramnames = self._cache.get_cacheable_params()
if paramnames is not None:
# check the paramnames
for pname in paramnames:
if pname not in all_paramnames:
msg = "Parameter %s is not cache-able. Cache-able parameters are %s" % \
(pname, all_paramnames)
raise ValueError(msg)
self._cache.set(fname, paramnames)
return self
def get_orbweight(self, polarized: bool = False) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
if not polarized:
return self._orb_weights
else:
return SpinParam(u=self._orb_weights_u, d=self._orb_weights_d)
def get_nuclei_energy(self) -> torch.Tensor:
# atomzs: (natoms,)
# atompos: (natoms, ndim)
# r12: (natoms, natom)
r12 = safe_cdist(self._atompos, self._atompos, add_diag_eps=True, diag_inf=True)
z12 = self._atomzs.unsqueeze(-2) * self._atomzs.unsqueeze(-1) # (natoms, natoms)
q_by_r = z12 / r12
return q_by_r.sum() * 0.5
def setup_grid(self) -> None:
grid_inp = self._grid_inp
logger.log("Constructing the integration grid")
self._grid = get_predefined_grid(self._grid_inp, self._atomzs_int, self._atompos,
dtype=self._dtype, device=self._device)
logger.log("Constructing the integration grid: done")
# # 0, 1, 2, 3, 4, 5
# nr = [20, 40, 60, 75, 100, 125][grid_inp]
# prec = [13, 17, 21, 29, 41, 59][grid_inp]
# radgrid = RadialGrid(nr, "chebyshev", "logm3",
# dtype=self._dtype, device=self._device)
# sphgrid = LebedevGrid(radgrid, prec=prec)
#
# natoms = self._atompos.shape[-2]
# sphgrids = [sphgrid for _ in range(natoms)]
# self._grid = BeckeGrid(sphgrids, self._atompos)
def get_grid(self) -> BaseGrid:
if self._grid is None:
raise RuntimeError("Please run mol.setup_grid() first before calling get_grid()")
return self._grid
def requires_grid(self) -> bool:
req_grid = self._vext is not None
return req_grid
def getparamnames(self, methodname: str, prefix: str = "") -> List[str]:
if methodname == "get_nuclei_energy":
params = [prefix + "_atompos"]
if torch.is_floating_point(self._atomzs):
params += [prefix + "_atomzs"]
return params
else:
raise KeyError("Unknown methodname: %s" % methodname)
################### properties ###################
@property
def atompos(self) -> torch.Tensor:
return self._atompos
@property
def atomzs(self) -> torch.Tensor:
return self._atomzs
@property
def atommasses(self) -> torch.Tensor:
# returns the atomic mass (only for non-isotope for now)
if torch.is_floating_point(self._atomzs):
raise RuntimeError("Atom masses are not available for floating point Z")
return torch.tensor([get_atom_mass(int(atomz)) for atomz in self._atomzs],
dtype=self._dtype, device=self._device)
@property
def spin(self) -> ZType:
return self._spin
@property
def charge(self) -> ZType:
return self._charge
@property
def numel(self) -> ZType:
return self._numel
@property
def efield(self) -> Optional[Tuple[torch.Tensor, ...]]:
return self._efield
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
class Mol(BaseSystem):
"""
Describe the system of an isolated molecule.
Arguments
---------
* moldesc: str or 2-elements tuple (atomzs, atompos)
Description of the molecule system.
If string, it can be described like "H 0 0 0; H 0.5 0.5 0.5".
If tuple, the first element of the tuple is the Z number of the atoms while
the second element is the position of the atoms.
* basis: str, CGTOBasis or list of str or CGTOBasis
The string describing the gto basis. If it is a list, then it must have
the same length as the number of atoms.
* grid: int
Describe the grid.
If it is an integer, then it uses the default grid with specified level
of accuracy.
Default: 3
* spin: int, float, torch.Tensor, or None
The difference between spin-up and spin-down electrons.
It must be an integer or None.
If None, then it is ``num_electrons % 2``.
For floating point atomzs and/or charge, the ``spin`` must be specified.
Default: None
* charge: int, float, or torch.Tensor
The charge of the molecule.
Default: 0
* orb_weights: SpinParam[torch.Tensor] or None
Specifiying the orbital occupancy (or weights) directly. If specified,
``spin`` and ``charge`` arguments are ignored.
* efield: Optional[torch.Tensor]
Uniform electric field of the system. If present, then the energy is
calculated based on potential at (0, 0, 0) = 0.
If None, then the electric field is assumed to be 0.
* dtype: torch.dtype
The data type of tensors in this class.
Default: torch.float64
* device: torch.device
The device on which the tensors in this class are stored.
Default: torch.device('cpu')
"""
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
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 get_hamiltonian(self) -> BaseHamilton:
return self._hamilton
def set_cache(self,
fname: str,
paramnames: Optional[List[str]] = None) -> BaseSystem:
"""
Setup the cache of some parameters specified by `paramnames` to be read/written
on a file.
If the file exists, then the parameters will not be recomputed, but just
loaded from the cache instead.
Arguments
---------
fname: str
The file to store the cache.
paramnames: Optional[List[str]]
List of parameter names to be read/write from the cache.
"""
all_paramnames = self._cache.get_cacheable_params()
if paramnames is not None:
# check the paramnames
for pname in paramnames:
if pname not in all_paramnames:
msg = "Parameter %s is not cache-able. Cache-able parameters are %s" % \
(pname, all_paramnames)
raise ValueError(msg)
self._cache.set(fname, paramnames)
return self
def get_orbweight(
self,
polarized: bool = False
) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
if not polarized:
return self._orb_weights
else:
return SpinParam(u=self._orb_weights_u, d=self._orb_weights_d)
def get_nuclei_energy(self) -> torch.Tensor:
# atomzs: (natoms,)
# atompos: (natoms, ndim)
# r12: (natoms, natom)
r12 = safe_cdist(self._atompos,
self._atompos,
add_diag_eps=True,
diag_inf=True)
z12 = self._atomzs.unsqueeze(-2) * self._atomzs.unsqueeze(
-1) # (natoms, natoms)
q_by_r = z12 / r12
return q_by_r.sum() * 0.5
def setup_grid(self) -> None:
grid_inp = self._grid_inp
self._grid = get_grid(self._grid_inp,
self._atomzs_int,
self._atompos,
dtype=self._dtype,
device=self._device)
# # 0, 1, 2, 3, 4, 5
# nr = [20, 40, 60, 75, 100, 125][grid_inp]
# prec = [13, 17, 21, 29, 41, 59][grid_inp]
# radgrid = RadialGrid(nr, "chebyshev", "logm3",
# dtype=self._dtype, device=self._device)
# sphgrid = LebedevGrid(radgrid, prec=prec)
#
# natoms = self._atompos.shape[-2]
# sphgrids = [sphgrid for _ in range(natoms)]
# self._grid = BeckeGrid(sphgrids, self._atompos)
def get_grid(self) -> BaseGrid:
if self._grid is None:
raise RuntimeError(
"Please run mol.setup_grid() first before calling get_grid()")
return self._grid
def getparamnames(self, methodname: str, prefix: str = "") -> List[str]:
pass
@property
def spin(self) -> ZType:
return self._spin
@property
def charge(self) -> ZType:
return self._charge
@property
def numel(self) -> ZType:
return self._numel
class Sol(BaseSystem):
"""
Describe the system of a solid (i.e. periodic boundary condition system).
Arguments
---------
* soldesc: str or 2-elements tuple
Description of the molecule system.
If string, it can be described like ``"H 1 0 0; H -1 0 0"``.
If tuple, the first element of the tuple is the Z number of the atoms while
the second element is the position of the atoms: ``(atomzs, atomposs)``.
* basis: str, CGTOBasis, list of str, or CGTOBasis
The string describing the gto basis. If it is a list, then it must have
the same length as the number of atoms.
* grid: int
Describe the grid.
If it is an integer, then it uses the default grid with specified level
of accuracy.
* spin: int, float, torch.Tensor, or None
The difference between spin-up and spin-down electrons.
It must be an integer or ``None``.
If ``None``, then it is ``num_electrons % 2``.
For floating point atomzs and/or charge, the ``spin`` must be specified.
* charge: int, float, or torch.Tensor
The charge of the molecule.
* orb_weights: SpinParam[torch.Tensor] or None
Specifiying the orbital occupancy (or weights) directly. If specified,
``spin`` and ``charge`` arguments are ignored.
* dtype: torch.dtype
The data type of tensors in this class.
* device: torch.device
The device on which the tensors in this class are stored.
"""
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 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_hamiltonian(self) -> BaseHamilton:
"""
Returns the Hamiltonian that corresponds to the system, i.e.
:class:`~dqc.hamilton.HamiltonCGTO_PBC`
"""
return self._hamilton
def set_cache(self,
fname: str,
paramnames: Optional[List[str]] = None) -> BaseSystem:
"""
Setup the cache of some parameters specified by `paramnames` to be read/written
on a file.
If the file exists, then the parameters will not be recomputed, but just
loaded from the cache instead.
Cache is usually used for repeated calculations where the cached parameters
are not changed (e.g. running multiple systems with slightly different environment.)
Arguments
---------
fname: str
The file to store the cache.
paramnames: list of str or None
List of parameter names to be read/write from the cache.
"""
self._cache.set(fname, paramnames)
return self
def get_orbweight(
self,
polarized: bool = False
) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
if not polarized:
return self._orb_weights
else:
return SpinParam(u=self._orb_weights_u, d=self._orb_weights_d)
def get_nuclei_energy(self) -> torch.Tensor:
# self._atomzs: (natoms,)
# self._atompos: (natoms, ndim)
# r12: (natoms, natoms)
r12_inf = safe_cdist(self._atompos,
self._atompos,
add_diag_eps=True,
diag_inf=True)
r12 = safe_cdist(self._atompos, self._atompos, add_diag_eps=True)
z12 = self._atomzs.unsqueeze(-2) * self._atomzs.unsqueeze(
-1) # (natoms, natoms)
precision = self._lattsum_opt.precision
eta = self._lattice.estimate_ewald_eta(precision) * 2
vol = self._lattice.volume()
rcut = scipy.special.erfcinv(
float(vol.detach()) * eta * eta / (2 * np.pi) * precision) / eta
gcut = scipy.special.erfcinv(
precision * np.sqrt(np.pi) / 2 / eta) * 2 * eta
# get the shift vector in real space and in reciprocal space
ls = self._lattice.get_lattice_ls(rcut=rcut,
exclude_zeros=True) # (nls, ndim)
# gv: (ngv, ndim), gvweights: (ngv,)
gv, gvweights = self._lattice.get_gvgrids(gcut=gcut,
exclude_zeros=True)
gv_norm2 = torch.einsum("gd,gd->g", gv, gv) # (ngv)
# get the shift in position
atpos_shift = self._atompos - ls.unsqueeze(-2) # (nls, natoms, ndim)
r12_ls = safe_cdist(atpos_shift,
self._atompos) # (nls, natoms, natoms)
# calculate the short range
short_range_comp1 = torch.erfc(
eta * r12_ls) / r12_ls # (nls, natoms, natoms)
short_range_comp2 = torch.erfc(eta * r12) / r12_inf # (natoms, natoms)
short_range1 = torch.sum(z12 * short_range_comp1) # scalar
short_range2 = torch.sum(z12 * short_range_comp2) # scalar
short_range = short_range1 + short_range2
# calculate the long range sum
coul_g = 4 * np.pi / gv_norm2 * gvweights # (ngv,)
# this part below is quicker, but raises warning from pytorch
si = torch.exp(
1j *
torch.matmul(self._atompos, gv.transpose(-2, -1))) # (natoms, ngv)
zsi = torch.einsum("a,ag->g", self._atomzs.to(si.dtype), si) # (ngv,)
zexpg2 = zsi * torch.exp(-gv_norm2 / (4 * eta * eta))
long_range = torch.einsum("a,a,a->", zsi.conj(), zexpg2,
coul_g.to(zsi.dtype)).real # (scalar)
# # alternative way to compute the long-range part
# r12_pair = self._atompos.unsqueeze(-2) - self._atompos # (natoms, natoms, ndim)
# long_range_exp = torch.exp(-gv_norm2 / (4 * eta * eta)) * coul_g # (ngv,)
# long_range_cos = torch.cos(torch.einsum("gd,abd->gab", gv, -r12_pair)) # (ngv, natoms, natoms)
# long_range = torch.sum(long_range_exp[:, None, None] * long_range_cos * z12) # scalar
# background interaction
vbar1 = -torch.sum(self._atomzs**2) * (2 * eta / np.sqrt(np.pi))
vbar2 = -torch.sum(self._atomzs)**2 * np.pi / (eta * eta * vol)
vbar = vbar1 + vbar2 # (scalar)
eii = short_range + long_range + vbar
return eii * 0.5
def setup_grid(self) -> None:
self._grid = get_predefined_grid(self._grid_inp,
self._atomzs,
self._atompos,
lattice=self._lattice,
dtype=self._dtype,
device=self._device)
def get_grid(self) -> BaseGrid:
if self._grid is None:
raise RuntimeError(
"Please run mol.setup_grid() first before calling get_grid()")
return self._grid
def requires_grid(self) -> bool:
return False
def getparamnames(self, methodname: str, prefix: str = "") -> List[str]:
pass
################### properties ###################
@property
def atompos(self) -> torch.Tensor:
return self._atompos
@property
def atomzs(self) -> torch.Tensor:
return self._atomzs
@property
def atommasses(self) -> torch.Tensor:
# returns the atomic mass (only for non-isotope for now)
return torch.tensor(
[get_atom_mass(int(atomz)) for atomz in self._atomzs],
dtype=self._dtype,
device=self._device)
@property
def spin(self) -> ZType:
return self._spin
@property
def charge(self) -> ZType:
return self._charge
@property
def numel(self) -> ZType:
return self._numel
@property
def efield(self) -> Optional[Tuple[torch.Tensor, ...]]:
# solid with external efield has not been implemented
return None