Take integer and index tensors from the configured int dtype

AGENTS.md

@@ -51,6 +51,12 @@ Practically:
 - **Never hardcode a dtype.** Take it from the config: `torchref.config.get_float_dtype()`,
   `get_int_dtype()`, `get_complex_dtype()`, or from an input tensor. Roughly 200 call sites
   already do this; follow them.
+- Integer and index tensors take `get_int_dtype()` too (int32 by default). Plain indexing,
+  `index_select` and `index_add_` accept it. A literal int dtype survives only where torch or
+  the arithmetic forces it, with a `# dtype-ok:` marker naming the constraint: `scatter`/`gather`
+  indices (int64 on torch < 2.8), `index_copy_`/`index_fill_`/`one_hot` (int64 always), packed
+  keys such as `i * n + j` that overflow int32, compiled kernels that `TORCH_CHECK` a dtype
+  (the Legendre shell kernel), and external-library contracts (TorchMD-Net).
 - `torch.float64` *is* a supported configuration (`TORCHREF_DTYPE_FLOAT=float64`) used as an
   eager numerical reference and in gradient checks. Code must **work** in float64, must not
   **require** it, and must not silently downcast (see `tests/integration/test_dtype_config_float64.py`).

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 
 Unreleased
 ----------
+- Integer and index tensors now take the configured int dtype (``get_int_dtype()``, ``TORCHREF_DTYPE_INT``, int32 by default) throughout the package. A hardcoded ``int64`` remains only where a torch op (``scatter``/``gather`` on torch < 2.8, ``index_copy_``), a compiled kernel, an int32 overflow, or an external library requires it, and each such site says which.
 - ``torchref.difference-map``, ``torchref.difference-refine`` and ``torchref.validate-ded`` gain ``--ded-weight {sigma_d,inverse_variance,none}`` and ``--sigma-d-gamma``. The difference MTZ now carries the unweighted ``DF``/``SIGDF`` on ``PHDELWT`` with one mean-one weight column per scheme, ``W_SD`` and ``W_IVW`` (MTZ type W), and the observed-to-model scale ``KSCALE``; ``DELFWT`` is no longer written, build the map with ``torchref.mtz2map -csf DF -cw W_IVW -cphi PHDELWT``. Registered in ``torchref.maps.ded_weights``
 - Added the ``sigma_D`` estimator (``torchref.refinement.model_error_estimation.sigma_d``): the expected true difference power per resolution shell, ``mean(dF_obs^2) - mean(sigma^2)`` with a fitted ``F_dark^gamma`` amplitude law and DerSimonian-Laird shrinkage of the signed shell power toward a decaying exponential in ``d*^2`` (fitted on all shells, so a dataset without a difference yields no power instead of the positive half of its noise), giving the Wiener weight ``S/(S + sigma^2)`` and, with a difference model, ``alpha``/``beta_model``. Inverse-variance weights suppress the strong reflections whose difference power is 10-70x that of weak ones; on independent half-datasets a Wiener weight with the true power raised map agreement 1.2-1.8x in effective patterns. The single-dataset estimate inherits the calibration of the reported sigmas, and on the campaign TD1 data (sigmas ~1.5x too large at high resolution) it emptied 60-90 % of the shells, so inverse variance stays the default; ``sigma_d`` reports its clamped-shell count and falls back to inverse variance with a warning when every shell is empty
 - Added the ``difference_sd`` collection target (``CollectionDifferenceSigmaDTarget``): the difference Gaussian centred on ``alpha * dF_calc`` with variance ``beta_model + sigma_diff^2`` from ``sigma_D`` fitted on the free set. Selected with ``torchref.difference-refine --difference-target difference_sd``; ``difference`` stays the default

tests/unit/model/test_disorder_field.py

@@ -14,6 +14,7 @@
 import pytest
 import torch
 
