torchref.base.electron_density.voxel_utils module

Voxel utility functions for electron density calculations.

Functions for finding relevant voxels around atoms for map building.

torchref.base.electron_density.voxel_utils.find_relevant_voxels(real_space_grid, xyz, radius_angstrom=4, inv_frac_matrix=None)[source]

Identify surrounding voxels of atoms in a real space grid.

This is a vectorized function that finds all voxels within a spherical radius around each atom position.

Parameters:

  • real_space_grid (torch.Tensor) – Real space grid containing xyz coordinates at each grid point, of shape (nx, ny, nz, 3).

  • xyz (torch.Tensor) – Atom coordinates in real space (Cartesian coordinates), of shape (N, 3) or (3,).

  • radius_angstrom (float, optional) – Radius around each atom in Angstroms. Default is 4.

  • inv_frac_matrix (torch.Tensor, optional) – Matrix to convert Cartesian to fractional coordinates of shape (3, 3). Required for proper handling of non-orthogonal cells.

Returns:

  • surrounding_coords (torch.Tensor) – Coordinates of surrounding voxels for each atom of shape (N, R, 3), where R is the number of voxels within the radius.

  • voxel_indices_wrapped (torch.Tensor) – Wrapped voxel indices of shape (N, R, 3).

Notes

Atom coordinates are NOT wrapped here - periodic boundary conditions are handled in smallest_diff() which finds the minimum image distance. We only wrap voxel indices to ensure they’re valid array indices.

torchref.base.electron_density.voxel_utils.excise_angstrom_radius_around_coord(real_space_grid, start_indices, radius_angstrom=4.0)[source]

Identify voxel indices within an Angstrom radius around specified grid positions.

Parameters:

  • real_space_grid (torch.Tensor) – Real space grid of shape (nx, ny, nz, 3) containing xyz coordinates.

  • start_indices (torch.Tensor) – Starting grid indices of shape (N, 3) or (3,).

  • radius_angstrom (float, optional) – Radius in Angstroms. Default is 4.0.

Returns:

Wrapped voxel indices of shape (N, R, 3), where R is the number of voxels within the radius.

Return type:

torch.Tensor

Notes

Periodic boundary conditions are handled by wrapping the indices to ensure they’re valid array indices.