def __init__(self, radgrid: RadialGrid, prec: int) -> None:
self._dtype = radgrid.dtype
self._device = radgrid.device
assert (prec % 2 == 1) and (3 <= prec <= 131),\
"Precision must be an odd number between 3 and 131"
# load the Lebedev grid points
lebedev_dsets = torch.tensor(LebedevLoader.load(prec),
dtype=self._dtype,
device=self._device)
wphitheta = lebedev_dsets[:, -1] # (nphitheta)
phi = lebedev_dsets[:, 0]
theta = lebedev_dsets[:, 1]
# get the radial grid
assert radgrid.coord_type == "radial"
r = radgrid.get_rgrid().unsqueeze(-1) # (nr, 1)
# get the cartesian coordinate
rsintheta = r * torch.sin(theta)
x = (rsintheta * torch.cos(phi)).view(-1, 1) # (nr * nphitheta, 1)
y = (rsintheta * torch.sin(phi)).view(-1, 1)
z = (r * torch.cos(theta)).view(-1, 1)
xyz = torch.cat((x, y, z), dim=-1) # (nr * nphitheta, ndim)
self._xyz = xyz
# calculate the dvolume (integration weights)
dvol_rad = radgrid.get_dvolume().unsqueeze(-1) # (nr, 1)
self._dvolume = (dvol_rad * wphitheta).view(-1) # (nr * nphitheta)
def test_radial_grid_dvol(grid_integrator, grid_transform):
ngrid = 40
dtype = torch.float64
radgrid = RadialGrid(ngrid, grid_integrator, grid_transform, dtype=dtype)
dvol = radgrid.get_dvolume() # (ngrid,)
rgrid = radgrid.get_rgrid() # (ngrid, ndim)
r = rgrid[:, 0]
# test gaussian integration
fcn = torch.exp(-r * r * 0.5) # (ngrid,)
int1 = (fcn * dvol).sum()
val1 = 2 * np.sqrt(2 * np.pi) * np.pi
assert torch.allclose(int1, int1 * 0 + val1)
def get_atomic_grid(grid_inp: Union[int, str],
atomz: int,
dtype: torch.dtype = _dtype,
device: torch.device = _device) -> BaseGrid:
"""
Returns an individual atomic grid centered at (0, 0, 0) given the grid input description.
"""
if isinstance(grid_inp, int):
# grid_inp as an int is deprecated (TODO: put a warning here)
# 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=dtype,
device=device)
return LebedevGrid(radgrid, prec=prec)
elif isinstance(grid_inp, str):
grid_str = grid_inp.lower().replace("-", "")
grid_cls = _get_grid_cls(grid_str)
return grid_cls(atomz, dtype=dtype, device=device)
else:
raise TypeError("Unknown type of grid_inp: %s" % type(grid_inp))
def __init__(self,
atomz: int,
dtype: torch.dtype = _dtype,
device: torch.device = _device):
# prepare the whole radial grid
ratom = get_atomic_radius(atomz)
grid_transform = DE2Transformation(alpha=self.de2_alphas.get(
atomz, 1.0),
rmin=1e-7,
rmax=15 * ratom)
radgrid = RadialGrid(self.nr,
grid_integrator="uniform",
grid_transform=grid_transform,
dtype=dtype,
device=device)
is_truncated = self._is_truncated(atomz)
if is_truncated:
# truncated
slices = self._get_truncate_slices(atomz)
precs = self._get_truncate_precs(atomz)
assert len(slices) == len(
precs), "Please report this bug to the github page"
# list of radial grid slices
radgrids: List[BaseGrid] = [radgrid[sl] for sl in slices]
else:
# not truncated
radgrids = [radgrid]
precs = [self.prec]
super().__init__(radgrids, precs)
def test_pbc_becke_grid_dvol(rgrid_integrator, rgrid_transform):
dtype = torch.float64
nr = 40
prec = 7
radgrid = RadialGrid(nr, rgrid_integrator, rgrid_transform, dtype=dtype)
sphgrid = LebedevGrid(radgrid, prec=prec)
atompos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype)
natoms = atompos.shape[0]
lattice = Lattice(torch.eye(3, dtype=dtype) * 3)
grid = PBCBeckeGrid([sphgrid] * natoms, atompos, lattice=lattice)
dvol = grid.get_dvolume() # (ngrid,)
rgrid = grid.get_rgrid() # (ngrid, ndim)
ls = lattice.get_lattice_ls(rcut=5) # (nls, ndim)
atomposs = (atompos.unsqueeze(1) + ls).reshape(-1, 1,
3) # (natoms * nls, ndim)
# test gaussian integration
fcn = torch.exp(-((rgrid - atomposs)**2).sum(dim=-1)).sum(dim=0) # (ngrid)
# fcn = rgrid[:, 0] * 0 + 1
int1 = (fcn * dvol).sum()
val1 = int1 * 0 + 2 * np.pi**1.5 # analytical function
# TODO: rtol is relatively large, maybe inspect the Becke integration grid?
assert torch.allclose(int1, int1 * 0 + val1, rtol=1e-2)
def test_lebedev_grid_dvol(rgrid_integrator, rgrid_transform):
dtype = torch.float64
nr = 40
prec = 7
radgrid = RadialGrid(nr, rgrid_integrator, rgrid_transform, dtype=dtype)
sphgrid = LebedevGrid(radgrid, prec=prec)
dvol = sphgrid.get_dvolume() # (ngrid,)
rgrid = sphgrid.get_rgrid() # (ngrid, ndim)
x = rgrid[:, 0]
y = rgrid[:, 1]
z = rgrid[:, 2]
# test gaussian integration
fcn = torch.exp(-(x * x + y * y + z * z) * 0.5)
int1 = (fcn * dvol).sum()
val1 = 2 * np.sqrt(2 * np.pi) * np.pi
assert torch.allclose(int1, int1 * 0 + val1)
def test_becke_grid_dvol(rgrid_integrator, rgrid_transform):
dtype = torch.float64
nr = 40
prec = 7
radgrid = RadialGrid(nr, rgrid_integrator, rgrid_transform, dtype=dtype)
sphgrid = LebedevGrid(radgrid, prec=prec)
atompos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype)
natoms = atompos.shape[0]
grid = BeckeGrid([sphgrid, sphgrid], atompos)
dvol = grid.get_dvolume() # (ngrid,)
rgrid = grid.get_rgrid() # (ngrid, ndim)
atompos = atompos.unsqueeze(1) # (natoms, 1, ndim)
# test gaussian integration
fcn = torch.exp(-((rgrid - atompos)**2).sum(dim=-1) * 0.5).sum(
dim=0) # (ngrid)
int1 = (fcn * dvol).sum()
val1 = 2 * (2 * np.sqrt(2 * np.pi) * np.pi)
# TODO: rtol is relatively large, maybe inspect the Becke integration grid?
assert torch.allclose(int1, int1 * 0 + val1, rtol=3e-3)
def rad_slices(self, atz: int, radgrid: RadialGrid) -> List[slice]:
ratom = self._radii_list[atz]
ralphas = self._alphas * ratom
rgrid = radgrid.get_rgrid().reshape(-1, 1) # (nr, 1)
if atz <= 2: # H & He
ralphas_i = ralphas[0]
elif atz <= 10:
ralphas_i = ralphas[1]
else:
ralphas_i = ralphas[2]
# place has value from 0 to 4 (inclusive)
place = torch.sum(rgrid > ralphas_i, dim=-1) # (nr,)
# convert it to slice
pl, counts = torch.unique_consecutive(place, return_counts=True)
idx = 0
res: List[slice] = []
precs = self._get_precs(atz)
for i in range(len(precs)):
c = int(counts[i])
res.append(slice(idx, idx + c, None))
idx += c
return res
def get_grid(atomzs: Union[List[int], torch.Tensor],
atompos: torch.Tensor,
*,
lattice: Optional[Lattice] = None,
nr: Union[int, Callable[[int], int]] = 99,
nang: Union[int, Callable[[int], int]] = 590,
radgrid_generator: str = "uniform",
radgrid_transform: str = "sg2-dasgupta",
atom_radii: str = "expected",
multiatoms_scheme: str = "becke",
truncate: Optional[str] = "dasgupta",
dtype: torch.dtype = _dtype,
device: torch.device = _device) -> BaseGrid:
# atompos: (natoms, ndim)
assert atompos.ndim == 2
assert atompos.shape[-2] == len(atomzs)
# convert the atomzs to a list of integers
if isinstance(atomzs, torch.Tensor):
assert atomzs.ndim == 1
atomzs_list = [a.item() for a in atomzs]
else:
atomzs_list = list(atomzs)
# get the atom radii list
atom_radii_options: Mapping[str, Union[List[float]]] = {
"expected": atom_expected_radii,
"bragg": atom_bragg_radii,
}
atom_radii_list = get_option("atom radii", atom_radii, atom_radii_options)
atomradii = torch.tensor([atom_radii_list[atomz] for atomz in atomzs_list],
dtype=dtype,
device=device)
# construct the radial grid transformation as a function of atom z
radgrid_tf_options = {
"sg2-dasgupta":
lambda atz: DE2Transformation(alpha=__sg2_dasgupta_alphas[atz],
rmin=1e-7,
rmax=15 * atom_radii_list[atz]),
"sg3-dasgupta":
lambda atz: DE2Transformation(alpha=__sg3_dasgupta_alphas[atz],
rmin=1e-7,
rmax=15 * atom_radii_list[atz]),
"logm3":
lambda atz: LogM3Transformation(ra=atom_radii_list[atz]),
"treutlerm4":
lambda atz: TreutlerM4Transformation(xi=__treutler_xi[atz], alpha=0.6),
}
radgrid_tf = get_option("radial grid transformation", radgrid_transform,
radgrid_tf_options)
# get the precisions
if isinstance(nang, int):
prec: Union[int,
Callable[[int],
int]] = get_option("number of angular points",
nang, __nang2prec)
else:
def _prec_fcn(atz: int) -> int:
assert callable(nang)
return get_option("number of angular points", nang(atz),
__nang2prec)
prec = _prec_fcn
# wrap up a function to get the nr
def _get_nr(nr: Union[int, Callable[[int], int]], atz: int):
if isinstance(nr, int):
return nr
else:
return nr(atz)
# get the truncation rule as a function to avoid unnecessary evaluation
trunc_options = {
"dasgupta":
lambda: DasguptaTrunc(nr),
"nwchem":
lambda: NWChemTrunc(atom_radii_list,
prec,
list(__nang2prec.values()),
dtype=dtype,
device=device),
"no":
lambda: NoTrunc(),
}
truncate_str = truncate if truncate is not None else "no"
trunc = get_option("truncation rule", truncate_str, trunc_options)()
sphgrids: List[BaseGrid] = []
sphgrids_dict: Dict[int, BaseGrid] = {}
for atz in atomzs_list:
if atz in sphgrids_dict:
sphgrids.append(sphgrids_dict[atz])
continue
nr_value = _get_nr(nr, atz)
radgrid = RadialGrid(nr_value,
grid_integrator=radgrid_generator,
grid_transform=radgrid_tf(atz),
dtype=dtype,
device=device)
if trunc.to_truncate(atz):
rad_slices = trunc.rad_slices(atz, radgrid)
radgrids: List[BaseGrid] = [radgrid[sl] for sl in rad_slices]
precs = trunc.precs(atz, radgrid)
sphgrid = TruncatedLebedevGrid(radgrids, precs)
else:
sphgrid = LebedevGrid(radgrid, prec=_get_nr(prec, atz))
sphgrids_dict[atz] = sphgrid
sphgrids.append(sphgrid)
# get the multi atoms grid
# the values are a function to avoid constructing it unnecessarily
if lattice is None:
multiatoms_options: Mapping[str, Callable[[], BaseGrid]] = {
"becke":
lambda: BeckeGrid(sphgrids, atompos, atomradii=atomradii),
"treutler":
lambda: BeckeGrid(sphgrids,
atompos,
atomradii=atomradii,
ratom_adjust="treutler"),
}
else:
assert isinstance(lattice, Lattice)
multiatoms_options = {
"becke":
lambda: PBCBeckeGrid(sphgrids, atompos, lattice=lattice
), # type: ignore
"treutler":
lambda: PBCBeckeGrid(
sphgrids,
atompos,
lattice=lattice, # type: ignore
ratom_adjust="treutler"),
}
grid = get_option("multiatoms scheme", multiatoms_scheme,
multiatoms_options)()
return grid