diff --git a/AGENTS.md b/AGENTS.md index d02be91..4c410f9 100644 --- a/AGENTS.md +++ b/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`). diff --git a/docs/changelog.rst b/docs/changelog.rst index 48350eb..4caae4e 100644 --- a/docs/changelog.rst +++ b/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 diff --git a/tests/unit/base/test_flat_grid_index.py b/tests/unit/base/test_flat_grid_index.py new file mode 100644 index 0000000..0e8255e --- /dev/null +++ b/tests/unit/base/test_flat_grid_index.py @@ -0,0 +1,23 @@ +"""Flat reciprocal-grid indices stay exact on grids larger than int32 can address. + +``h*Ny*Nz + k*Nz + l`` is a packed key: formed in the default int32 it wraps past +2**31 voxels and addresses aliased voxels, so it is built in int64 whatever dtype the +Miller indices arrive in. +""" + +import pytest +import torch + +from torchref.base.reciprocal.symmetry import _equiv_hkls_to_flat_indices +from torchref.config import get_int_dtype + +pytestmark = pytest.mark.unit + + +def test_flat_indices_are_exact_past_int32(): + n = 2048 # 2048**3 voxels; the helper never allocates the grid + hkl = torch.tensor([[[1000, -3, 7], [-1, 0, 2047]]], dtype=get_int_dtype()) + flat = _equiv_hkls_to_flat_indices(hkl, n, n, n) + expected = [(h % n) * n * n + (k % n) * n + (l % n) for h, k, l in hkl[0].tolist()] + assert flat.dtype == torch.int64 + assert flat.tolist() == expected diff --git a/tests/unit/model/test_disorder_field.py b/tests/unit/model/test_disorder_field.py index dcb8f11..506447d 100644 --- a/tests/unit/model/test_disorder_field.py +++ b/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] diff --git a/torchref/base/electron_density/kernels/cpu/jit_reference.py b/torchref/base/electron_density/kernels/cpu/jit_reference.py index 257fa6d..be65d44 100644 --- a/torchref/base/electron_density/kernels/cpu/jit_reference.py +++ b/torchref/base/electron_density/kernels/cpu/jit_reference.py @@ -135,8 +135,12 @@ def forward( # Scatter add to density map ny: int = density_map.shape[1] nz: int = density_map.shape[2] + # Compiled by TorchScript: no Tensor.new_tensor, no config dtype getters. 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 + [ny * nz, nz, 1], + device=voxel_indices.device, + # dtype-ok: int64 strides make the flat voxel index int64; scatter_add_ requires int64 on torch < 2.8 + 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 diff --git a/torchref/base/electron_density/kernels/cpu/variable_radius.py b/torchref/base/electron_density/kernels/cpu/variable_radius.py index bb19ae9..3feab13 100644 --- a/torchref/base/electron_density/kernels/cpu/variable_radius.py +++ b/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,12 +92,13 @@ 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 - center_1d = ((center_idx[:, 0] % nx) * (ny * nz) - + (center_idx[:, 1] % ny) * nz + (center_idx[:, 2] % nz)) + # dtype-ok: the flat voxel index overflows int32 above 2**31 voxels + c = center_idx.to(torch.int64) + center_1d = (c[:, 0] % nx) * (ny * nz) + (c[:, 1] % ny) * nz + (c[:, 2] % nz) order, spans = _bucket_by_radius(radius_per_atom, center_1d) return order, spans, center_idx[order], w0[order] @@ -111,8 +116,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 +157,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) diff --git a/torchref/base/electron_density/map_building.py b/torchref/base/electron_density/map_building.py index d00d766..4fe71c1 100644 --- a/torchref/base/electron_density/map_building.py +++ b/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 diff --git a/torchref/base/electron_density/solvent_mask.py b/torchref/base/electron_density/solvent_mask.py index 2e36e70..018e39a 100644 --- a/torchref/base/electron_density/solvent_mask.py +++ b/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 diff --git a/torchref/base/electron_density/voxel_utils.py b/torchref/base/electron_density/voxel_utils.py index 3fc903d..1551c25 100644 --- a/torchref/base/electron_density/voxel_utils.py +++ b/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 diff --git a/torchref/base/french_wilson.py b/torchref/base/french_wilson.py index a03c6df..fbdcc96 100644 --- a/torchref/base/french_wilson.py +++ b/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)) diff --git a/torchref/base/metrics/binwise_scale.py b/torchref/base/metrics/binwise_scale.py index 2d803ef..0bfffc1 100644 --- a/torchref/base/metrics/binwise_scale.py +++ b/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 diff --git a/torchref/base/reciprocal/grid_operations.py b/torchref/base/reciprocal/grid_operations.py index a631a64..2aabbd1 100644 --- a/torchref/base/reciprocal/grid_operations.py +++ b/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 @@ -46,15 +48,16 @@ def place_on_grid( device = structure_factor.device 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 + # dtype-ok: the flat index h*Ny*Nz + k*Nz + l overflows int32 above 2**31 voxels + hkls = hkls.to(device=device, dtype=torch.int64) + h = hkls[:, 0] + k = hkls[:, 1] + l = hkls[:, 2] 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 # (N,) grid = torch.zeros((B, Nx * Ny * Nz), dtype=dtype, device=device) grid = grid.index_add(1, lin, structure_factor) # (B, Nx*Ny*Nz) @@ -62,7 +65,7 @@ def place_on_grid( 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 vals_conj = torch.conj(structure_factor) grid = grid.index_add(1, lin_sym, vals_conj) @@ -101,9 +104,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) diff --git a/torchref/base/reciprocal/symmetry.py b/torchref/base/reciprocal/symmetry.py index 390a13b..04c6ecd 100644 --- a/torchref/base/reciprocal/symmetry.py +++ b/torchref/base/reciprocal/symmetry.py @@ -47,11 +47,12 @@ def _equiv_hkls_to_flat_indices( torch.Tensor Flat indices, shape ``(n_ops * N,)``, dtype ``int64``, wrapped modulo the grid. """ - all_hkl = equiv_hkls.reshape(-1, 3) + # dtype-ok: the flat index h*Ny*Nz + k*Nz + l overflows int32 above 2**31 voxels + all_hkl = equiv_hkls.reshape(-1, 3).to(torch.int64) 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 class ReciprocalSymmetryExtractor(DeviceMixin): diff --git a/torchref/base/scattering/scattering_table.py b/torchref/base/scattering/scattering_table.py index fbc26aa..8a3d833 100644 --- a/torchref/base/scattering/scattering_table.py +++ b/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()) diff --git a/torchref/cli/mtz2map.py b/torchref/cli/mtz2map.py index e98fe12..621a833 100644 --- a/torchref/cli/mtz2map.py +++ b/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) diff --git a/torchref/cli/validate_ded.py b/torchref/cli/validate_ded.py index 18b6d9a..ea89022 100644 --- a/torchref/cli/validate_ded.py +++ b/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, diff --git a/torchref/experimental/alignment/frf/data_mr.py b/torchref/experimental/alignment/frf/data_mr.py index 3ddc744..9d5d506 100644 --- a/torchref/experimental/alignment/frf/data_mr.py +++ b/torchref/experimental/alignment/frf/data_mr.py @@ -19,6 +19,8 @@ import torch +from torchref.config import get_int_dtype + _PROFILE = bool(os.environ.get("FRF_PROFILE")) #: Byte budget for the per-chunk transients in :func:`bessel_sh_expand`. @@ -143,7 +145,7 @@ def spherical_bessel_table( inv_threshold = 1.0 / threshold # Rescales applied so far, per element. Every element's ladder sits in the # single frame 2**(-_BESSEL_RESCALE_EXP * n_rescales). - n_rescales = torch.zeros_like(x64, dtype=torch.int32) # dtype-ok: small integer counter + n_rescales = torch.zeros_like(x64, dtype=get_int_dtype()) for n in range(n_start, 0, -1): j_low = (2.0 * n + 1.0) * inv_x * j_mid - j_high @@ -162,7 +164,7 @@ def spherical_bessel_table( j_high = j_high * factor if n - 1 <= u_max: j_table[n - 1:] = j_table[n - 1:] * factor - n_rescales = n_rescales + over.to(torch.int32) # dtype-ok: small integer counter + n_rescales = n_rescales + over.to(get_int_dtype()) true_j0 = torch.sin(x64) * inv_x true_j0 = torch.where(x64 < 1e-30, torch.ones_like(x64), true_j0) @@ -301,15 +303,15 @@ def bessel_sh_expand( n_list.append(n) u_list.append(u) w_list.append(math.sqrt(float(2 * u + 1))) - l_idx = torch.tensor(l_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 - n_idx = torch.tensor(n_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 - u_idx = torch.tensor(u_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + l_idx = torch.tensor(l_list, dtype=get_int_dtype(), device=device) + n_idx = torch.tensor(n_list, dtype=get_int_dtype(), device=device) + u_idx = torch.tensor(u_list, dtype=get_int_dtype(), device=device) w_vec = torch.tensor(w_list, dtype=comp_real, device=device) # Only even degrees l ∈ [2, lmax_even] carry signal (odd-l and l=0 are zeroed # by Patterson centrosymmetry). Compute / contract Y_lm on these rows only — # the assembly + einsum are the bottleneck, so this ~halves them. The full # c_nlm keeps the (L, ...) shape with odd/zero rows left at zero. - even_l_idx = torch.tensor(even_ls, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + even_l_idx = torch.tensor(even_ls, dtype=get_int_dtype(), device=device) M = s_vectors.shape[0] einsum_dtype = complex_dtype @@ -347,8 +349,10 @@ def _tick(t0): s_key = s_vectors.detach().cpu().to(torch.float64) # dtype-ok: exact clustering key on the host; the device never sees it s_mag_key = s_key.norm(dim=-1).clamp(min=1e-30) cos_key = (s_key[..., 2] / s_mag_key).clamp(min=-1.0, max=1.0) - k_s = (s_mag_key * _GROUP_SCALE_S).round().to(torch.int64) # dtype-ok: exact clustering key - k_c = (cos_key * _GROUP_SCALE_COS).round().to(torch.int64) + _GROUP_SCALE_COS # dtype-ok: exact clustering key + # dtype-ok: clustering key k_s*(2e7+1)+k_c overflows int32 + k_s = (s_mag_key * _GROUP_SCALE_S).round().to(torch.int64) + # dtype-ok: clustering key k_s*(2e7+1)+k_c overflows int32 + k_c = (cos_key * _GROUP_SCALE_COS).round().to(torch.int64) + _GROUP_SCALE_COS key = (k_s * (2 * _GROUP_SCALE_COS + 1) + k_c).to(s_vectors.device) uniq_key, inverse = torch.unique(key, return_inverse=True) n_clusters = int(uniq_key.shape[0]) @@ -379,7 +383,8 @@ def _group_mean(values, index, n_groups): # index to meet device values. inv_s = inv_s.to(device) n_shells = int(uniq_ks.shape[0]) - shell_of_cluster = torch.zeros(n_clusters, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + # dtype-ok: the legendre_shell kernel TORCH_CHECKs int64 shell labels + shell_of_cluster = torch.zeros(n_clusters, dtype=torch.int64, device=device) shell_of_cluster[inverse] = inv_s shell_smag = _group_mean(s_mag_all.to(comp_real), inv_s, n_shells) diff --git a/torchref/experimental/alignment/frf/dense_calc.py b/torchref/experimental/alignment/frf/dense_calc.py index d95921c..2c11c01 100644 --- a/torchref/experimental/alignment/frf/dense_calc.py +++ b/torchref/experimental/alignment/frf/dense_calc.py @@ -21,7 +21,7 @@ import torch -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype if TYPE_CHECKING: from torchref.model import ModelFT @@ -80,9 +80,9 @@ def dense_calc_via_box( nmax = int(math.ceil(a / d_min)) idx = torch.arange(-nmax, nmax + 1, device=dev) H, K, Lg = torch.meshgrid(idx, idx, idx, indexing="ij") - hkl = torch.stack( - [H.reshape(-1), K.reshape(-1), Lg.reshape(-1)], dim=-1 - ).to(torch.long) # dtype-ok: Miller indices are integers + hkl = torch.stack([H.reshape(-1), K.reshape(-1), Lg.reshape(-1)], dim=-1).to( + get_int_dtype() + ) # Cubic box: |s| = |hkl| / a. real = get_float_dtype() smag = hkl.to(real).norm(dim=-1) / a diff --git a/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py b/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py index e7d5674..07f5478 100644 --- a/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py +++ b/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py @@ -246,7 +246,8 @@ def shell_offsets(shell: torch.Tensor, n_shells: int) -> torch.Tensor: write their accumulator rows without atomics. """ counts = torch.bincount(shell, minlength=n_shells) - offsets = torch.zeros(n_shells + 1, dtype=torch.long, device=shell.device) # dtype-ok: index tensor; index_add_/gather need int64 + # dtype-ok: the kernel TORCH_CHECKs int64 offsets + offsets = torch.zeros(n_shells + 1, dtype=torch.int64, device=shell.device) torch.cumsum(counts, dim=0, out=offsets[1:]) return offsets diff --git a/torchref/experimental/alignment/frf/peak_finder.py b/torchref/experimental/alignment/frf/peak_finder.py index 19692db..ddb8e56 100644 --- a/torchref/experimental/alignment/frf/peak_finder.py +++ b/torchref/experimental/alignment/frf/peak_finder.py @@ -29,6 +29,8 @@ import torch +from torchref.config import get_int_dtype + from ....base.alignment.rotation import rotation_matrix_euler_zyz from .types import AdaptiveRotationFunction, RotationPeak @@ -55,7 +57,7 @@ def _so3_greedy_nms( """ n = values.shape[0] if n == 0: - return torch.empty(0, dtype=torch.int64, device=values.device) # dtype-ok: index tensor; index_add_/gather need int64 + return torch.empty(0, dtype=get_int_dtype(), device=values.device) # The greedy walk is inherently sequential and latency-bound; on GPU a # per-iteration `.item()` sync would dominate. Move the (tiny) candidate # rotations to CPU once and run the loop there with no device syncs, a @@ -97,7 +99,7 @@ def _so3_greedy_nms( count += 1 if count >= keep_at_most: break - return torch.tensor(kept_idx, dtype=torch.int64, device=values.device) # dtype-ok: index tensor; index_add_/gather need int64 + return torch.tensor(kept_idx, dtype=get_int_dtype(), device=values.device) def find_rotation_peaks( diff --git a/torchref/experimental/alignment/frf/preprocessing.py b/torchref/experimental/alignment/frf/preprocessing.py index 6397781..b8eed82 100644 --- a/torchref/experimental/alignment/frf/preprocessing.py +++ b/torchref/experimental/alignment/frf/preprocessing.py @@ -298,8 +298,8 @@ def fit_relative_wilson_b( F2_calc = (F_calc * F_calc).to(real) s2_obs = (s_mag * s_mag).to(real) - counts_obs = torch.zeros(n_shells, dtype=torch.int64, device=s_mag.device) # dtype-ok: per-shell counts - counts_calc = torch.zeros(n_shells, dtype=torch.int64, device=s_mag.device) # dtype-ok: per-shell counts + counts_obs = torch.zeros(n_shells, dtype=shell_idx_obs.dtype, device=s_mag.device) + counts_calc = torch.zeros(n_shells, dtype=shell_idx_calc.dtype, device=s_mag.device) sum_F2obs = torch.zeros(n_shells, dtype=real, device=s_mag.device) sum_F2calc = torch.zeros(n_shells, dtype=real, device=s_mag.device) sum_s2 = torch.zeros(n_shells, dtype=real, device=s_mag.device) diff --git a/torchref/experimental/alignment/frf/sitelist_ang.py b/torchref/experimental/alignment/frf/sitelist_ang.py index c297c3c..be9bcb7 100644 --- a/torchref/experimental/alignment/frf/sitelist_ang.py +++ b/torchref/experimental/alignment/frf/sitelist_ang.py @@ -40,6 +40,8 @@ import torch +from torchref.config import get_int_dtype + from ....config import canonical_device from ....symmetry.symmetry import find_fft_friendly_size from .types import AdaptiveRotationFunction @@ -98,7 +100,7 @@ def build_dense_map_per_beta( (n_beta, fft_size, fft_size), dtype=S.dtype, device=device, ) m_vals = torch.arange(-(L - 1), L, device=device) - idx = (m_vals % fft_size).to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 + idx = (m_vals % fft_size).to(get_int_dtype()) pad[:, idx.unsqueeze(1), idx.unsqueeze(0)] = S # 3. Forward 2D FFT — torch convention: @@ -210,8 +212,10 @@ def build_adaptive_sample_list( # original dict scan, but no host sync / Python loop. # Hash the two rounded fracs (each in [0, 1e6]) into one int64 so we # can use the fast 1-D unique instead of a 2-D row lexsort. - a_round = (alpha_frac * 1_000_000).round().to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 - g_round = (gamma_frac * 1_000_000).round().to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 + # dtype-ok: a_round*1_000_001+g_round overflows int32 + a_round = (alpha_frac * 1_000_000).round().to(torch.int64) + # dtype-ok: a_round*1_000_001+g_round overflows int32 + g_round = (gamma_frac * 1_000_000).round().to(torch.int64) key_hash = a_round * 1_000_001 + g_round _, uniq_idx = torch.unique(key_hash, return_inverse=True) n = uniq_idx.shape[0] @@ -234,7 +238,7 @@ def build_adaptive_sample_list( alphas = torch.cat(alphas_list).to(device) gammas = torch.cat(gammas_list).to(device) betas_flat = torch.cat(betas_list).to(device) - beta_starts_t = torch.tensor(beta_starts, dtype=torch.int64, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + beta_starts_t = torch.tensor(beta_starts, dtype=get_int_dtype(), device=device) b = torch.arange(bmax, dtype=torch.float64, device=cpu) # dtype-ok: sample-list geometry follows the accumulator's width betas_rad = (b * grid_sampling_deg * deg2rad).to(device=device, dtype=dtype) @@ -256,8 +260,8 @@ def _bilinear_interp_periodic( N = M.shape[-1] af = (alpha_frac % 1.0) * N gf = (gamma_frac % 1.0) * N - a0 = torch.floor(af).to(torch.int64) % N # dtype-ok: index tensor; index_add_/gather need int64 - g0 = torch.floor(gf).to(torch.int64) % N # dtype-ok: index tensor; index_add_/gather need int64 + a0 = torch.floor(af).to(get_int_dtype()) % N + g0 = torch.floor(gf).to(get_int_dtype()) % N a1 = (a0 + 1) % N g1 = (g0 + 1) % N da = (af - torch.floor(af)).to(M.real.dtype) diff --git a/torchref/experimental/alignment/sh.py b/torchref/experimental/alignment/sh.py index 0da057f..6b94c1f 100644 --- a/torchref/experimental/alignment/sh.py +++ b/torchref/experimental/alignment/sh.py @@ -29,6 +29,8 @@ import torch +from torchref.config import get_int_dtype + from ...config import get_float_dtype @@ -134,10 +136,10 @@ def _bar_legendre_recurrence( if keep_l is None: rows = torch.arange(L, device=device) else: - rows = keep_l.to(device=device, dtype=torch.long) # dtype-ok: index tensor; index_add_/gather need int64 + rows = keep_l.to(device=device, dtype=get_int_dtype()) # l -> its position in the output, or -1 when it is not kept. - where = torch.full((L,), -1, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 - where[rows] = torch.arange(rows.numel(), device=device) + where = torch.full((L,), -1, dtype=get_int_dtype(), device=device) + where[rows] = torch.arange(rows.numel(), device=device, dtype=where.dtype) where_list = where.tolist() out = torch.zeros((*batch_shape, rows.numel(), L), dtype=dtype, device=device) @@ -338,7 +340,7 @@ def fit_overall_anisotropy( F, s, idx, cen = F[ok], s[ok], idx[ok], cen[ok] I = F * F - count = torch.zeros(P, dtype=torch.int64, device=F.device) # dtype-ok: index tensor; index_add_/gather need int64 + count = torch.zeros(P, dtype=idx.dtype, device=F.device) total = torch.zeros(P, dtype=work, device=F.device) count.index_add_(0, idx, torch.ones_like(idx)) total.index_add_(0, idx, I) @@ -536,7 +538,7 @@ def compute_patterson_shell_variance( valid = shell_idx >= 0 patt_v = patt[valid] idx_v = shell_idx[valid] - count = torch.zeros(P, dtype=torch.int64, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + count = torch.zeros(P, dtype=idx_v.dtype, device=device) count.index_add_(0, idx_v, torch.ones_like(idx_v)) sum1 = torch.zeros(P, dtype=dtype, device=device) sum2 = torch.zeros(P, dtype=dtype, device=device) diff --git a/torchref/experimental/alignment/translation.py b/torchref/experimental/alignment/translation.py index 9791397..0fdb41b 100644 --- a/torchref/experimental/alignment/translation.py +++ b/torchref/experimental/alignment/translation.py @@ -38,7 +38,12 @@ import torch from torchref.base.targets.xray_likelihoods import rice_per_refl -from torchref.config import get_complex_dtype, get_default_device, get_float_dtype +from torchref.config import ( + get_complex_dtype, + get_default_device, + get_float_dtype, + get_int_dtype, +) from torchref.scaling import WilsonNormaliser from torchref.scaling.weighting import (inverse_variance_weight, normalise_weight, snr_from_amplitude) @@ -151,7 +156,7 @@ def build( rec_basis = real_cell.reciprocal_basis_matrix.to(device=dev, dtype=real) s_mag = (hkl_i.to(real) @ rec_basis).norm(dim=-1) - hkl_l = hkl_i.round().to(torch.int64) # dtype-ok: Miller indices are integers + hkl_l = hkl_i.round().to(get_int_dtype()) # friedel=False: Wilson's = eps*Sigma counts the operations mapping # h to itself, which add coherently and set the mean. The Friedel-folded # branch changes the distribution instead, and that is centricity -- @@ -249,7 +254,7 @@ def prepare_candidate( # h_R[i, n, d] = sum_e hkl[n, e] sym_R[i, e, d]: the h.S convention. h_R = torch.einsum("ne,ied->ind", hkl, sym_R) phase = torch.exp((2j * math.pi) * torch.einsum("ne,ie->in", hkl, sym_t).to(cplx)) - hkl_SN = h_R.reshape(-1, 3).round().to(torch.int64).to(model_p1.xyz().device) # dtype-ok: Miller indices are integers + hkl_SN = h_R.reshape(-1, 3).round().to(get_int_dtype()).to(model_p1.xyz().device) with torch.no_grad(): F_all = model_p1(hkl_SN).to(device).reshape(S, N).to(cplx) G_raw = F_all * phase @@ -395,7 +400,8 @@ def fast_translation_function( G = cand.G.to(device=device, dtype=cplx) S, N = G.shape coeff = obs.coeff.to(device=device, dtype=cplx) - h_R_int = cand.h_R.round().to(torch.int64) # dtype-ok: Miller indices are integers + # dtype-ok: the flat translation-grid index overflows int32 above 2**31 grid points + h_R_int = cand.h_R.round().to(torch.int64) # The pair (j, i) is the conjugate of (i, j) at -dh, so the map is twice # the real part of the upper triangle's transform plus the diagonal, which diff --git a/torchref/experimental/ensemble/ensemble_amber_kl.py b/torchref/experimental/ensemble/ensemble_amber_kl.py index ebb5cfc..1994cf8 100644 --- a/torchref/experimental/ensemble/ensemble_amber_kl.py +++ b/torchref/experimental/ensemble/ensemble_amber_kl.py @@ -53,6 +53,7 @@ import numpy as np import torch +from torchref.config import get_int_dtype from torchref.experimental.targets.amber_target import AMBER14_STANDARD, AmberTarget if TYPE_CHECKING: @@ -136,7 +137,7 @@ def __init__( self.register_buffer( "_member_atom_idx", torch.as_tensor( - atom_idx_np, dtype=torch.long, device=self._model.device # dtype-ok: atom index tensor for indexing; PyTorch requires int64 + atom_idx_np, dtype=get_int_dtype(), device=self._model.device ), ) else: diff --git a/torchref/experimental/ensemble/quasi_crystal_amber.py b/torchref/experimental/ensemble/quasi_crystal_amber.py index ee9986c..bf90cd9 100644 --- a/torchref/experimental/ensemble/quasi_crystal_amber.py +++ b/torchref/experimental/ensemble/quasi_crystal_amber.py @@ -60,6 +60,7 @@ import numpy as np import torch +from torchref.config import get_int_dtype from torchref.experimental.targets.amber_target import ( AmberTarget, _OpenMMAMBERFunction, @@ -564,19 +565,19 @@ def _ensure_torch_buffers(self, device: torch.device, dtype: torch.dtype) -> Non # Index pairs (long) for the scatter from model atoms into OMM slots. self._src_model_idx_torch = torch.from_numpy(self._src_model_idx_np).to( device=device, - dtype=torch.long, # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 + dtype=get_int_dtype(), ) self._dst_omm_idx_torch = torch.from_numpy(self._dst_omm_idx_np).to( device=device, - dtype=torch.long, # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 + dtype=get_int_dtype(), ) # Index of ensemble-model atoms (in the FULL EnsembleModel layout) # that survived the special-position filter — used in forward to # subset ``xyz_per_member`` before applying the layout transform. self._keep_atom_idx_torch = torch.from_numpy(self._keep_atom_idx_np).to( - device=device, dtype=torch.long - ) # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 + device=device, dtype=get_int_dtype() + ) self._omm_to_model = self._omm_to_model.to(device) self._buffers_device = device @@ -610,7 +611,6 @@ def _compose_full_omm_xyz(self, supercell_xyz_nm: torch.Tensor) -> torch.Tensor: ) return supercell_xyz_nm.index_select(1, self._omm_to_model).reshape(-1, 3) - # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ diff --git a/torchref/experimental/ensemble/wilson_prior.py b/torchref/experimental/ensemble/wilson_prior.py index 44490f9..8fce242 100644 --- a/torchref/experimental/ensemble/wilson_prior.py +++ b/torchref/experimental/ensemble/wilson_prior.py @@ -172,7 +172,8 @@ def _build_bin_assignment(self) -> None: order = torch.argsort(res) n = res.numel() nbins = min(self.nbins, max(1, n // 50)) - bin_assign = torch.empty(n, dtype=torch.long, device=res.device) # dtype-ok: bin-assignment tensor used as scatter_add index; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + bin_assign = torch.empty(n, dtype=torch.long, device=res.device) edges = torch.linspace(0, n, nbins + 1, device=res.device).round().long() for b in range(nbins): start = int(edges[b].item()) diff --git a/torchref/experimental/monolithic_refinement/density_scaler.py b/torchref/experimental/monolithic_refinement/density_scaler.py index afe76f9..1f76c1d 100644 --- a/torchref/experimental/monolithic_refinement/density_scaler.py +++ b/torchref/experimental/monolithic_refinement/density_scaler.py @@ -36,7 +36,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_default_device, get_float_dtype, get_int_dtype from torchref.scaling.scaler import Scaler from torchref.scaling.solvent import SolventModel from torchref.experimental.monolithic_refinement.density_solvent import ( @@ -127,7 +127,7 @@ def get_rec_solvent(self, hkl): Not detached: ``F_sol`` follows the moving atoms so gradients reach ``xyz``/``adp``. The scaler applies the contrast and falloff on top. """ - return self.density(hkl.to(torch.long)) # dtype-ok: hkl cast to long for density lookup indexing; PyTorch requires int64 + return self.density(hkl.to(get_int_dtype())) def update_solvent(self): """No-op: the density mask is rebuilt live on every scaler forward.""" diff --git a/torchref/experimental/targets/forcefield_target.py b/torchref/experimental/targets/forcefield_target.py index 09f0e1d..62290cc 100644 --- a/torchref/experimental/targets/forcefield_target.py +++ b/torchref/experimental/targets/forcefield_target.py @@ -178,7 +178,8 @@ def forward(self) -> torch.Tensor: Z = self.model.Z # Shape: (n_atoms,) # Ensure Z is long tensor - if Z.dtype != torch.long: # dtype-ok: dtype guard comparison against torch.long, not an allocation + # dtype-ok: TorchMD-Net expects a LongTensor Z; external library contract + if Z.dtype != torch.long: Z = Z.long() # Create batch tensor (single structure = all zeros) diff --git a/torchref/experimental/targets/sampled_ml_phase_target.py b/torchref/experimental/targets/sampled_ml_phase_target.py index e505240..7d4fa1c 100644 --- a/torchref/experimental/targets/sampled_ml_phase_target.py +++ b/torchref/experimental/targets/sampled_ml_phase_target.py @@ -18,6 +18,7 @@ import torch from typing import TYPE_CHECKING, Dict, Tuple +from torchref.config import get_int_dtype from torchref.refinement.targets.base import Target from torchref.refinement.targets.xray import XrayTarget from torchref.utils.stats import ( @@ -129,7 +130,9 @@ def __init__( self.name = "xray_sampled_ml_work" if use_work_set else "xray_sampled_ml_test" # Register tunable parameters as buffers for state_dict access - self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) # dtype-ok: scalar sample-count buffer; categorical count, not model-precision data + self.register_buffer( + "_n_samples", torch.tensor(n_samples, dtype=get_int_dtype()) + ) self.register_buffer("_sigma_model_log", torch.tensor(sigma_model_log)) self.register_buffer("_use_analytical", torch.tensor(use_analytical)) self.register_buffer("_use_antithetic", torch.tensor(use_antithetic)) @@ -545,7 +548,9 @@ def __init__( self.add_module("_scaler_dark", scaler_dark) # Tunable parameters as buffers - self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) # dtype-ok: scalar sample-count buffer; categorical count, not model-precision data + self.register_buffer( + "_n_samples", torch.tensor(n_samples, dtype=get_int_dtype()) + ) self.register_buffer("_sigma_model_log", torch.tensor(sigma_model_log)) self.use_work_set = use_work_set diff --git a/torchref/io/datasets/fcalc_data.py b/torchref/io/datasets/fcalc_data.py index d4a4a2a..8b2947f 100644 --- a/torchref/io/datasets/fcalc_data.py +++ b/torchref/io/datasets/fcalc_data.py @@ -12,7 +12,7 @@ import pandas as pd import torch -from torchref.config import get_float_dtype, normalize_device +from torchref.config import get_float_dtype, get_int_dtype, normalize_device from torchref.symmetry import Cell, SpaceGroup, SpaceGroupLike from .base import CrystalDataset @@ -133,7 +133,7 @@ def from_cell_and_resolution( # make_miller_array returns unique HKL for the asymmetric unit only. hkl_list = gemmi.make_miller_array(gemmi_cell, gemmi_sg, d_min) - hkl = torch.tensor(hkl_list, dtype=torch.int32, device=device) # dtype-ok: hkl Miller indices; fixed int32 crystallographic representation, not model-precision data + hkl = torch.tensor(hkl_list, dtype=get_int_dtype(), device=device) resolution = get_d_spacing(hkl.float(), cell_tensor) diff --git a/torchref/io/datasets/reflection_data.py b/torchref/io/datasets/reflection_data.py index 58722ed..9893923 100644 --- a/torchref/io/datasets/reflection_data.py +++ b/torchref/io/datasets/reflection_data.py @@ -17,7 +17,7 @@ from torchref.base import math_torch from torchref.base.french_wilson import FrenchWilson -from torchref.config import dtypes, normalize_device +from torchref.config import dtypes, get_int_dtype, normalize_device from torchref.io import cif, mtz from torchref.io.datasets.base import CrystalDataset from torchref.symmetry import Cell, SpaceGroup @@ -306,7 +306,7 @@ def _subset_indices(self, kind: str) -> torch.Tensor: n = 0 if self.hkl is None else len(self.hkl) device = self.device if n == 0: - empty = torch.empty(0, dtype=torch.long, device=device) # dtype-ok: empty index tensor; PyTorch requires int64 for indexing + empty = torch.empty(0, dtype=get_int_dtype(), device=device) self._subset_cache = { "work": empty, "free": empty, @@ -428,7 +428,7 @@ def _reindex_per_reflection( n_src = len(self.hkl) if self.hkl is not None else 0 new_hkl = new_hkl.to(dtype=dtypes.int, device=self.device) n_out = len(new_hkl) - index_map = index_map.to(device=self.device, dtype=torch.long) # dtype-ok: index map used for indexing/gather; PyTorch requires int64 + index_map = index_map.to(device=self.device, dtype=get_int_dtype()) present = index_map >= 0 src_idx = index_map[present] @@ -448,7 +448,7 @@ def _reindex_per_reflection( # Present rows keep their signed (anomalous) index; missing rows # fall back to the canonical reference HKL (never a 0,0,0 row). out = new_hkl.clone() - out[present] = val[src_idx] + out[present] = val[src_idx].to(out.dtype) else: fill = self._REINDEX_FILL.get(name, 0) out = torch.full( @@ -1357,13 +1357,15 @@ def mean_res_per_bin(self) -> torch.Tensor: mean_resolutions = torch.scatter_add( mean_resolutions, 0, - self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add/index; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), self.resolution[mask], ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add/index; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), torch.ones_like(self.resolution[mask], dtype=dtypes.int), ) mean_resolutions = mean_resolutions / count_per_bin.clamp(min=1).float() @@ -1392,12 +1394,17 @@ def mean_F_per_bin(self) -> torch.Tensor: count_per_bin = torch.zeros(self._n_bins, dtype=dtypes.int, device=self.device) mask = self.masks() mean_F = torch.scatter_add( - mean_F, 0, self.bin_indices[mask].to(torch.int64), self.F[mask] # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 + mean_F, + 0, + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), + self.F[mask], ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), torch.ones_like(self.F[mask], dtype=dtypes.int), ) mean_F = mean_F / count_per_bin.clamp(min=1).float() @@ -1426,12 +1433,17 @@ def mean_sigma_per_bin(self) -> Optional[torch.Tensor]: count_per_bin = torch.zeros(self._n_bins, dtype=dtypes.int, device=self.device) mask = self.masks() mean_sigma = torch.scatter_add( - mean_sigma, 0, self.bin_indices[mask].to(torch.int64), self.F_sigma[mask] # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 + mean_sigma, + 0, + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), + self.F_sigma[mask], ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + self.bin_indices[mask].to(torch.int64), torch.ones_like(self.F_sigma[mask], dtype=dtypes.int), ) mean_sigma = mean_sigma / count_per_bin.clamp(min=1).float() @@ -2507,9 +2519,9 @@ def _build_anomalous_dataframe( uniq = hkl[self._group_representative_rows(inverse, M)] # The (+) member is the unconjugated row, (-) is the Friedel-flagged row. - arange = torch.arange(N) - plus_idx = torch.full((M,), -1, dtype=torch.long) # dtype-ok: Friedel-mate index map (-1 sentinel) for indexing; PyTorch requires int64 - minus_idx = torch.full((M,), -1, dtype=torch.long) # dtype-ok: Friedel-mate index map (-1 sentinel) for indexing; PyTorch requires int64 + arange = torch.arange(N, dtype=get_int_dtype()) + plus_idx = torch.full((M,), -1, dtype=get_int_dtype()) + minus_idx = torch.full((M,), -1, dtype=get_int_dtype()) # A Bijvoet mate only counts as present if it is a real, positive # observation. Stacked anomalous input (rs.stack_anomalous) carries a # row for every *absent* mate with a NaN intensity, which French-Wilson @@ -2947,7 +2959,7 @@ def remap( ---------- new_hkl : torch.Tensor, shape (M, 3) New Miller indices. - index_mapping : torch.Tensor, shape (M,), dtype int64 + index_mapping : torch.Tensor, shape (M,), integer dtype Maps new indices to original: ``new[i] = old[index_mapping[i]]`` Values of -1 indicate missing reflections (filled with defaults). phase_shifts : torch.Tensor, optional, shape (M,) diff --git a/torchref/model/disorder_field.py b/torchref/model/disorder_field.py index a2911d3..582997e 100644 --- a/torchref/model/disorder_field.py +++ b/torchref/model/disorder_field.py @@ -88,7 +88,7 @@ def farthest_point_anchors(xyz: torch.Tensor, n_nodes: int) -> torch.Tensor: chosen.append(nxt) d2_nearest = torch.minimum(d2_nearest, ((xyz - xyz[nxt]) ** 2).sum(-1)) - anchors = torch.tensor(chosen, dtype=torch.int64, device=xyz.device) # dtype-ok: anchor atom indices; torch indexing requires int64 + anchors = torch.tensor(chosen, dtype=get_int_dtype(), device=xyz.device) # Lloyd relaxation, snapping to real atoms so an anchor is always an atom index. for _ in range(10): @@ -126,7 +126,7 @@ def density_anchor_rows(xyz: torch.Tensor, n_nodes: int): """ seeds = farthest_point_anchors(xyz, n_nodes) assign = torch.cdist(xyz, xyz[seeds]).argmin(dim=1) - atom_idx = torch.arange(xyz.shape[0], dtype=torch.int64, device=xyz.device) # dtype-ok: arange atom indices; index requires int64 + atom_idx = torch.arange(xyz.shape[0], dtype=get_int_dtype(), device=xyz.device) # A seed whose cluster somehow came out empty still needs a position. present = torch.bincount(assign, minlength=seeds.shape[0]) > 0 @@ -715,12 +715,12 @@ def __init__( if anchor_rows is None: anchor_atom = farthest_point_anchors(xyz, n_nodes) anchor_node = torch.arange( - anchor_atom.shape[0], dtype=torch.int64, device=device # dtype-ok: arange anchor indices; index requires int64 + anchor_atom.shape[0], dtype=get_int_dtype(), device=device ) else: anchor_atom, anchor_node = anchor_rows - anchor_atom = anchor_atom.to(device=device, dtype=torch.int64) # dtype-ok: anchor_atom indices cast; index requires int64 - anchor_node = anchor_node.to(device=device, dtype=torch.int64) # dtype-ok: anchor_node indices cast; index requires int64 + anchor_atom = anchor_atom.to(device=device, dtype=get_int_dtype()) + anchor_node = anchor_node.to(device=device, dtype=get_int_dtype()) n_k = int(anchor_node.max()) + 1 node_pos = self._segment_mean(xyz, anchor_atom, anchor_node, n_k) diff --git a/torchref/model/model.py b/torchref/model/model.py index 56d9a8d..879eff4 100644 --- a/torchref/model/model.py +++ b/torchref/model/model.py @@ -25,6 +25,7 @@ canonical_device, get_default_device, get_float_dtype, + get_int_dtype, normalize_device, ) from torchref.io import cif, pdb @@ -334,7 +335,6 @@ def _iso_covers_all(self) -> bool: def _aniso_is_empty(self) -> bool: return self._sf_partition()[3] - # ========================================================================= # Cell, SpaceGroup, and Symmetry properties # ========================================================================= @@ -435,7 +435,7 @@ def _build_z_tensor(self) -> torch.Tensor: for elem in self.pdb["element"] ] self.register_buffer( - "_Z", torch.tensor(z_values, dtype=torch.int32, device=self.device) # dtype-ok: atomic-number Z categorical codes buffer; fixed int32 lookup keys + "_Z", torch.tensor(z_values, dtype=get_int_dtype(), device=self.device) ) return self._Z @@ -946,7 +946,7 @@ def _create_occupancy_groups(self, pdb_df, initial_occ): altloc_groups = [] refinable_mask = torch.zeros(n_atoms, dtype=torch.bool) - sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) # dtype-ok: arange atom indices (sharing groups); index requires long + sharing_groups_tensor = torch.arange(n_atoms, dtype=get_int_dtype()) collapsed_idx = 0 # First pass: altlocs. ALL atoms of one conformation must share a collapsed @@ -1015,7 +1015,7 @@ def _create_occupancy_groups(self, pdb_df, initial_occ): # Compact to contiguous indices 0..n_collapsed-1. unique_indices = torch.unique(sharing_groups_tensor, sorted=True) - index_map = torch.zeros(n_atoms, dtype=torch.long) # dtype-ok: index_map atom-index remap; indexing requires long + index_map = torch.zeros(n_atoms, dtype=get_int_dtype()) for new_idx, old_idx in enumerate(unique_indices): mask = sharing_groups_tensor == old_idx sharing_groups_tensor[mask] = new_idx @@ -2029,7 +2029,7 @@ def register_alternative_conformations(self): for altloc in unique_altlocs: altloc_atoms = group[group["altloc"] == altloc] indices = torch.tensor( - altloc_atoms["index"].tolist(), dtype=torch.long # dtype-ok: altloc atom indices; indexing requires long + altloc_atoms["index"].tolist(), dtype=get_int_dtype() ) conformation_tensors.append(indices) @@ -2070,7 +2070,6 @@ def shake_adp(self, stddev: float): new_adp, refinable_mask=self.adp.refinable_mask, name="adp" ) - def _new_model_from_df(self, df, *, strip_H=None, add_hydrogens=False): """Build a fresh model of the same class from a DataFrame. @@ -2221,7 +2220,6 @@ def hydrogenate(self, verbose: int = 0, optimize: bool = True) -> "Model": augmented = augment_atom_table(self.pdb, plan, restraints.topology) return self._new_model_from_df(augmented, strip_H=False) - def state_dict(self, destination=None, prefix="", keep_vars=False): """ Return a dictionary containing the complete state of the Model. @@ -3000,11 +2998,15 @@ def _complete_riding_waters(self, frames): self.pdb, plan, restraints.topology ) frames = generated.remap(old_rows).fill_planned_rows(new_rows) - source = torch.empty(len(augmented), dtype=torch.long, device=self.device) + source = torch.empty(len(augmented), dtype=get_int_dtype(), device=self.device) old_index = torch.as_tensor(old_rows, device=self.device) new_index = torch.as_tensor(new_rows, device=self.device) - source[old_index] = torch.arange(len(self.pdb), device=self.device) - source[new_index] = torch.as_tensor(plan.parent, device=self.device) + source[old_index] = torch.arange( + len(self.pdb), device=self.device, dtype=source.dtype + ) + source[new_index] = torch.as_tensor( + plan.parent, device=self.device, dtype=source.dtype + ) xyz = ( self.xyz.to_mixed_tensor() if hasattr(self.xyz, "to_mixed_tensor") diff --git a/torchref/model/parameter_wrappers.py b/torchref/model/parameter_wrappers.py index 1991cab..288919c 100644 --- a/torchref/model/parameter_wrappers.py +++ b/torchref/model/parameter_wrappers.py @@ -14,7 +14,7 @@ import torch from torch import nn -from torchref.config import get_float_dtype, normalize_device +from torchref.config import get_float_dtype, get_int_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -1474,11 +1474,13 @@ def _setup_sharing_groups_and_expansion( # Use sharing_groups directly as the expansion mask if sharing_groups is None: # No sharing - each atom maps to its own index - expansion_mask = torch.arange(n_atoms, dtype=torch.long, device=device) # dtype-ok: arange expansion_mask atom indices; index requires long + # dtype-ok: expansion_mask is a scatter_add_ index; int64 required on torch < 2.8 + expansion_mask = torch.arange(n_atoms, dtype=torch.long, device=device) self._collapsed_shape = n_atoms else: # Use the provided index tensor - expansion_mask = sharing_groups.to(device=device, dtype=torch.long) # dtype-ok: expansion_mask atom/group indices for scatter; requires long + # dtype-ok: expansion_mask is a scatter_add_ index; int64 required on torch < 2.8 + expansion_mask = sharing_groups.to(device=device, dtype=torch.long) self._collapsed_shape = expansion_mask.max().item() + 1 self.register_buffer("expansion_mask", expansion_mask) @@ -1500,10 +1502,10 @@ def _setup_sharing_groups_and_expansion( for conf_atoms in conf_groups: if isinstance(conf_atoms, (list, tuple)): conf_atoms = torch.tensor( - conf_atoms, dtype=torch.long, device=device # dtype-ok: conf_atoms atom indices; indexing requires long + conf_atoms, dtype=get_int_dtype(), device=device ) else: - conf_atoms = conf_atoms.to(device=device, dtype=torch.long) # dtype-ok: conf_atoms atom indices cast; indexing requires long + conf_atoms = conf_atoms.to(device=device, dtype=get_int_dtype()) # Get collapsed index for first atom collapsed_idx = expansion_mask[conf_atoms[0]].item() @@ -1531,7 +1533,7 @@ def _setup_sharing_groups_and_expansion( # Store as dictionary with keys like 'linked_occ_2', 'linked_occ_3', etc. for n_conf, groups in linked_occupancies.items(): # Shape: (N_groups, n_conf) - tensor = torch.tensor(groups, dtype=torch.long, device=device) # dtype-ok: linked-occupancy group index buffer; indexing requires long + tensor = torch.tensor(groups, dtype=get_int_dtype(), device=device) self.register_buffer(f"linked_occ_{n_conf}", tensor) # Store which sizes we have @@ -1539,7 +1541,9 @@ def _setup_sharing_groups_and_expansion( # Create count buffer for vectorized collapse operations # counts[i] = number of atoms that map to collapsed index i - counts = torch.zeros(self._collapsed_shape, dtype=torch.long, device=device) # dtype-ok: count accumulator; scatter_add source is long ones, dtype must match + counts = torch.zeros( + self._collapsed_shape, dtype=expansion_mask.dtype, device=device + ) counts.scatter_add_(0, expansion_mask, torch.ones_like(expansion_mask)) self.register_buffer("collapse_counts", counts) @@ -2017,7 +2021,7 @@ def from_residue_groups( grouped = pdb_dataframe.groupby(["resname", "resseq", "chainid", "altloc"]) n_atoms = len(initial_values) - sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) # dtype-ok: arange atom indices (sharing groups); index requires long + sharing_groups_tensor = torch.arange(n_atoms, dtype=get_int_dtype()) # Singletons keep their arange ids (0..n_atoms-1); start multi-atom # group ids past that range so a group id can never collide with a # singleton's leftover arange id (the torch.unique compaction below diff --git a/torchref/model/riding_xyz.py b/torchref/model/riding_xyz.py index 12771a6..d2133c6 100644 --- a/torchref/model/riding_xyz.py +++ b/torchref/model/riding_xyz.py @@ -30,6 +30,7 @@ place_local_frame, rotate_vectors, ) +from torchref.config import get_int_dtype from torchref.model.parameter_wrappers import MixedTensor from torchref.topology.hydrogens import HydrogenFrames @@ -65,9 +66,7 @@ def _register_rows(self, n_full: int, frames: HydrogenFrames, device) -> None: if len(rows) and is_riding[rows[rows >= 0]].any(): raise ValueError(f"{name} must reference stored rows, not riding ones") - long = dict( - dtype=torch.int64, device=device - ) # dtype-ok: row index buffers; int64 index required + long = dict(dtype=get_int_dtype(), device=device) self.register_buffer("base_row", torch.as_tensor(base, **long)) self.register_buffer("h_row", torch.as_tensor(h, **long)) self.register_buffer( @@ -96,25 +95,21 @@ def _rebuild_row_cache(self) -> None: n_full = int(base.numel() + self.h_row.numel()) self._n_full = n_full full_to_base = torch.full( - (max(n_full, 1),), -1, dtype=torch.int64, device=device - ) # dtype-ok: index map; int64 + (max(n_full, 1),), -1, dtype=get_int_dtype(), device=device + ) full_to_base[base] = torch.arange( - base.numel(), dtype=torch.int64, device=device - ) # dtype-ok: index map; int64 + base.numel(), dtype=get_int_dtype(), device=device + ) self._parent_bidx = full_to_base[self.parent_row.clamp(min=0)].clamp(min=0) self._n1_bidx = full_to_base[self.n1_row.clamp(min=0)].clamp(min=0) self._n2_bidx = full_to_base[self.n2_row.clamp(min=0)].clamp(min=0) # ``cat([base, derived])[gather]`` lays the full table out in one gather. - order = torch.empty( - n_full, dtype=torch.int64, device=device - ) # dtype-ok: gather index; int64 - order[base] = torch.arange( - base.numel(), dtype=torch.int64, device=device - ) # dtype-ok: gather index; int64 + order = torch.empty(n_full, dtype=get_int_dtype(), device=device) + order[base] = torch.arange(base.numel(), dtype=get_int_dtype(), device=device) order[self.h_row] = base.numel() + torch.arange( self.h_row.numel(), - dtype=torch.int64, - device=device, # dtype-ok: gather index; int64 + dtype=get_int_dtype(), + device=device, ) self._gather_order = order @@ -225,9 +220,7 @@ def __init__( for buffer in ("base_row", "h_row", "parent_row", "n1_row", "n2_row"): self.register_buffer( buffer, - torch.zeros( - 0, dtype=torch.int64, device=self.device - ), # dtype-ok: empty row-index buffer; int64 + torch.zeros(0, dtype=get_int_dtype(), device=self.device), ) self.register_buffer( "frame_valid", torch.zeros(0, dtype=torch.bool, device=self.device) @@ -261,8 +254,8 @@ def __init__( is_riding = np.zeros(n_full, dtype=bool) is_riding[np.asarray(frames.h_row, dtype=np.int64)] = True base_rows = torch.as_tensor( - np.nonzero(~is_riding)[0], dtype=torch.int64, device=device - ) # dtype-ok: row index; int64 + np.nonzero(~is_riding)[0], dtype=get_int_dtype(), device=device + ) if refinable_mask is None: base_mask = None @@ -380,8 +373,10 @@ def _orientation_selection(self, full_mask): parents = getattr(self, "_" + kind + "_parents") rows = getattr(self, "_" + kind + "_h") groups = getattr(self, "_" + kind + "_inverse") - selected = full_mask[parents].to(torch.int32) - selected.index_add_(0, groups, full_mask[self.h_row[rows]].to(torch.int32)) + selected = full_mask[parents].to(get_int_dtype()) + selected.index_add_( + 0, groups, full_mask[self.h_row[rows]].to(get_int_dtype()) + ) selections.append(selected > 0) return selections diff --git a/torchref/model/rigid_xyz.py b/torchref/model/rigid_xyz.py index 0d0a304..a22fee3 100644 --- a/torchref/model/rigid_xyz.py +++ b/torchref/model/rigid_xyz.py @@ -27,7 +27,7 @@ from torch import nn from torchref.base.alignment.rotation import rotation_matrix_euler_xyz -from torchref.config import get_float_dtype, normalize_device +from torchref.config import get_float_dtype, get_int_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -73,7 +73,7 @@ def __init__( dtype = dtype if dtype is not None else get_float_dtype() self.register_buffer("original_xyz", torch.empty(0, 3, device=device, dtype=dtype)) self.register_buffer( - "chain_indices", torch.empty(0, dtype=torch.long, device=device) # dtype-ok: empty chain_indices buffer; indexing requires long + "chain_indices", torch.empty(0, dtype=get_int_dtype(), device=device) ) self.register_buffer("chain_centers", torch.empty(0, 3, device=device, dtype=dtype)) self.register_buffer( diff --git a/torchref/refinement/base_refinement.py b/torchref/refinement/base_refinement.py index 96ee579..0d9dcac 100644 --- a/torchref/refinement/base_refinement.py +++ b/torchref/refinement/base_refinement.py @@ -8,7 +8,7 @@ import torch from torch.nn import Module as nnModule -from torchref.config import normalize_device +from torchref.config import get_int_dtype, normalize_device from torchref.io import ReflectionData from torchref.model.model_ft import ModelFT from torchref.refinement.logger import Logger @@ -441,7 +441,7 @@ def mark(idx): return # 4. freeze xyz of those atoms (same path as freeze_selection) - model.xyz_mask[torch.tensor(freeze_idx, dtype=torch.long)] = False # dtype-ok: freeze index used to index xyz_mask; PyTorch requires int64 + model.xyz_mask[torch.tensor(freeze_idx, dtype=get_int_dtype())] = False model.apply_mask_to_parameter("xyz") if self.verbose > 0: shown = frozen_res[:20] + (["..."] if len(frozen_res) > 20 else []) diff --git a/torchref/refinement/model_error_estimation/_shells.py b/torchref/refinement/model_error_estimation/_shells.py index 3d6e9a7..297507c 100644 --- a/torchref/refinement/model_error_estimation/_shells.py +++ b/torchref/refinement/model_error_estimation/_shells.py @@ -14,6 +14,8 @@ import torch +from torchref.config import get_int_dtype + @lru_cache(maxsize=8) def segment_layout(lengths: tuple[int, ...], device_str: str): @@ -23,12 +25,10 @@ def segment_layout(lengths: tuple[int, ...], device_str: str): ``lengths`` is a tuple so it can be a cache key. """ device = torch.device(device_str) - # dtype-ok: segment lengths for cumsum offsets/gather index; PyTorch requires int64 - L = torch.tensor(lengths, dtype=torch.long, device=device) + L = torch.tensor(lengths, dtype=get_int_dtype(), device=device) total = int(L.sum()) max_len = int(L.max()) if L.numel() else 0 - # dtype-ok: zero offset concatenated into gather index; PyTorch requires int64 - zero = torch.zeros(1, dtype=torch.long, device=device) + zero = torch.zeros(1, dtype=get_int_dtype(), device=device) starts = torch.cat([zero, L.cumsum(0)[:-1]]) ar = torch.arange(max_len, device=device).reshape(1, max_len) # Clamp keeps the gather in bounds for the padding slots; `mask` zeroes them anyway. diff --git a/torchref/refinement/optimizers/curvature.py b/torchref/refinement/optimizers/curvature.py index 74b5bba..84d28fe 100644 --- a/torchref/refinement/optimizers/curvature.py +++ b/torchref/refinement/optimizers/curvature.py @@ -23,6 +23,7 @@ import torch +from torchref.config import get_int_dtype from torchref.utils import use_portable @@ -36,7 +37,7 @@ def _sample_probe( """Draw one Hutchinson probe vector of length ``numel``.""" if probe == "rademacher": r = torch.randint( - 0, 2, (numel,), generator=generator, device=device, dtype=torch.int64 # dtype-ok: randint {0,1} bernoulli draw, immediately cast to float dtype; width irrelevant + 0, 2, (numel,), generator=generator, device=device, dtype=get_int_dtype() ) return r.to(dtype).mul_(2.0).sub_(1.0) # {0,1} -> {-1,+1} if probe == "gaussian": diff --git a/torchref/refinement/targets/adp/rigid_bond.py b/torchref/refinement/targets/adp/rigid_bond.py index 07a639a..1a58df0 100644 --- a/torchref/refinement/targets/adp/rigid_bond.py +++ b/torchref/refinement/targets/adp/rigid_bond.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Dict from torchref.base.targets.adp import adp_rigid_bond_aniso_math +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -122,7 +123,7 @@ def _bond_pairs(self) -> torch.Tensor: chunks.append(idx_) if chunks: return torch.cat(chunks, dim=0).contiguous() - return torch.empty(0, 2, dtype=torch.long, device=self.model.xyz().device) # dtype-ok: empty (0,2) atom-pair index tensor; PyTorch requires int64 + return torch.empty(0, 2, dtype=get_int_dtype(), device=self.model.xyz().device) def _compute_aniso_rigid_bond(self) -> torch.Tensor: """Rigid-bond NLL from ``Δz = l^T U_1 l - l^T U_2 l`` along each bond. diff --git a/torchref/refinement/targets/adp/similarity.py b/torchref/refinement/targets/adp/similarity.py index 191d237..bea1fac 100644 --- a/torchref/refinement/targets/adp/similarity.py +++ b/torchref/refinement/targets/adp/similarity.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Dict from torchref.base.targets.adp import adp_simu_math, adp_simu_aniso_math +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -94,8 +95,9 @@ def _get_pair_indices(self) -> torch.Tensor: if chunks: cached = torch.cat(chunks, dim=0).contiguous() else: - cached = torch.empty(0, 2, dtype=torch.long, # dtype-ok: empty (0,2) atom-pair index tensor; PyTorch requires int64 - device=self.model.xyz().device) + cached = torch.empty( + 0, 2, dtype=get_int_dtype(), device=self.model.xyz().device + ) self._simu_pair_indices_cache = cached return cached diff --git a/torchref/refinement/targets/difference.py b/torchref/refinement/targets/difference.py index d0a75ab..21ef310 100644 --- a/torchref/refinement/targets/difference.py +++ b/torchref/refinement/targets/difference.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple from .base import Target +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -204,10 +205,10 @@ def _match_reflections(self): device = hkl_light.device self._matched_indices_light = torch.tensor( - matched_light, dtype=torch.long, device=device # dtype-ok: matched atom indices used for indexing; PyTorch requires int64 + matched_light, dtype=get_int_dtype(), device=device ) self._matched_indices_dark = torch.tensor( - matched_dark, dtype=torch.long, device=device # dtype-ok: matched atom indices used for indexing; PyTorch requires int64 + matched_dark, dtype=get_int_dtype(), device=device ) # Store common HKL (using light indices, they should be identical) diff --git a/torchref/refinement/targets/geometry/chiral.py b/torchref/refinement/targets/geometry/chiral.py index 1cc6dce..d6e6925 100644 --- a/torchref/refinement/targets/geometry/chiral.py +++ b/torchref/refinement/targets/geometry/chiral.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Dict from torchref.base.targets.chiral import chiral_math +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -82,9 +83,9 @@ def get_violations(self, threshold: float = 0.5) -> Dict[str, torch.Tensor]: if "chiral" not in self.restraints.restraints: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing - 0, 4 - ), + "indices": torch.tensor( + [], dtype=get_int_dtype(), device=device + ).reshape(0, 4), "volumes": torch.tensor([], device=device), "ideal_volumes": torch.tensor([], device=device), "deviations": torch.tensor([], device=device), diff --git a/torchref/refinement/targets/geometry/non_bonded.py b/torchref/refinement/targets/geometry/non_bonded.py index 637be37..bc7a3f0 100644 --- a/torchref/refinement/targets/geometry/non_bonded.py +++ b/torchref/refinement/targets/geometry/non_bonded.py @@ -9,6 +9,7 @@ import torch from typing import TYPE_CHECKING, Dict, Tuple +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -346,9 +347,9 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: if "vdw" not in self.restraints.restraints: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing - 0, 2 - ), + "indices": torch.tensor( + [], dtype=get_int_dtype(), device=device + ).reshape(0, 2), "violations": torch.tensor([], device=device), "distances": torch.tensor([], device=device), "min_distances": torch.tensor([], device=device), @@ -359,9 +360,9 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: if indices is None or len(indices) == 0: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing - 0, 2 - ), + "indices": torch.tensor( + [], dtype=get_int_dtype(), device=device + ).reshape(0, 2), "violations": torch.tensor([], device=device), "distances": torch.tensor([], device=device), "min_distances": torch.tensor([], device=device), diff --git a/torchref/refinement/targets/similarity.py b/torchref/refinement/targets/similarity.py index d01c3ef..700883c 100644 --- a/torchref/refinement/targets/similarity.py +++ b/torchref/refinement/targets/similarity.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Dict from .base import Target +from torchref.config import get_int_dtype from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -68,10 +69,10 @@ def __init__( # path (the one ``load_state_dict`` uses) would have no such buffers at all. # ``_build_atom_map`` overwrites them rather than creating them. self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 + "_idx_dark", torch.zeros(0, dtype=get_int_dtype(), device=self.device) ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 + "_idx_light", torch.zeros(0, dtype=get_int_dtype(), device=self.device) ) if model_dark is not None and model_light is not None: self._build_atom_map() @@ -140,10 +141,10 @@ def _build_atom_map(self): "dark and light models" ) self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 + "_idx_dark", torch.zeros(0, dtype=get_int_dtype(), device=self.device) ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 + "_idx_light", torch.zeros(0, dtype=get_int_dtype(), device=self.device) ) return @@ -166,13 +167,13 @@ def _build_atom_map(self): self.register_buffer( "_idx_dark", torch.tensor( - merged["_idx_dark"].values, dtype=torch.long, device=self.device # dtype-ok: atom index tensor used for indexing; PyTorch requires int64 + merged["_idx_dark"].values, dtype=get_int_dtype(), device=self.device ), ) self.register_buffer( "_idx_light", torch.tensor( - merged["_idx_light"].values, dtype=torch.long, device=self.device # dtype-ok: atom index tensor used for indexing; PyTorch requires int64 + merged["_idx_light"].values, dtype=get_int_dtype(), device=self.device ), ) diff --git a/torchref/scaling/collection_scaler.py b/torchref/scaling/collection_scaler.py index 2b7d918..a1d99f7 100644 --- a/torchref/scaling/collection_scaler.py +++ b/torchref/scaling/collection_scaler.py @@ -14,7 +14,7 @@ import torch.nn as nn from torchref.base.metrics.rfactor import rfactor_work_free -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype from torchref.scaling.scaler_base import ( DEFAULT_SCALE_TARGET, SCALE_TARGETS, @@ -192,7 +192,8 @@ def _calc_initial_scale_joint(self): pos_mask = torch.ones_like(fobs, dtype=torch.bool) mask = (work_mask & pos_mask).to(torch.bool) - bins = self.bins[mask].to(torch.int64) # dtype-ok: bin indices for scatter/index_select; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + bins = self.bins[mask].to(torch.int64) log_ratios = ( torch.log(fobs_clamped[mask]) - torch.log(fcalc_amp[mask]) ).to(self.device) @@ -205,7 +206,7 @@ def _calc_initial_scale_joint(self): per_bin = scales / (counts + 1e-6) with torch.no_grad(): - target = per_bin.detach()[self.bins.to(torch.int64)] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 + target = per_bin.detach()[self.bins.to(get_int_dtype())] design = self._iso_design.to(target.dtype) coeff = torch.linalg.lstsq(design, target.unsqueeze(1)).solution.squeeze(1) self.c_iso = nn.Parameter(coeff.detach()) diff --git a/torchref/scaling/scaler_base.py b/torchref/scaling/scaler_base.py index b8350e4..786e67e 100644 --- a/torchref/scaling/scaler_base.py +++ b/torchref/scaling/scaler_base.py @@ -21,7 +21,7 @@ rfactor_work_free, ) from torchref.base.reciprocal import get_scattering_vectors -from torchref.config import get_complex_dtype, get_float_dtype +from torchref.config import get_complex_dtype, get_float_dtype, get_int_dtype from torchref.utils.autograd_ops import gather_with_index_add from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin @@ -264,7 +264,9 @@ def calc_initial_scale(self, fcalc: torch.Tensor): initial_log_scale.detach().cpu().numpy(), ) with torch.no_grad(): - target = initial_log_scale.detach().to(self.device)[self.bins.to(torch.int64)] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 + target = initial_log_scale.detach().to(self.device)[ + self.bins.to(get_int_dtype()) + ] design = self._iso_design.to(target.dtype) coeff = torch.linalg.lstsq(design, target.unsqueeze(1)).solution.squeeze(1) self.c_iso = nn.Parameter(coeff.detach().to(self.device)) @@ -422,7 +424,8 @@ def get_binwise_mean_intensity(self, fcalc: torch.Tensor): mean_calc_intensity = torch.zeros(self.nbins, device=self.device, dtype=fobs.dtype) counts = torch.zeros(self.nbins, device=self.device, dtype=fobs.dtype) counts_vals = torch.ones_like(F_calc, device=self.device, dtype=fobs.dtype) - bins_sel = self.bins.to(torch.int64)[sel] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 + # dtype-ok: scatter_add index; int64 required on torch < 2.8 + bins_sel = self.bins.to(torch.int64)[sel] mean_obs_intensity = torch.scatter_add( mean_obs_intensity, 0, bins_sel, intensities[sel] ) diff --git a/torchref/scaling/solvent.py b/torchref/scaling/solvent.py index 87f906e..c7a4e0d 100644 --- a/torchref/scaling/solvent.py +++ b/torchref/scaling/solvent.py @@ -10,7 +10,7 @@ get_scattering_vectors, ifft, ) -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin from torchref.utils.device_resolution import resolve_device @@ -387,7 +387,7 @@ def get_solvent_mask(self): # grids, where the SF code's 1024 would OOM (denser intermediates). ATOM_CHUNK = 256 - grid_dims = torch.tensor(grid_shape, dtype=torch.long, device=device) # dtype-ok: grid dims for voxel index arithmetic; PyTorch requires int64 + grid_dims = torch.tensor(grid_shape, dtype=get_int_dtype(), device=device) grid_shape_float = grid_dims.float() inv_grid = 1.0 / grid_shape_float G = frac.T @ frac # metric tensor: r²_cart = diff_frac · G · diff_frac @@ -464,12 +464,12 @@ def get_solvent_mask(self): protein_voxels = ( torch.cat(protein_chunks, dim=0) if protein_chunks - else torch.empty((0, 3), dtype=torch.long, device=device) # dtype-ok: empty (0,3) voxel index tensor; PyTorch requires int64 for indexing + else torch.empty((0, 3), dtype=get_int_dtype(), device=device) ) boundary_voxels = ( torch.cat(boundary_chunks, dim=0) if boundary_chunks - else torch.empty((0, 3), dtype=torch.long, device=device) # dtype-ok: empty (0,3) voxel index tensor; PyTorch requires int64 for indexing + else torch.empty((0, 3), dtype=get_int_dtype(), device=device) ) del protein_chunks, boundary_chunks diff --git a/torchref/scaling/wilson.py b/torchref/scaling/wilson.py index e16b048..e476201 100644 --- a/torchref/scaling/wilson.py +++ b/torchref/scaling/wilson.py @@ -32,7 +32,7 @@ import torch -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype from torchref.scaling.basis import chebyshev_design __all__ = ["WilsonNormaliser"] @@ -267,7 +267,6 @@ def _solve_intercept( out[0] = out[0] + torch.log(ratio) return out - def _irls( self, X: torch.Tensor, @@ -455,7 +454,7 @@ def from_hkl( branches feed two different parameters of the same likelihood. """ work = get_float_dtype() - hkl_l = hkl.to(torch.long) # dtype-ok: Miller indices are integers + hkl_l = hkl.to(get_int_dtype()) # The cell may carry the configured default device while the reflections # are somewhere else; the caller should not have to reconcile them. rec = cell.reciprocal_basis_matrix.to(device=hkl_l.device, dtype=work) diff --git a/torchref/symmetry/map_symmetry.py b/torchref/symmetry/map_symmetry.py index afbeb4a..b58b65b 100644 --- a/torchref/symmetry/map_symmetry.py +++ b/torchref/symmetry/map_symmetry.py @@ -21,6 +21,7 @@ import torch +from torchref.config import get_int_dtype from torchref.utils.device_mixin import DeviceMixin @@ -149,7 +150,7 @@ def _index_grid(self, op_index: int) -> torch.Tensor: transformed = transformed - torch.floor(transformed) shape_t = torch.tensor([nx, ny, nz], dtype=dtype, device=device) - indices = torch.round(transformed * shape_t).to(torch.int64) # dtype-ok: rounded voxel grid indices; int64 index tensor required + indices = torch.round(transformed * shape_t).to(get_int_dtype()) indices[:, 0] %= nx indices[:, 1] %= ny indices[:, 2] %= nz diff --git a/torchref/symmetry/reciprocal_symmetry.py b/torchref/symmetry/reciprocal_symmetry.py index f595ee7..1954fae 100644 --- a/torchref/symmetry/reciprocal_symmetry.py +++ b/torchref/symmetry/reciprocal_symmetry.py @@ -24,8 +24,7 @@ import numpy as np import torch -from torchref.config import get_float_dtype - +from torchref.config import get_float_dtype, get_int_dtype def _expand_hkl( @@ -82,7 +81,7 @@ def _expand_hkl( for i in range(n_ops): # h' = h @ R^T hkl_transformed = torch.round(torch.matmul(hkl_float, recip_matrices[i].T)).to( - torch.int32 # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation + get_int_dtype() ) # Phase shift from translation: -2π h·t, for h' = hR under the convention # F(h) = Σ_j f_j exp(+2πi h·x_j). Do NOT "simplify" the sign: the wrong sign @@ -124,10 +123,10 @@ def _expand_hkl( # Build output tensors expanded_hkl = torch.tensor( - [list(k) for k in unique_dict.keys()], dtype=torch.int32, device=device # dtype-ok: unique Miller indices (hkl); fixed-width int32 representation + [list(k) for k in unique_dict.keys()], dtype=get_int_dtype(), device=device ) phase_shifts = torch.tensor(unique_phases, dtype=get_float_dtype(), device=device) - orig_idx_tensor = torch.tensor(orig_indices, dtype=torch.int64, device=device) # dtype-ok: reflection index mapping; int64 index tensor required + orig_idx_tensor = torch.tensor(orig_indices, dtype=get_int_dtype(), device=device) if remove_absences and sym.number != 1: keep_mask = ~sym.is_absent(expanded_hkl) @@ -169,7 +168,7 @@ def _complete_hkl( ------- complete_hkl : torch.Tensor, shape (M, 3), dtype int32 All possible Miller indices within resolution (minus systematic absences). - input_indices : torch.Tensor, shape (M,), dtype int64 + input_indices : torch.Tensor, shape (M,), integer dtype Index mapping complete → input, or -1 where missing. Use as ``F_complete[~missing] = F_input[input_indices[~missing]]``. missing_mask : torch.Tensor, shape (M,), dtype bool @@ -198,7 +197,7 @@ def _complete_hkl( all_hkl_np = all_hkl.cpu().numpy() n_complete = len(all_hkl) - input_indices = torch.full((n_complete,), -1, dtype=torch.int64, device=device) # dtype-ok: reflection index buffer (-1 sentinel); int64 index required + input_indices = torch.full((n_complete,), -1, dtype=get_int_dtype(), device=device) missing_mask = torch.ones(n_complete, dtype=torch.bool, device=device) for i, hkl in enumerate(all_hkl_np): @@ -274,7 +273,7 @@ def get_canonical_hkl(hkl_single): for i in range(n_ops): # h' = h @ R^T hkl_trans = torch.round(torch.matmul(hkl_single, recip_matrices[i].T)).to( - torch.int32 # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation + get_int_dtype() ) equivalents.append(hkl_trans) @@ -311,7 +310,7 @@ def get_canonical_hkl(hkl_single): R = recip_matrices[equiv_idx] t = translations[equiv_idx] - hkl_trans = torch.round(torch.matmul(hkl_single, R.T)).to(torch.int32) # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation + hkl_trans = torch.round(torch.matmul(hkl_single, R.T)).to(get_int_dtype()) # -2π h·t, same convention as expand_hkl (see the derivation there). phase_shift = -2.0 * np.pi * torch.matmul(hkl_single, t) @@ -333,9 +332,9 @@ def get_canonical_hkl(hkl_single): asu_list = sorted(asu_reflections.keys()) n_asu = len(asu_list) - hkl_asu = torch.tensor(asu_list, dtype=torch.int32, device=device) # dtype-ok: ASU Miller indices (hkl); fixed-width int32 representation + hkl_asu = torch.tensor(asu_list, dtype=get_int_dtype(), device=device) reduction_indices = torch.full( - (n_asu, n_equiv), -1, dtype=torch.int64, device=device # dtype-ok: reduction index map (-1 sentinel); int64 index tensor required + (n_asu, n_equiv), -1, dtype=get_int_dtype(), device=device ) phase_shifts = torch.zeros((n_asu, n_equiv), dtype=get_float_dtype(), device=device) @@ -446,7 +445,7 @@ def _canonicalize_hkl( empty_hkl = torch.empty((0, 3), dtype=hkl_dtype, device=device) empty_f = torch.empty(0, dtype=get_float_dtype(), device=device) empty_b = torch.empty(0, dtype=torch.bool, device=device) - empty_i = torch.empty(0, dtype=torch.int64, device=device) # dtype-ok: empty index tensor; int64 index dtype required + empty_i = torch.empty(0, dtype=get_int_dtype(), device=device) return empty_hkl, empty_f, empty_b, empty_i # The ASU lookup tables are numpy-backed, so the operations come across to CPU @@ -549,11 +548,9 @@ def _canonicalize_hkl( # Lexicographic sort by (h, k, l) via composite key h_max = int(canonical_hkl.abs().max().item()) + 1 base = 2 * h_max + 1 - sort_key = ( - canonical_hkl[:, 0].to(torch.int64) * base * base # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing - + canonical_hkl[:, 1].to(torch.int64) * base # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing - + canonical_hkl[:, 2].to(torch.int64) # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing - ) + # dtype-ok: composite sort key h*base^2+k*base+l overflows int32 for large Miller indices + hkl64 = canonical_hkl.to(torch.int64) + sort_key = hkl64[:, 0] * base * base + hkl64[:, 1] * base + hkl64[:, 2] sort_indices = torch.argsort(sort_key) return ( diff --git a/torchref/symmetry/symmetry.py b/torchref/symmetry/symmetry.py index 4cb767d..428366c 100644 --- a/torchref/symmetry/symmetry.py +++ b/torchref/symmetry/symmetry.py @@ -29,7 +29,7 @@ import torch -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype from torchref.utils.device_mixin import DeviceMixin if TYPE_CHECKING: @@ -352,11 +352,11 @@ def expand_reciprocal(self, hkl: torch.Tensor) -> torch.Tensor: Returns ------- torch.Tensor - Shape ``(n_ops, N, 3)``, rounded to ``int64``. Rounding is exact for valid - operations on integer indices and only mops up float error. + Shape ``(n_ops, N, 3)``, rounded to the configured int dtype. Rounding is + exact for valid operations on integer indices and only mops up float error. """ equivalents = self.reciprocal.apply_rotations(hkl) - return torch.round(equivalents).to(torch.int64) # dtype-ok: rounded Miller equivalents; int64 for exact integer compare/index + return torch.round(equivalents).to(get_int_dtype()) # ========================================================================= # Reflection predicates @@ -381,7 +381,7 @@ def is_centric(self, hkl: torch.Tensor) -> torch.Tensor: with torch.no_grad(): flat = hkl.reshape(-1, 3) equivalents = self.expand_reciprocal(flat) # (n_ops, N, 3) - target = -flat.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + target = -flat.to(device=equivalents.device, dtype=get_int_dtype()) centric = (equivalents == target).all(dim=-1).any(dim=0) return centric.reshape(original_shape).to(hkl.device) @@ -405,7 +405,7 @@ def is_absent(self, hkl: torch.Tensor) -> torch.Tensor: with torch.no_grad(): flat = hkl.reshape(-1, 3) equivalents = self.expand_reciprocal(flat) # (n_ops, N, 3) - target = flat.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + target = flat.to(device=equivalents.device, dtype=get_int_dtype()) maps_to_self = (equivalents == target).all(dim=-1) # (n_ops, N) h_dot_t = torch.matmul( @@ -465,7 +465,7 @@ def epsilon(self, hkl: torch.Tensor, *, friedel: bool = True) -> torch.Tensor: float_dtype = get_float_dtype() with torch.no_grad(): equivalents = self.expand_reciprocal(hkl) # (n_ops, N, 3) - target = hkl.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + target = hkl.to(device=equivalents.device, dtype=get_int_dtype()) fixes = (equivalents == target).all(dim=-1) if friedel: fixes = fixes | (equivalents == -target).all(dim=-1) @@ -494,7 +494,7 @@ def grid_requirements(self) -> dict: # ``Fraction(float)`` would need a tolerance where this is exact. numerators = torch.round( self.translations.detach().cpu().double() * _TRANSLATION_DENOMINATOR - ).to(torch.int64) # dtype-ok: integer translation numerators for exact Fraction recovery + ).to(get_int_dtype()) for op_numerators in numerators.tolist(): for axis, numerator in enumerate(op_numerators): diff --git a/torchref/topology/atom_graph.py b/torchref/topology/atom_graph.py index 0520431..c012187 100644 --- a/torchref/topology/atom_graph.py +++ b/torchref/topology/atom_graph.py @@ -17,6 +17,7 @@ import numpy as np import torch +from torchref.config import get_int_dtype from torchref.topology.edges import EdgeBlock from torchref.utils.device_mixin import DeviceMixin @@ -42,8 +43,8 @@ def _build_csr(bonds: torch.Tensor, n_atoms: int) -> Tuple[torch.Tensor, torch.T device = bonds.device if bonds.numel() == 0: return ( - torch.zeros(n_atoms + 1, dtype=torch.int64, device=device), # dtype-ok: CSR indptr offset array; int64 index required - torch.zeros(0, dtype=torch.int64, device=device), # dtype-ok: empty CSR neighbor index array; int64 index required + torch.zeros(n_atoms + 1, dtype=get_int_dtype(), device=device), + torch.zeros(0, dtype=get_int_dtype(), device=device), ) src = torch.cat([bonds[:, 0], bonds[:, 1]]) @@ -55,9 +56,9 @@ def _build_csr(bonds: torch.Tensor, n_atoms: int) -> Tuple[torch.Tensor, torch.T src, dst = pairs[:, 0], pairs[:, 1] counts = torch.bincount(src, minlength=n_atoms) - indptr = torch.zeros(n_atoms + 1, dtype=torch.int64, device=device) # dtype-ok: CSR indptr offset array; int64 index required + indptr = torch.zeros(n_atoms + 1, dtype=get_int_dtype(), device=device) torch.cumsum(counts, dim=0, out=indptr[1:]) - return indptr, dst.to(torch.int64) # dtype-ok: CSR neighbor (dst) index array; int64 index required + return indptr, dst.to(get_int_dtype()) def _extend_paths( @@ -81,13 +82,17 @@ def _extend_paths( """ device = paths.device if paths.numel() == 0: - return torch.zeros((0, paths.shape[1] + 1), dtype=torch.int64, device=device) # dtype-ok: empty BFS path index array; int64 index required + return torch.zeros( + (0, paths.shape[1] + 1), dtype=get_int_dtype(), device=device + ) last, prev = paths[:, -1], paths[:, -2] counts = indptr[last + 1] - indptr[last] total = int(counts.sum()) if total == 0: - return torch.zeros((0, paths.shape[1] + 1), dtype=torch.int64, device=device) # dtype-ok: empty BFS path index array; int64 index required + return torch.zeros( + (0, paths.shape[1] + 1), dtype=get_int_dtype(), device=device + ) row = torch.repeat_interleave(torch.arange(len(paths), device=device), counts) # Offset of each slot within its own neighbour list. @@ -204,16 +209,18 @@ def implicit_h_count(self) -> Optional[torch.Tensor]: return None is_h = self.is_hydrogen bonds = self.bonds.indices - present = torch.zeros(self.n_atoms, dtype=torch.int64, device=bonds.device) # dtype-ok: bincount output; int64 + present = torch.zeros(self.n_atoms, dtype=get_int_dtype(), device=bonds.device) if bonds.numel(): heavy_of_h = torch.cat( [bonds[is_h[bonds[:, 1]] & ~is_h[bonds[:, 0]], 0], bonds[is_h[bonds[:, 0]] & ~is_h[bonds[:, 1]], 1]] ) if heavy_of_h.numel(): - present = torch.bincount(heavy_of_h, minlength=self.n_atoms) + present = torch.bincount(heavy_of_h, minlength=self.n_atoms).to( + present.dtype + ) known = self.template_h_count >= 0 - missing = self.template_h_count.to(torch.int64) - present + missing = self.template_h_count - present return torch.where(known, missing.clamp(min=0), torch.zeros_like(missing)) def subset(self, remap: torch.Tensor, residue_remap: torch.Tensor) -> "AtomGraph": diff --git a/torchref/topology/build.py b/torchref/topology/build.py index 8797029..901ff92 100644 --- a/torchref/topology/build.py +++ b/torchref/topology/build.py @@ -12,6 +12,7 @@ import pandas as pd import torch +from torchref.config import get_int_dtype from torchref.topology.builders import ( InterResidueAngleBuilder, InterResidueBondBuilder, @@ -711,7 +712,7 @@ def _block_with_values( per_origin, arity, edge_type, payload ) block = EdgeBlock( - indices=torch.as_tensor(indices, dtype=torch.int64, device=device), # dtype-ok: atom index tensor for restraint edges; int64 index required + indices=torch.as_tensor(indices, dtype=get_int_dtype(), device=device), origin_bounds=bounds, ) values = { @@ -997,7 +998,7 @@ def build_topology_with_values( np.arange(n_res, dtype=np.int64), nodes["atom_end"] - nodes["atom_start"], ), - dtype=torch.int64, # dtype-ok: atom index tensor; int64 index required + dtype=get_int_dtype(), device=device, ), bonds=bond_block, @@ -1007,7 +1008,7 @@ def build_topology_with_values( planes=plane_blocks, energy_type=energy_type, template_h_count=torch.as_tensor( - template_h_count, dtype=torch.int8, device=device + template_h_count, dtype=get_int_dtype(), device=device ), ) diff --git a/torchref/topology/builders.py b/torchref/topology/builders.py index 1450262..f554ba8 100644 --- a/torchref/topology/builders.py +++ b/torchref/topology/builders.py @@ -613,8 +613,10 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -720,8 +722,10 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -841,8 +845,10 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), "periods": torch.tensor(periods, dtype=get_int_dtype(), device=device), } @@ -932,7 +938,7 @@ def build( key = f"{n_atoms}_atoms" result[key] = { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -1054,7 +1060,7 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "ideal_volumes": torch.tensor( ideal_volumes, dtype=get_float_dtype(), device=device ), @@ -1309,8 +1315,10 @@ def finalize( sigmas = np.where(sigmas == 0, min_sigma, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -1412,8 +1420,10 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -1534,8 +1544,10 @@ def finalize( sigmas = np.where(sigmas == 0, min_sigma, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -1648,8 +1660,10 @@ def build( sigmas = np.where(sigmas == 0, 1e-4, sigmas) return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } @@ -1785,8 +1799,10 @@ def finalize_disulfide( periods = periods[sort_order] return { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 - "references": torch.tensor(references, dtype=get_float_dtype(), device=device), + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), + "references": torch.tensor( + references, dtype=get_float_dtype(), device=device + ), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), "periods": torch.tensor(periods, dtype=get_int_dtype(), device=device), } @@ -1952,7 +1968,7 @@ def build( indices = indices[order] periods = periods[order] result["phi"] = { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "periods": torch.tensor(periods, dtype=get_int_dtype(), device=device), } @@ -1965,7 +1981,7 @@ def build( indices = indices[order] periods = periods[order] result["psi"] = { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "periods": torch.tensor(periods, dtype=get_int_dtype(), device=device), } @@ -1984,7 +2000,7 @@ def build( periods = periods[order] is_proline = is_proline[order] result["omega"] = { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "references": torch.tensor( references, dtype=get_float_dtype(), device=device ), @@ -2023,13 +2039,13 @@ def build( stypes = stypes[order] result["ramachandran"] = { "phi_indices": torch.tensor( - phi_idx, dtype=torch.long, device=device # dtype-ok: phi atom-index tensor for dihedral; int64 required + phi_idx, dtype=get_int_dtype(), device=device ), "psi_indices": torch.tensor( - psi_idx, dtype=torch.long, device=device # dtype-ok: psi atom-index tensor for dihedral; int64 required + psi_idx, dtype=get_int_dtype(), device=device ), "surface_type": torch.tensor( - stypes, dtype=torch.long, device=device # dtype-ok: categorical rama surface-type code used as advanced index; int64 + stypes, dtype=get_int_dtype(), device=device ), } @@ -2077,34 +2093,34 @@ def build( planes_by_size: Dict[int, List[Tuple[np.ndarray, np.ndarray]]] = {} for res_i_idx, res_next_idx in pairs: - for map_i in conf_maps[res_i_idx]: - for map_next in conf_maps[res_next_idx]: + for map_i in conf_maps[res_i_idx]: + for map_next in conf_maps[res_next_idx]: - for plane_data in link_data.planes: - comp_ids = plane_data["comp_ids"] - atom_names = plane_data["atoms"] - sigmas = plane_data["sigmas"] + for plane_data in link_data.planes: + comp_ids = plane_data["comp_ids"] + atom_names = plane_data["atoms"] + sigmas = plane_data["sigmas"] - plane_indices = [] - plane_sigmas = [] - all_found = True + plane_indices = [] + plane_sigmas = [] + all_found = True - for i, (comp_id, atom_name, sigma) in enumerate( + for i, (comp_id, atom_name, sigma) in enumerate( zip(comp_ids, atom_names, sigmas) ): - atom_map = map_i if comp_id == "1" else map_next - if atom_name in atom_map: - plane_indices.append(atom_map[atom_name]) - plane_sigmas.append(sigma) - else: - all_found = False - break - - if all_found and len(plane_indices) >= 3: - n_atoms = len(plane_indices) - if n_atoms not in planes_by_size: - planes_by_size[n_atoms] = [] - planes_by_size[n_atoms].append( + atom_map = map_i if comp_id == "1" else map_next + if atom_name in atom_map: + plane_indices.append(atom_map[atom_name]) + plane_sigmas.append(sigma) + else: + all_found = False + break + + if all_found and len(plane_indices) >= 3: + n_atoms = len(plane_indices) + if n_atoms not in planes_by_size: + planes_by_size[n_atoms] = [] + planes_by_size[n_atoms].append( ( np.array(plane_indices, dtype=np.int64), np.array(plane_sigmas, dtype=np.float64), @@ -2128,7 +2144,7 @@ def build( key = f"{n_atoms}_atoms" result[key] = { - "indices": torch.tensor(indices, dtype=torch.long, device=device), # dtype-ok: atom-index restraint tensor; torch indexing requires int64 + "indices": torch.tensor(indices, dtype=get_int_dtype(), device=device), "sigmas": torch.tensor(sigmas, dtype=get_float_dtype(), device=device), } diff --git a/torchref/topology/edges.py b/torchref/topology/edges.py index c70a5f2..50b3877 100644 --- a/torchref/topology/edges.py +++ b/torchref/topology/edges.py @@ -17,6 +17,7 @@ import numpy as np import torch +from torchref.config import get_int_dtype from torchref.utils.device_mixin import DeviceMixin #: Origin order per edge type. Fixes the block layout so a rebuild on the same @@ -153,7 +154,7 @@ class EdgeBlock(DeviceMixin): def empty(cls, arity: int, device=None) -> "EdgeBlock": """An edge-free block of the given arity.""" return cls( - indices=torch.zeros((0, arity), dtype=torch.int64, device=device), # dtype-ok: empty edge index tensor (0,arity); int64 index required + indices=torch.zeros((0, arity), dtype=get_int_dtype(), device=device), origin_bounds={}, ) @@ -190,7 +191,7 @@ def from_origins( if len(indices) == 0: return cls.empty(arity, device=device) return cls( - indices=torch.as_tensor(indices, dtype=torch.int64, device=device), # dtype-ok: edge atom index tensor; int64 index required + indices=torch.as_tensor(indices, dtype=get_int_dtype(), device=device), origin_bounds=bounds, ) diff --git a/torchref/topology/hydrogens.py b/torchref/topology/hydrogens.py index 5ec8530..b4b41ed 100644 --- a/torchref/topology/hydrogens.py +++ b/torchref/topology/hydrogens.py @@ -28,7 +28,8 @@ import numpy as np import torch -from torchref.config import get_float_dtype + +from torchref.config import get_float_dtype, get_int_dtype #: Standard heavy-atom valences, one of the two budgets that cap how many hydrogens a #: parent may take. Elements not listed fall back to 4 and are then bounded only by the @@ -1145,18 +1146,16 @@ def sorted_by_row(self) -> "HydrogenFrames": def to_tensors(self, device=None) -> Dict[str, torch.Tensor]: """Return frame and orientation arrays as tensors, keyed by field name.""" return { - "h_row": torch.as_tensor( - self.h_row, dtype=torch.int64, device=device - ), # dtype-ok: row index; int64 required + "h_row": torch.as_tensor(self.h_row, dtype=get_int_dtype(), device=device), "parent_row": torch.as_tensor( - self.parent_row, dtype=torch.int64, device=device - ), # dtype-ok: row index; int64 required + self.parent_row, dtype=get_int_dtype(), device=device + ), "n1_row": torch.as_tensor( - self.n1_row, dtype=torch.int64, device=device - ), # dtype-ok: row index; int64 required + self.n1_row, dtype=get_int_dtype(), device=device + ), "n2_row": torch.as_tensor( - self.n2_row, dtype=torch.int64, device=device - ), # dtype-ok: row index; int64 required + self.n2_row, dtype=get_int_dtype(), device=device + ), "frame_valid": torch.as_tensor( self.frame_valid, dtype=torch.bool, device=device ), diff --git a/torchref/topology/nonbonded.py b/torchref/topology/nonbonded.py index 404e7a6..8128f60 100644 --- a/torchref/topology/nonbonded.py +++ b/torchref/topology/nonbonded.py @@ -15,7 +15,7 @@ import numpy as np import torch -from torchref.config import dtypes, get_float_dtype +from torchref.config import dtypes, get_float_dtype, get_int_dtype if TYPE_CHECKING: from torchref.symmetry.cell import Cell @@ -85,8 +85,8 @@ def prefilter_symop_offsets( valid_ops.append(op_idx) valid_offsets.append([dx, dy, dz]) - op_indices = torch.tensor(valid_ops, dtype=torch.long, device=device) # dtype-ok: symmetry-operator index tensor; int64 - cell_offsets = torch.tensor(valid_offsets, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata + op_indices = torch.tensor(valid_ops, dtype=get_int_dtype(), device=device) + cell_offsets = torch.tensor(valid_offsets, dtype=get_int_dtype(), device=device) return op_indices, cell_offsets @@ -146,7 +146,7 @@ def assign_to_grid( gd = grid_dims.to(device=device, dtype=fdtype) cell_ijk = (frac_wrapped * gd[None, None, :]).long() cell_ijk = cell_ijk.clamp( - min=torch.zeros(3, dtype=torch.long, device=device), # dtype-ok: clamp min-bound for long grid-index tensor; matches int64 + min=torch.zeros(3, dtype=get_int_dtype(), device=device), max=(grid_dims - 1).to(device), ) @@ -188,14 +188,12 @@ def build_cell_list( unique_cells, counts = torch.unique_consecutive( sorted_cells, return_counts=True ) - starts = torch.zeros(len(unique_cells) + 1, dtype=torch.long, device=device) # dtype-ok: CSR boundary/offset array; int64 required + starts = torch.zeros(len(unique_cells) + 1, dtype=get_int_dtype(), device=device) starts[1:] = counts.cumsum(0) - cell_lookup = torch.full( - (n_grid_total,), -1, dtype=torch.long, device=device # dtype-ok: grid-cell to index lookup table; used for indexing, int64 - ) + cell_lookup = torch.full((n_grid_total,), -1, dtype=get_int_dtype(), device=device) cell_lookup[unique_cells] = torch.arange( - len(unique_cells), dtype=torch.long, device=device # dtype-ok: index values written into lookup table; int64 + len(unique_cells), dtype=get_int_dtype(), device=device ) return sort_order, unique_cells, starts, cell_lookup @@ -233,7 +231,7 @@ def _get_canonical_offsets_14(device: torch.device) -> torch.Tensor: offsets.append([dx, dy, dz]) assert len(offsets) == 14, f"expected 14 canonical offsets, got {len(offsets)}" _NEIGHBOR_OFFSETS_14 = torch.tensor( - offsets, dtype=torch.long, device=device # dtype-ok: grid neighbor-cell offset deltas used to compute index; int64 + offsets, dtype=get_int_dtype(), device=device ) return _NEIGHBOR_OFFSETS_14 @@ -493,7 +491,7 @@ def find_pairs_periodic_grid_v2( all_pair_combo_j.append(cj) if not all_pair_atom_i: - empty = torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty atom-pair index placeholder; int64 required + empty = torch.tensor([], dtype=get_int_dtype(), device=device) return empty, empty, empty return ( @@ -517,11 +515,13 @@ def exclusion_set_to_hash( Hash: min(i,j) * max_idx + max(i,j), sorted for searchsorted. """ if not exclusion_set: - return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty exclusion-hash placeholder; int64 + # dtype-ok: packed pair key min*max_idx+max overflows int32 beyond ~46k atoms; searchsorted needs both sides int64 + return torch.tensor([], dtype=torch.long, device=device) arr = np.array(list(exclusion_set), dtype=np.int64) hashes = arr[:, 0] * max_idx + arr[:, 1] # already (min, max) hashes.sort() - return torch.tensor(hashes, dtype=torch.long, device=device) # dtype-ok: packed pair-hash key for searchsorted; int64 avoids overflow + # dtype-ok: packed pair key min*max_idx+max overflows int32 beyond ~46k atoms; searchsorted needs both sides int64 + return torch.tensor(hashes, dtype=torch.long, device=device) def filter_pairs( @@ -629,11 +629,11 @@ def build_vdw_restraints_gpu( sg = SG(sg) empty_result = { - "indices": torch.zeros(0, 2, dtype=torch.long, device=device), # dtype-ok: atom-pair index tensor; torch indexing requires int64 + "indices": torch.zeros(0, 2, dtype=get_int_dtype(), device=device), "min_distances": torch.zeros(0, dtype=get_float_dtype(), device=device), "sigmas": torch.zeros(0, dtype=get_float_dtype(), device=device), - "symop_indices": torch.zeros(0, dtype=torch.long, device=device), # dtype-ok: symmetry-operator index tensor; int64 - "cell_offsets": torch.zeros(0, 3, dtype=torch.long, device=device), # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata + "symop_indices": torch.zeros(0, dtype=get_int_dtype(), device=device), + "cell_offsets": torch.zeros(0, 3, dtype=get_int_dtype(), device=device), } # Step 1: prefilter symop combos @@ -654,12 +654,15 @@ def build_vdw_restraints_gpu( identity_indices = is_identity.nonzero(as_tuple=True)[0] if len(identity_indices) == 0: # Identity not in valid combos — should not happen, but add it - op_indices = torch.cat([ - torch.zeros(1, dtype=torch.long, device=device), op_indices # dtype-ok: identity prepended to symop-index tensor; int64 - ]) - cell_offsets_valid = torch.cat([ - torch.zeros(1, 3, dtype=torch.long, device=device), cell_offsets_valid # dtype-ok: identity prepended to cell-offset tensor; int64 - ]) + op_indices = torch.cat( + [torch.zeros(1, dtype=get_int_dtype(), device=device), op_indices] + ) + cell_offsets_valid = torch.cat( + [ + torch.zeros(1, 3, dtype=get_int_dtype(), device=device), + cell_offsets_valid, + ] + ) identity_combo = 0 M = len(op_indices) else: @@ -774,7 +777,9 @@ def build_vdw_restraints_gpu( "valid_op_indices": op_indices, "valid_cell_offsets": cell_offsets_valid, "grid_dims": grid_dims, - "identity_combo": torch.tensor(identity_combo, dtype=torch.long, device=device), # dtype-ok: combo index scalar into symop/offset arrays; int64 + "identity_combo": torch.tensor( + identity_combo, dtype=get_int_dtype(), device=device + ), } if verbose > 0: @@ -856,7 +861,7 @@ def find_h_vdw_pairs_gpu( xyz_all = torch.cat([xyz_heavy, xyz_h], dim=0) # (N_all, 3) n_all = xyz_all.shape[0] - empty = torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index placeholder tensor; int64 required + empty = torch.tensor([], dtype=get_int_dtype(), device=device) if n_all == 0: return empty, empty, empty diff --git a/torchref/topology/residue_graph.py b/torchref/topology/residue_graph.py index 0a68a21..0cc77c1 100644 --- a/torchref/topology/residue_graph.py +++ b/torchref/topology/residue_graph.py @@ -15,6 +15,8 @@ import numpy as np import torch +from torchref.config import get_int_dtype + #: SG-SG separation below which two cysteines are taken to be disulfide-bonded. DISULFIDE_MAX_DISTANCE = 2.5 @@ -272,7 +274,7 @@ def find_disulfide_links( rows = list(sg_rows) if len(rows) < 2: return [] - idx = torch.as_tensor(rows, dtype=torch.int64, device=xyz.device) # dtype-ok: residue-atom index tensor; int64 index required + idx = torch.as_tensor(rows, dtype=get_int_dtype(), device=xyz.device) dist = torch.cdist(xyz[idx], xyz[idx]) close = (dist > DISULFIDE_MIN_DISTANCE) & (dist < DISULFIDE_MAX_DISTANCE) diff --git a/torchref/topology/restraint_sets.py b/torchref/topology/restraint_sets.py index 45c079a..a3ae432 100644 --- a/torchref/topology/restraint_sets.py +++ b/torchref/topology/restraint_sets.py @@ -17,7 +17,7 @@ import numpy as np import torch -from torchref.config import get_float_dtype +from torchref.config import get_float_dtype, get_int_dtype #: Origins making up each edge type's ``all`` group -- what the geometry targets read. #: ``None`` means every origin present. ``phi`` and ``psi`` are conformationally free @@ -41,7 +41,7 @@ def to_tensor(values, prop: str, device=None) -> torch.Tensor: if isinstance(values, torch.Tensor): return values.to(device=device) if device is not None else values if prop in _INTEGER_PROPERTIES: - dtype = torch.int64 # dtype-ok: dtype var for index tensors; int64 index required + dtype = get_int_dtype() elif prop in _BOOL_PROPERTIES: dtype = torch.bool else: diff --git a/torchref/topology/restraints.py b/torchref/topology/restraints.py index 0c49868..f1b65c7 100644 --- a/torchref/topology/restraints.py +++ b/torchref/topology/restraints.py @@ -26,17 +26,16 @@ import torch from torch.nn import Module +from torchref.config import get_float_dtype, get_int_dtype from torchref.topology.monomer.cif import ( find_cif_file_in_library, read_cif, read_link_definitions, ) -from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin - class Restraints(DeviceMixin, DebugMixin, Module): """ Restraints handler for crystallographic model refinement. @@ -223,13 +222,6 @@ def get_vdw_radii(self) -> torch.Tensor: # Restraint storage # ========================================================================= - - - - - - - @property def restraints(self) -> dict: """Restraint groups as ``[edge type][origin][property]``. @@ -339,7 +331,6 @@ def _load_cif_dictionaries(self, cif_path): f"and will have no restraints applied: {self.missing_residues}" ) - def _load_rama_surfaces(self, device: torch.device): """Load pre-computed Ramachandran NLL surfaces as a buffer.""" from torchref.topology.ramachandran import load_nll_surfaces @@ -391,12 +382,6 @@ def build_restraints(self): self.debug_on_error(e, context="Restraints.build_restraints") raise - - - - - - def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): """Atom pairs within ``cutoff`` of each other, as (M, 2) rows with i < j. @@ -408,7 +393,7 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): n_atoms = xyz.shape[0] if n_atoms == 0: - return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required + return torch.tensor([], dtype=get_int_dtype(), device=device).reshape(0, 2) # Work on CPU to avoid per-iteration GPU kernel launch overhead coords = xyz.detach().cpu() @@ -433,13 +418,13 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): sorted_flat, return_counts=True ) n_unique = len(unique_cells) - starts = torch.zeros(n_unique + 1, dtype=torch.long) # dtype-ok: grid-cell CSR start offsets; int64 index required + starts = torch.zeros(n_unique + 1, dtype=get_int_dtype()) starts[1:] = counts.cumsum(0) # Lookup: flat_cell -> index in unique_cells (-1 if empty) n_grid = gx * gyz - cell_lookup = torch.full((n_grid,), -1, dtype=torch.long) # dtype-ok: cell lookup table (-1 sentinel); int64 index required - cell_lookup[unique_cells] = torch.arange(n_unique) + cell_lookup = torch.full((n_grid,), -1, dtype=get_int_dtype()) + cell_lookup[unique_cells] = torch.arange(n_unique, dtype=cell_lookup.dtype) # 14 unique neighbour offsets: self (0,0,0) + 13 forward neighbours. # "Forward" = first non-zero component is positive, avoiding double counting. @@ -524,9 +509,9 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): if pair_chunks: all_pairs = np.concatenate(pair_chunks, axis=0) - return torch.from_numpy(all_pairs).to(dtype=torch.long, device=device) # dtype-ok: atom-pair index array from numpy; int64 index required + return torch.from_numpy(all_pairs).to(dtype=get_int_dtype(), device=device) else: - return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required + return torch.tensor([], dtype=get_int_dtype(), device=device).reshape(0, 2) def _expand_with_symmetry_mates(self, xyz, cutoff): """Append symmetry-mate positions to ASU ``xyz`` for neighbour search. @@ -651,7 +636,6 @@ def h_topo(self): """Access riding hydrogen topology (None if not built).""" return getattr(self, "_h_topo", None) - def _build_h_exclusion_hash(self, h_topo, device): """Sorted 1-D hash tensor of H-specific 1-2 and 1-3 exclusions. @@ -659,7 +643,8 @@ def _build_h_exclusion_hash(self, h_topo, device): ``torch.searchsorted`` lookup. """ if h_topo is None or h_topo.n_hydrogens == 0: - return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index tensor; int64 index required + # dtype-ok: packed pair key min*max_idx+max overflows int32 beyond ~46k atoms; searchsorted needs both sides int64 + return torch.tensor([], dtype=torch.long, device=device) n_heavy = len(self.pdb) n_h = h_topo.n_hydrogens @@ -684,13 +669,15 @@ def _build_h_exclusion_hash(self, h_topo, device): exclusions.add((min(h_combined, nb), max(h_combined, nb))) if not exclusions: - return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index tensor; int64 index required + # dtype-ok: packed pair key min*max_idx+max overflows int32 beyond ~46k atoms; searchsorted needs both sides int64 + return torch.tensor([], dtype=torch.long, device=device) arr = np.array(list(exclusions), dtype=np.int64) max_idx = max(n_heavy + n_h, int(arr.max()) + 1) hashes = arr[:, 0] * max_idx + arr[:, 1] hashes.sort() - return torch.tensor(hashes, dtype=torch.long, device=device) # dtype-ok: grid-cell hash values used as keys/index; int64 required + # dtype-ok: packed pair key min*max_idx+max overflows int32 beyond ~46k atoms; searchsorted needs both sides int64 + return torch.tensor(hashes, dtype=torch.long, device=device) def _build_vdw_restraints( self, cutoff=6.0, sigma=0.2, inter_residue_only=True, use_spatial_hash=True @@ -893,17 +880,23 @@ def _build_vdw_restraints_legacy( if dist_sq < cutoff_sq: pairs_list.append([i, j]) nearby_pairs = ( - torch.tensor(pairs_list, dtype=torch.long, device=device) # dtype-ok: atom-pair index tensor; int64 index required + torch.tensor(pairs_list, dtype=get_int_dtype(), device=device) if pairs_list - else torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required + else torch.tensor([], dtype=get_int_dtype(), device=device).reshape( + 0, 2 + ) ) empty_result = { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape(0, 2), # dtype-ok: empty atom-pair index tensor; int64 index required + "indices": torch.tensor([], dtype=get_int_dtype(), device=device).reshape( + 0, 2 + ), "min_distances": torch.tensor([], dtype=get_float_dtype(), device=device), "sigmas": torch.tensor([], dtype=get_float_dtype(), device=device), - "symop_indices": torch.tensor([], dtype=torch.long, device=device), # dtype-ok: empty symop index tensor; int64 index required - "cell_offsets": torch.tensor([], dtype=torch.long, device=device).reshape(0, 3), # dtype-ok: empty cell-offset index tensor; int64 index required + "symop_indices": torch.tensor([], dtype=get_int_dtype(), device=device), + "cell_offsets": torch.tensor( + [], dtype=get_int_dtype(), device=device + ).reshape(0, 3), } if len(nearby_pairs) == 0: @@ -1035,7 +1028,7 @@ def _build_vdw_restraints_legacy( # Store results final_pairs = np.stack([final_i1, final_i2], axis=1) self._vdw = { - "indices": torch.tensor(final_pairs, dtype=torch.long, device=device), # dtype-ok: final atom-pair index tensor; int64 index required + "indices": torch.tensor(final_pairs, dtype=get_int_dtype(), device=device), "min_distances": torch.tensor( min_distances, dtype=get_float_dtype(), device=device ), @@ -1043,10 +1036,10 @@ def _build_vdw_restraints_legacy( (len(final_pairs),), sigma, dtype=get_float_dtype(), device=device ), "symop_indices": torch.tensor( - final_symop, dtype=torch.long, device=device # dtype-ok: symop index tensor; int64 index required + final_symop, dtype=get_int_dtype(), device=device ), "cell_offsets": torch.tensor( - final_offsets, dtype=torch.long, device=device # dtype-ok: cell-offset index tensor; int64 index required + final_offsets, dtype=get_int_dtype(), device=device ), } @@ -1162,8 +1155,6 @@ def get_count(rtype, origin): f"torsions={n_torsions}, peptide_bonds={n_bonds_peptide})" ) - - def bond_lengths(self, idx, xyz: torch.Tensor = None): """ Compute current bond lengths from atomic coordinates. @@ -1505,7 +1496,6 @@ def _wrap_torsion_periodicity(self, diff_rad, periods): # All periods are 0 or 1, simple wrapping return torch.remainder(diff_rad + torch.pi, 2.0 * torch.pi) - torch.pi - def torsion_deviations_with_sigmas(self, xyz: torch.Tensor = None): """ Compute torsion deviations (wrapped for periodicity) and sigmas. @@ -1538,9 +1528,6 @@ def torsion_deviations_with_sigmas(self, xyz: torch.Tensor = None): return deviations_rad, sigmas_deg - - - def adp_b_differences(self, adp: torch.Tensor = None): """ Compute B-factor differences between bonded atoms. @@ -1571,4 +1558,3 @@ def adp_b_differences(self, adp: torch.Tensor = None): if diffs_list: return torch.cat(diffs_list, dim=0) return torch.tensor([], device=b_factors.device) - diff --git a/torchref/topology/riding.py b/torchref/topology/riding.py index f36a681..29e1872 100644 --- a/torchref/topology/riding.py +++ b/torchref/topology/riding.py @@ -22,7 +22,7 @@ import numpy as np import torch -from torchref.config import dtypes, normalize_device +from torchref.config import dtypes, get_int_dtype, normalize_device from torchref.utils.device_resolution import resolve_device from torchref.utils.device_mixin import DeviceMixin @@ -431,17 +431,19 @@ def build_hydrogen_topology( fdtype = dtypes.float if n_h_total == 0: - topo.h_parent_idx = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: parent atom-index tensor (empty); int64 required + topo.h_parent_idx = torch.zeros(0, dtype=get_int_dtype(), device=device) topo.h_bond_length = torch.zeros(0, dtype=fdtype, device=device) topo.h_vdw_radius = torch.zeros(0, dtype=fdtype, device=device) - topo.h_placement_type = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: categorical H placement-type code (empty) - topo.h_slot_in_parent = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: slot index into parent (empty); int64 + topo.h_placement_type = torch.zeros(0, dtype=get_int_dtype(), device=device) + topo.h_slot_in_parent = torch.zeros(0, dtype=get_int_dtype(), device=device) topo.parent_neighbor_idx = torch.zeros( - 0, MAX_HEAVY_NB, dtype=torch.long, device=device # dtype-ok: parent neighbor atom-index tensor (empty); int64 required + 0, MAX_HEAVY_NB, dtype=get_int_dtype(), device=device ) - topo.parent_neighbor_count = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: per-parent neighbor count (empty); structural int - topo.h_chainid_enc = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: categorical chain-id encoding (empty) - topo.h_resseq = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: residue sequence id (empty); categorical + topo.parent_neighbor_count = torch.zeros( + 0, dtype=get_int_dtype(), device=device + ) + topo.h_chainid_enc = torch.zeros(0, dtype=get_int_dtype(), device=device) + topo.h_resseq = torch.zeros(0, dtype=get_int_dtype(), device=device) return topo # Sort all topology arrays by placement type for contiguous slicing @@ -466,22 +468,26 @@ def build_hydrogen_topology( idxs = np.where(mask)[0] type_bounds[t] = (int(idxs[0]), int(idxs[-1]) + 1) - topo.h_parent_idx = torch.tensor(acc_parent_idx, dtype=torch.long, device=device) # dtype-ok: parent atom-index tensor; torch indexing requires int64 + topo.h_parent_idx = torch.tensor( + acc_parent_idx, dtype=get_int_dtype(), device=device + ) topo.h_bond_length = torch.tensor(acc_bond_length, dtype=fdtype, device=device) topo.h_vdw_radius = torch.full((n_h_total,), 1.20, dtype=fdtype, device=device) topo.h_placement_type = torch.tensor( - acc_placement_type, dtype=torch.long, device=device # dtype-ok: categorical H placement-type code; used for sort/slice + acc_placement_type, dtype=get_int_dtype(), device=device ) - topo.h_slot_in_parent = torch.tensor(acc_slot, dtype=torch.long, device=device) # dtype-ok: slot index into parent neighbor slots; int64 + topo.h_slot_in_parent = torch.tensor(acc_slot, dtype=get_int_dtype(), device=device) topo.parent_neighbor_idx = torch.tensor( - np.stack(acc_nb_idx), dtype=torch.long, device=device # dtype-ok: parent neighbor atom-index tensor; int64 required + np.stack(acc_nb_idx), dtype=get_int_dtype(), device=device ) topo.parent_neighbor_count = torch.tensor( - acc_nb_count, dtype=torch.long, device=device # dtype-ok: per-parent neighbor count; structural int metadata + acc_nb_count, dtype=get_int_dtype(), device=device ) topo.type_bounds = type_bounds # dict: type_code -> (start, end) - topo.h_chainid_enc = torch.tensor(acc_chainid_enc, dtype=torch.long, device=device) # dtype-ok: categorical chain-id encoding - topo.h_resseq = torch.tensor(acc_resseq, dtype=torch.long, device=device) # dtype-ok: residue sequence id; categorical + topo.h_chainid_enc = torch.tensor( + acc_chainid_enc, dtype=get_int_dtype(), device=device + ) + topo.h_resseq = torch.tensor(acc_resseq, dtype=get_int_dtype(), device=device) if verbose > 0: print(f" Hydrogen topology: {n_h_total} riding H atoms") @@ -736,8 +742,10 @@ def build_h_candidate_pairs( if n_h == 0: for name in ("cand_idx_i", "cand_idx_j", "cand_symop_idx"): - setattr(h_topo, name, torch.zeros(0, dtype=torch.long, device=device)) # dtype-ok: candidate atom/symop index tensors (empty); int64 required - h_topo.cand_cell_offset = torch.zeros(0, 3, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors (empty); symmetry metadata + setattr(h_topo, name, torch.zeros(0, dtype=get_int_dtype(), device=device)) + h_topo.cand_cell_offset = torch.zeros( + 0, 3, dtype=get_int_dtype(), device=device + ) h_topo.cand_min_dist = torch.zeros(0, dtype=dtypes.float, device=device) return @@ -836,15 +844,17 @@ def _same_res(chain_a, resseq_a, chain_b, resseq_b): if not acc_idx_i: for name in ("cand_idx_i", "cand_idx_j", "cand_symop_idx"): - setattr(h_topo, name, torch.zeros(0, dtype=torch.long, device=device)) # dtype-ok: candidate atom/symop index tensors (empty); int64 required - h_topo.cand_cell_offset = torch.zeros(0, 3, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors (empty); symmetry metadata + setattr(h_topo, name, torch.zeros(0, dtype=get_int_dtype(), device=device)) + h_topo.cand_cell_offset = torch.zeros( + 0, 3, dtype=get_int_dtype(), device=device + ) h_topo.cand_min_dist = torch.zeros(0, dtype=dtypes.float, device=device) return - cand_i = torch.tensor(acc_idx_i, dtype=torch.long, device=device) # dtype-ok: combined atom-index tensor; torch indexing requires int64 - cand_j = torch.tensor(acc_idx_j, dtype=torch.long, device=device) # dtype-ok: combined atom-index tensor; torch indexing requires int64 - cand_sym = torch.tensor(acc_symop, dtype=torch.long, device=device) # dtype-ok: symmetry-operator index; int64 - cand_off = torch.tensor(np.stack(acc_offset), dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata + cand_i = torch.tensor(acc_idx_i, dtype=get_int_dtype(), device=device) + cand_j = torch.tensor(acc_idx_j, dtype=get_int_dtype(), device=device) + cand_sym = torch.tensor(acc_symop, dtype=get_int_dtype(), device=device) + cand_off = torch.tensor(np.stack(acc_offset), dtype=get_int_dtype(), device=device) # Apply 1-2 / 1-3 exclusions for intra-ASU candidates if h_excl_hash is not None and len(h_excl_hash) > 0: @@ -853,7 +863,8 @@ def _same_res(chain_a, resseq_a, chain_b, resseq_b): max_idx = n_heavy + n_h norm_i = torch.minimum(cand_i, cand_j) norm_j = torch.maximum(cand_i, cand_j) - pair_hash = norm_i * max_idx + norm_j + # dtype-ok: packed pair key overflows int32; searchsorted needs int64 like the table + pair_hash = norm_i.to(torch.int64) * max_idx + norm_j.to(torch.int64) ins = torch.searchsorted(h_excl_hash, pair_hash).clamp( max=len(h_excl_hash) - 1 ) diff --git a/torchref/topology/topology.py b/torchref/topology/topology.py index 5141999..ab777ac 100644 --- a/torchref/topology/topology.py +++ b/torchref/topology/topology.py @@ -14,6 +14,7 @@ import numpy as np import torch +from torchref.config import get_int_dtype from torchref.topology.atom_graph import AtomGraph from torchref.topology.residue_graph import ResidueGraph from torchref.utils.device_mixin import DeviceMixin @@ -89,7 +90,7 @@ def subset(self, keep) -> "Topology": mask = torch.as_tensor(keep) if mask.dtype != torch.bool: selected = torch.zeros(self.n_atoms, dtype=torch.bool) - selected[mask.to(torch.int64)] = True # dtype-ok: boolean-mask->index cast for scatter select; int64 index required + selected[mask.to(get_int_dtype())] = True mask = selected mask = mask.to(device=self.atoms.residue_of.device) @@ -97,8 +98,10 @@ def subset(self, keep) -> "Topology": raise ValueError("subset would keep no atoms") n_kept = int(mask.sum()) - remap = torch.full((self.n_atoms,), -1, dtype=torch.int64, device=mask.device) # dtype-ok: atom remap index array (-1 sentinel); int64 index required - remap[mask] = torch.arange(n_kept, dtype=torch.int64, device=mask.device) # dtype-ok: arange remap indices; int64 index required + remap = torch.full( + (self.n_atoms,), -1, dtype=get_int_dtype(), device=mask.device + ) + remap[mask] = torch.arange(n_kept, dtype=get_int_dtype(), device=mask.device) # A residue survives if any of its atoms does. Counting per residue also # gives the new atom ranges, contiguous because the atom order is unchanged. @@ -112,10 +115,10 @@ def subset(self, keep) -> "Topology": atom_start = atom_end - counts residue_remap = torch.full( - (self.n_residues,), -1, dtype=torch.int64, device=mask.device # dtype-ok: residue remap index array (-1 sentinel); int64 index required + (self.n_residues,), -1, dtype=get_int_dtype(), device=mask.device ) residue_remap[torch.as_tensor(residue_keep, device=mask.device)] = torch.arange( - int(residue_keep.sum()), dtype=torch.int64, device=mask.device # dtype-ok: arange residue remap indices; int64 index required + int(residue_keep.sum()), dtype=get_int_dtype(), device=mask.device ) return Topology(