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)