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 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