+from torchref.config import get_int_dtype
 from torchref.model.disorder_field import (
     DisorderFieldTensor,
     build_neighbor_list,
@@ -52,7 +53,7 @@ def test_anchor_selection_is_deterministic(coords):
     b = farthest_point_anchors(coords, 10)
     assert torch.equal(a, b)
     assert a.shape[0] == 10
-    assert a.dtype == torch.int64
+    assert a.dtype == get_int_dtype()
     # Anchors are atom indices, and distinct.
     assert int(a.max()) < coords.shape[0]
     assert torch.unique(a).shape[0] == a.shape[0]

torchref/base/electron_density/kernels/cpu/jit_reference.py

@@ -135,9 +135,8 @@ def forward(
         # Scatter add to density map
         ny: int = density_map.shape[1]
         nz: int = density_map.shape[2]
-        strides = torch.tensor(
-            [ny * nz, nz, 1], device=voxel_indices.device, dtype=torch.long  # dtype-ok: CPU-kernel strides for flat voxel index arithmetic; indexing requires long
-        )
+        # dtype-ok: int64 strides make the flat voxel index int64; scatter_add_ requires int64 on torch < 2.8
+        strides = voxel_indices.new_tensor([ny * nz, nz, 1], dtype=torch.long)
         index_flat = torch.sum(voxel_indices.to(torch.long) * strides, dim=-1).view(-1)  # dtype-ok: voxel indices flattened for scatter; indexing requires long
 
         density_map.view(-1).scatter_add_(0, index_flat, density.reshape(-1))

torchref/base/electron_density/kernels/cpu/variable_radius.py

@@ -30,6 +30,7 @@
 import torch
 
 from torchref.base.electron_density.radius_policy import _u6_to_u3
+from torchref.config import get_int_dtype
 
 _PI = math.pi
 _PI_SQ = _PI * _PI
@@ -51,8 +52,11 @@ def _bucket_by_radius(radius: torch.Tensor, center_1d: torch.Tensor):
         order_parts.append(idx)
         spans.append((float(r), cursor, cursor + idx.numel()))
         cursor += idx.numel()
-    order = (torch.cat(order_parts) if order_parts
-             else torch.zeros(0, dtype=torch.long, device=radius.device))  # dtype-ok: empty voxel-index fallback; must stay long for indexing
+    order = (
+        torch.cat(order_parts)
+        if order_parts
+        else torch.zeros(0, dtype=get_int_dtype(), device=radius.device)
+    )
     return order, spans
 
 
@@ -88,7 +92,7 @@ def _canonical_setup(xyz, inv_frac, frac, grid_dims, radius_per_atom, dtype):
     nx, ny, nz = grid_dims
     grid_f = torch.tensor(grid_dims, device=device, dtype=dtype)
     xyz_frac = (xyz @ inv_frac.T) % 1.0
-    center_idx = torch.round(xyz_frac * grid_f).to(torch.long)  # dtype-ok: rounded voxel center indices; torch indexing requires long
+    center_idx = torch.round(xyz_frac * grid_f).to(get_int_dtype())
     # w0: atom position relative to its anchor node, in Cartesian. This is what
     # centres the sphere on the atom rather than on the node.
     w0 = (xyz_frac - center_idx.to(dtype) / grid_f) @ frac.T
@@ -111,8 +115,9 @@ def add_isotropic_plain_var(density_map, xyz, adp, occ, A, B,
     device, dtype = xyz.device, density_map.dtype
     nx, ny, nz = (int(s) for s in density_map.shape)
     grid_dims = (nx, ny, nz)
-    strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long)  # dtype-ok: strides for flat voxel-index arithmetic; indexing requires long
-    grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long)  # dtype-ok: grid_shape for flat voxel-index arithmetic; indexing requires long
+    # dtype-ok: int64 strides make the flat voxel index int64; scatter_add requires int64 on torch < 2.8
+    strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long)
+    grid_shape = torch.tensor(grid_dims, device=device, dtype=get_int_dtype())
 
     order, spans, center_idx, w0 = _canonical_setup(
         xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype)
