torchref.base.electron_density.main module

Central electron density building with automatic backend selection.

Dispatches to the optimal implementation based on device and available backends:

GPU backends (ISO_MAP_ENGINE_GPU):

“separable_triton” — Tier 0, separable 1-D table lookups (default) “fused_triton” — Tier 1, fused Triton kernel “original” — Tier 2, find_relevant_voxels + vectorized_add_to_map “auto” — try separable → fused → original (legacy default)

CPU backends (ISO_MAP_ENGINE_CPU):

“separable” — separable Gaussian splatting (default) “separable_compiled” — torch.compile’d separable “fused” — fused fractional-space kernel “original” — find_relevant_voxels + vectorized_add_to_map (same as GPU Tier 2)

Override at import time via environment variables:

TORCHREF_ISO_MAP_ENGINE_GPU=auto
TORCHREF_ISO_MAP_ENGINE_CPU=fused

or at runtime by setting the module-level variables directly:

import torchref.base.electron_density.main as ed
ed.ISO_MAP_ENGINE_GPU = "original"
ed.ISO_MAP_ENGINE_CPU = "fused"
torchref.base.electron_density.main.build_electron_density(real_space_grid, xyz_iso, adp_iso, occ_iso, A_iso, B_iso, inv_frac_matrix, frac_matrix, radius_angstrom, voxel_size, xyz_aniso=None, u_aniso=None, occ_aniso=None, A_aniso=None, B_aniso=None, dtype=torch.float32)[source]

Build an electron density map from atomic parameters.

Selects the fastest available backend automatically. On CUDA, tries the fused Triton kernel first (eliminates find_relevant_voxels), then falls back to two-step Triton or JIT. On CPU, uses the JIT kernel.

Parameters:

  • real_space_grid (torch.Tensor) – Coordinate grid, shape (nx, ny, nz, 3).

  • xyz_iso (torch.Tensor) – Isotropic atom positions, shape (n_iso, 3).

  • adp_iso (torch.Tensor) – Isotropic B-factors, shape (n_iso,).

  • occ_iso (torch.Tensor) – Isotropic occupancies, shape (n_iso,).

  • A_iso (torch.Tensor) – ITC92 coefficients, shape (n_iso, 5).

  • B_iso (torch.Tensor) – ITC92 coefficients, shape (n_iso, 5).

  • inv_frac_matrix (torch.Tensor) – Cartesian-to-fractional matrix, shape (3, 3).

  • frac_matrix (torch.Tensor) – Fractional-to-Cartesian matrix, shape (3, 3).

  • radius_angstrom (float) – Radius around each atom in Angstroms.

  • voxel_size (torch.Tensor) – Voxel dimensions, shape (3,).

  • xyz_aniso (torch.Tensor, optional) – Anisotropic atom positions, shape (n_aniso, 3).

  • u_aniso (torch.Tensor, optional) – Anisotropic U parameters, shape (n_aniso, 6).

  • occ_aniso (torch.Tensor, optional) – Anisotropic occupancies, shape (n_aniso,).

  • A_aniso (torch.Tensor, optional) – ITC92 coefficients for anisotropic atoms, shape (n_aniso, 5).

  • B_aniso (torch.Tensor, optional) – ITC92 coefficients for anisotropic atoms, shape (n_aniso, 5).

  • dtype (torch.dtype, optional) – Float dtype for the density map. Default torch.float32.

Returns:

Electron density map, shape (nx, ny, nz).

Return type:

torch.Tensor