torchref.base.reciprocal.symmetry module

Reciprocal space symmetry application for structure factor calculation.

This module provides an alternative to the MapSymmetry-based approach for structure factor calculation. Instead of symmetrizing the density map in real space (“early symmetry”), we apply symmetry directly in reciprocal space after FFT (“late symmetry”).

Late symmetry is approximately 5x faster for structure factor calculation because it avoids the expensive map symmetrization step.

Mathematical Background

For a symmetry operation {R|t} (rotation R, translation t):

F_sym(h) = sum_ops exp(2*pi*i * h.t) * F_P1(R^T @ h)

Where: - F_sym(h) = structure factor with symmetry at Miller index h - F_P1(h’) = P1 structure factor at h’ = R^T @ h - exp(2*pi*i * h.t) = translation phase shift

Since R is an integer matrix (crystallographic), R^T @ h gives integer indices that can be extracted directly from the reciprocal space grid.

Terminology

  • Early symmetry: Apply symmetry to density map before FFT (MapSymmetry approach)

  • Late symmetry: Apply symmetry in reciprocal space after FFT (this module)

torchref.base.reciprocal.symmetry.compute_symmetry_equivalent_hkls(hkl, rotation_matrices)[source]

Compute symmetry-equivalent HKLs: h’ = R^T @ h for each operation.

In reciprocal space, Miller indices transform as h’ = h @ R (or equivalently h’ = R^T @ h when treating h as a column vector).

Parameters:

  • hkl (torch.Tensor, shape (N, 3)) – Miller indices.

  • rotation_matrices (torch.Tensor, shape (n_ops, 3, 3)) – Real-space rotation matrices. These are transposed internally for reciprocal space transformation.

Returns:

Equivalent HKLs for each symmetry operation.

Return type:

torch.Tensor, shape (n_ops, N, 3)

torchref.base.reciprocal.symmetry.compute_translation_phases(hkl, translations)[source]

Compute phase shifts: exp(2*pi*i * h.t) for each operation.

The translation component of a symmetry operation causes a phase shift in the structure factor.

Parameters:
Returns:

Complex phase factors exp(2*pi*i * h.t).

Return type:

torch.Tensor, shape (n_ops, N)

torchref.base.reciprocal.symmetry.extract_structure_factors_with_symmetry(reciprocal_grid, hkl, rotation_matrices, translations)[source]

Extract structure factors with symmetry applied in reciprocal space.

This is the main function that replaces the MapSymmetry approach for structure factor extraction. Instead of symmetrizing the density map and then extracting F(hkl), we extract F at all symmetry-equivalent positions and sum with phases.

Parameters:
Returns:

Complex structure factors with symmetry applied.

Return type:

torch.Tensor, shape (N,)

class torchref.base.reciprocal.symmetry.ReciprocalSymmetryExtractor(hkl, symmetry, grid_shape, device=None)[source]

Bases: DeviceMixin

Class-based interface for reciprocal space symmetry extraction.

This provides a more efficient interface when computing structure factors multiple times with the same symmetry and HKLs (e.g., during refinement). Precomputes equivalent HKLs, phase factors, and flat grid indices so that each call reduces to a single gather + multiply + sum (~3 GPU kernels instead of ~28).

Parameters:

  • hkl (torch.Tensor, shape (N, 3)) – Target Miller indices.

  • symmetry (SpaceGroup) – SpaceGroup object containing rotation matrices and translations.

  • grid_shape (tuple of int) – Reciprocal grid dimensions (Nx, Ny, Nz).

  • device (torch.device, optional) – Device for computation.

Examples

>>> extractor = ReciprocalSymmetryExtractor(hkl, symmetry, grid_shape=(209, 86, 67))
>>> f_calc = extractor.extract_from_grid(reciprocal_grid)
__init__(hkl, symmetry, grid_shape, device=None)[source]
__call__(density_map)[source]

Compute structure factors from P1 density map.

Parameters:

density_map (torch.Tensor, shape (Nx, Ny, Nz)) – P1 electron density map (NO symmetry applied).

Returns:

Complex structure factors with symmetry applied.

Return type:

torch.Tensor, shape (N,)

extract(density_map)[source]

Extract structure factors from P1 density map.

Parameters:

density_map (torch.Tensor, shape (Nx, Ny, Nz)) – P1 electron density map (NO symmetry applied).

Returns:

Complex structure factors with symmetry applied.

Return type:

torch.Tensor, shape (N,)

extract_from_grid(reciprocal_grid)[source]

Extract structure factors from precomputed reciprocal grid.

Uses precomputed flat indices for a single vectorized gather, avoiding per-symop Python loops and kernel launches.

Parameters:

reciprocal_grid (torch.Tensor, shape (Nx, Ny, Nz)) – Complex reciprocal space grid from FFT.

Returns:

Complex structure factors with symmetry applied.

Return type:

torch.Tensor, shape (N,)