@@ -151,8 +156,9 @@ def add_anisotropic_plain_var(density_map, xyz, u, occ, A, B,
     device, dtype = xyz.device, density_map.dtype
     nx, ny, nz = (int(s) for s in density_map.shape)
     grid_dims = (nx, ny, nz)
-    strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long)  # dtype-ok: strides for flat voxel-index arithmetic; indexing requires long
-    grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long)  # dtype-ok: grid_shape for flat voxel-index arithmetic; indexing requires long
+    # dtype-ok: int64 strides make the flat voxel index int64; scatter_add requires int64 on torch < 2.8
+    strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long)
+    grid_shape = torch.tensor(grid_dims, device=device, dtype=get_int_dtype())
 
     order, spans, center_idx, w0 = _canonical_setup(
         xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype)

torchref/base/electron_density/map_building.py

@@ -28,7 +28,8 @@ def scatter_add_nd(source, index, map):
     """Vectorized n-dimensional scatter-add: ``source`` ``(N,)`` into ``map``
     ``(d1..dn)`` at ``index`` ``(N, ndim)``, returning the modified map.
     """
-    map_shape = torch.tensor(map.shape, device=index.device, dtype=torch.int64)  # dtype-ok: map_shape for stride/flat-index arithmetic feeding scatter_add; requires int64
+    # dtype-ok: int64 shape/strides make the flat index int64; scatter_add_ requires int64 on torch < 2.8
+    map_shape = torch.tensor(map.shape, device=index.device, dtype=torch.int64)
 
     # Convert n-dimensional indices to flat indices
     # For shape (d1, d2, d3, ..., dn), flat_index = i0 * (d1*d2*...*dn) + i1 * (d2*d3*...*dn) + ... + in

torchref/base/electron_density/solvent_mask.py

@@ -8,7 +8,7 @@
 import numpy as np
 import torch
 
-from torchref.config import dtypes
+from torchref.config import dtypes, get_int_dtype
 from torchref.base.coordinates.periodic_boundary import smallest_diff
 from .map_building import scatter_add_nd
 
@@ -113,7 +113,7 @@ def add_to_phenix_mask(
     )  # (N_atoms, N_voxels)
 
     # Flatten for scatter operations
-    voxel_indices_flat = voxel_indices.reshape(-1, 3).to(torch.long)  # dtype-ok: voxel indices for grid indexing; requires long
+    voxel_indices_flat = voxel_indices.reshape(-1, 3).to(get_int_dtype())
 
     # Create protein core mask using scatter_add
     int_dtype = dtypes.int

torchref/base/electron_density/voxel_utils.py

@@ -6,7 +6,7 @@
 
 import torch
 
-from torchref.config import dtypes
+from torchref.config import dtypes, get_int_dtype
 
 
 def find_relevant_voxels(real_space_grid, xyz, radius_angstrom=4, inv_frac_matrix=None):
