torchref.symmetry.grid_utils module

Utilities for determining FFT-compatible grid sizes for crystallographic symmetry.

For interpolation-free symmetry expansion, grid dimensions must be compatible with the symmetry operations. Specifically: - Screw axes require specific divisibility constraints - Grid sizes should also be FFT-friendly (factors of 2, 3, 5)

This module provides the canonical implementations of FFT-friendly grid utilities. The spacegroup module imports and re-exports these functions for convenience.

torchref.symmetry.grid_utils.get_symmetry_grid_requirements(space_group)[source]

Get grid size requirements for a given space group.

This is a convenience wrapper around spacegroup.get_grid_requirements().

Returns a dict with keys ‘nx_mod’, ‘ny_mod’, ‘nz_mod’ indicating the required divisibility for each axis.

Parameters:

space_group (str) – Space group symbol (e.g., ‘P21’, ‘P212121’, ‘P41’, etc.).

Returns:

{‘nx_mod’: int, ‘ny_mod’: int, ‘nz_mod’: int} Required divisibility for each axis.

Return type:

dict

torchref.symmetry.grid_utils.find_fft_friendly_size(n, divisibility=1)[source]

Find the nearest FFT-friendly size >= n that satisfies divisibility constraint.

FFT-friendly means factors only of 2, 3, and 5 (radix-2,3,5 FFT algorithms).

Parameters:

  • n (int) – Minimum grid size.

  • divisibility (int, default 1) – Required divisibility (e.g., 2 for screw axes).

Returns:

Optimal grid size.

Return type:

int

torchref.symmetry.grid_utils.is_fft_friendly(n)[source]

Check if a number has only factors of 2, 3, and 5.

These are optimal for radix-2,3,5 FFT algorithms.

torchref.symmetry.grid_utils.calculate_optimal_grid_size(cell_params, max_res, space_group)[source]

Calculate optimal grid size for a given unit cell and space group.

Grid sizes are chosen to:

  1. Satisfy Shannon-Nyquist sampling (3x oversampling relative to max_res)

  2. Respect symmetry requirements (screw axis divisibility)

  3. Be FFT-friendly (factors of 2, 3, 5 only)

Parameters:

  • cell_params (array-like, shape (6,)) – Unit cell [a, b, c, alpha, beta, gamma].

  • max_res (float) – Maximum resolution in Angstroms.

  • space_group (str) – Space group symbol.

Returns:

Optimal grid dimensions (nx, ny, nz).

Return type:

tuple

torchref.symmetry.grid_utils.check_grid_compatibility(grid_shape, space_group)[source]

Check if a grid is compatible with the space group symmetry.

This is a convenience wrapper around spacegroup.check_grid_compatibility().

Parameters:

  • grid_shape (tuple) – Grid dimensions (nx, ny, nz).

  • space_group (str) – Space group symbol.

Returns:

Dictionary with the following keys:

  • ’compatible’ : bool

  • ’issues’ : list of str (description of problems)

  • ’requirements’ : dict (required divisibility)

  • ’can_use_direct_indexing’ : bool (True if interpolation not needed)

  • ’fft_friendly’ : bool

Return type:

dict

torchref.symmetry.grid_utils.recommend_grid_size(current_shape, space_group)[source]

Recommend a nearby compatible grid size.

Parameters:

  • current_shape (tuple) – Current (nx, ny, nz).

  • space_group (str) – Space group symbol.

Returns:

Recommended (nx, ny, nz).

Return type:

tuple