@@ -49,13 +49,13 @@ def find_relevant_voxels(real_space_grid, xyz, radius_angstrom=4, inv_frac_matri
         # This ensures atoms outside the unit cell are correctly wrapped
         xyz_frac = torch.matmul(inv_frac_matrix, xyz.T).T  # (N, 3)
         xyz_frac = xyz_frac % 1.0  # Wrap to [0, 1]
-        center_idx = torch.round(xyz_frac * grid_shape.unsqueeze(0)).to(torch.int64)  # dtype-ok: rounded voxel center indices; torch indexing requires int64
+        center_idx = torch.round(xyz_frac * grid_shape.unsqueeze(0)).to(get_int_dtype())
     else:
         # Fallback for orthogonal cells (less accurate for non-orthogonal)
         voxelsize = real_space_grid[3, 3, 3] - real_space_grid[2, 2, 2]
         center_idx = torch.round(
             (xyz - grid_origin.unsqueeze(0)) / voxelsize.unsqueeze(0)
-        ).to(torch.int64)  # dtype-ok: voxel index cast; torch indexing requires int64
+        ).to(get_int_dtype())
 
     voxel_indices_wrapped = excise_angstrom_radius_around_coord(
         real_space_grid, center_idx, radius_angstrom

torchref/base/french_wilson.py

@@ -1194,7 +1194,7 @@ def estimate_mean_intensity_by_resolution(
 
     # Use scatter_add to compute sum of intensities per bin
     bin_sums = torch.zeros(actual_n_bins, dtype=I.dtype, device=I.device)
-    bin_counts = torch.zeros(actual_n_bins, dtype=torch.long, device=I.device)  # dtype-ok: count accumulator; scatter_add source is long ones, dtype must match
+    bin_counts = torch.zeros(actual_n_bins, dtype=bin_indices.dtype, device=I.device)
     bin_sums.scatter_add_(0, bin_indices, I_sorted)
     bin_counts.scatter_add_(0, bin_indices, torch.ones_like(bin_indices))
 

torchref/base/metrics/binwise_scale.py

@@ -58,7 +58,8 @@ def binwise_scale(
     Fo = Fo.reshape(-1)
     device, dtype = Fc.device, Fc.dtype
 
-    bins = bins.reshape(-1).to(device=device, dtype=torch.int64)  # dtype-ok: resolution-bin indices used as scatter_add index; requires int64
+    # dtype-ok: scatter_add index; int64 required on torch < 2.8
+    bins = bins.reshape(-1).to(device=device, dtype=torch.int64)
     if nbins is None:
         nbins = int(bins.max().item()) + 1 if bins.numel() else 0
 

torchref/base/reciprocal/grid_operations.py

@@ -8,6 +8,8 @@
 
 import torch
 
+from torchref.config import get_int_dtype
+
 
 def place_on_grid(
     hkls, structure_factor, grid_size, enforce_hermitian: bool = True
@@ -47,22 +49,22 @@ def place_on_grid(
     dtype = structure_factor.dtype
     Nx, Ny, Nz = [int(x) for x in grid_size]
     hkls = hkls.to(device=device)
-    h = hkls[:, 0].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
-    k = hkls[:, 1].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
-    l = hkls[:, 2].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
+    h = hkls[:, 0].to(get_int_dtype())
+    k = hkls[:, 1].to(get_int_dtype())
+    l = hkls[:, 2].to(get_int_dtype())
 
     hi = torch.remainder(h, Nx)
     ki = torch.remainder(k, Ny)
     li = torch.remainder(l, Nz)
-    lin = (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64)  # (N,)  # dtype-ok: flat grid index (lin) for scatter/gather; requires int64
+    lin = (hi * (Ny * Nz) + ki * Nz + li).to(get_int_dtype())  # (N,)
     grid = torch.zeros((B, Nx * Ny * Nz), dtype=dtype, device=device)
     grid = grid.index_add(1, lin, structure_factor)  # (B, Nx*Ny*Nz)
 
     if enforce_hermitian:
         hi_sym = torch.remainder(-h, Nx)
         ki_sym = torch.remainder(-k, Ny)
         li_sym = torch.remainder(-l, Nz)
-        lin_sym = (hi_sym * (Ny * Nz) + ki_sym * Nz + li_sym).to(torch.int64)  # dtype-ok: symmetry flat grid index (lin_sym) for scatter/gather; requires int64
+        lin_sym = (hi_sym * (Ny * Nz) + ki_sym * Nz + li_sym).to(get_int_dtype())
         vals_conj = torch.conj(structure_factor)
         grid = grid.index_add(1, lin_sym, vals_conj)
 
@@ -101,9 +103,9 @@ def extract_structure_factor_from_grid(reciprocal_grid, hkls) -> torch.Tensor:
 
     # Same wrapping convention as place_on_grid.
     hkls = hkls.to(device=device)
-    h = hkls[:, 0].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
-    k = hkls[:, 1].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
-    l = hkls[:, 2].to(torch.int64)  # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long
+    h = hkls[:, 0].to(get_int_dtype())
+    k = hkls[:, 1].to(get_int_dtype())
+    l = hkls[:, 2].to(get_int_dtype())
 
     hi = torch.remainder(h, Nx)
     ki = torch.remainder(k, Ny)

torchref/base/reciprocal/symmetry.py

@@ -22,7 +22,7 @@ class through
 
 import torch
 
-from torchref.config import canonical_device
+from torchref.config import canonical_device, get_int_dtype
 from torchref.utils.autograd_ops import gather_with_index_add
 from torchref.utils.device_mixin import DeviceMixin
 
@@ -45,13 +45,14 @@ def _equiv_hkls_to_flat_indices(
     Returns
     -------
     torch.Tensor
-        Flat indices, shape ``(n_ops * N,)``, dtype ``int64``, wrapped modulo the grid.
+        Flat indices, shape ``(n_ops * N,)``, in the configured int dtype, wrapped
+        modulo the grid.
     """
     all_hkl = equiv_hkls.reshape(-1, 3)
     hi = torch.remainder(all_hkl[:, 0], Nx)
     ki = torch.remainder(all_hkl[:, 1], Ny)
     li = torch.remainder(all_hkl[:, 2], Nz)
-    return (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64)  # dtype-ok: flat HKL grid index; int64 avoids overflow, used for indexing
+    return (hi * (Ny * Nz) + ki * Nz + li).to(get_int_dtype())
 
 
 class ReciprocalSymmetryExtractor(DeviceMixin):

torchref/base/scattering/scattering_table.py

@@ -12,7 +12,7 @@
 
 import torch
 
-from torchref.config import get_float_dtype
+from torchref.config import get_float_dtype, get_int_dtype
 
 # Global cache for the loaded table
 _TABLE_CACHE: Optional[dict] = None
@@ -167,8 +167,7 @@ def get_scattering_params_by_z(
 
     table = load_scattering_table(device=device, dtype=dtype)
 
-    # Long, not the caller's int32: torch indexing requires it.
-    z_idx = z_tensor.to(device=device, dtype=torch.long)  # dtype-ok: z cast to long for scattering-table index lookup; indexing requires long
+    z_idx = z_tensor.to(device=device, dtype=get_int_dtype())
 
     A = table["A"][z_idx]
     B = table["B"][z_idx]
@@ -254,4 +253,4 @@ def elements_to_z(elements: list, normalize: bool = True) -> torch.Tensor:
         z = element_to_z.get(elem, 0)
         z_values.append(z)
 
-    return torch.tensor(z_values, dtype=torch.int32)  # dtype-ok: atomic-number Z categorical codes; fixed int32 lookup keys
+    return torch.tensor(z_values, dtype=get_int_dtype())

torchref/cli/mtz2map.py

@@ -18,13 +18,13 @@
 import numpy as np
 import torch
 
-from torchref.config import get_float_dtype
 from torchref.cli._common import (
     add_general_args,
     add_resolution_args,
     register_timing,
     parse_device_str,
 )
+from torchref.config import get_float_dtype, get_int_dtype
 
 
 def main():
@@ -246,7 +246,7 @@ def main():
               f"{d_spacings.max():.2f} - {d_spacings.min():.2f} A")
 
     # --- Convert to torch ---
-    hkl_t = torch.tensor(hkl, dtype=torch.int32, device=device)  # dtype-ok: hkl Miller indices fed to symmetry expand; fixed int32 crystallographic representation
+    hkl_t = torch.tensor(hkl, dtype=get_int_dtype(), device=device)
     amp_t = torch.tensor(amplitudes, dtype=get_float_dtype(), device=device)
     phi_t = torch.tensor(phases_deg, dtype=get_float_dtype(), device=device) * (np.pi / 180.0)
 

torchref/cli/validate_ded.py

@@ -46,6 +46,7 @@
     validate_cif_files,
     validate_files,
 )
+from torchref.config import get_int_dtype
 from torchref.maps.ded_weights import (
     DEFAULT_SCHEME,
     DedWeightFallbackWarning,
@@ -89,7 +90,7 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device):
         inv_frac_matrix=inv_frac,
     )
 
-    mask = torch.zeros(grid_shape, dtype=torch.int32, device=device)  # dtype-ok: integer solvent-mask accumulator (mask>0); categorical count, not model-precision data
+    mask = torch.zeros(grid_shape, dtype=get_int_dtype(), device=device)
     mask = add_to_solvent_mask(
         surrounding_coords,
         voxel_indices,