torchref.model.sf_fft module

SfFFT - Structure Factor calculation via FFT (Fast Fourier Transform).

This module provides a PyTorch nn.Module that handles: - Grid setup for real-space electron density calculations - Building electron density maps from atomic parameters - FFT-based conversion to structure factors

The SfFFT class can be used standalone (stateless) or with stored grid state (stateful), making it flexible for various use cases.

Note: FFT is provided as a backward compatibility alias for SfFFT.

class torchref.model.sf_fft.SfFFT(cell=None, spacegroup=None, max_res=1.5, radius_angstrom=3.0, dtype_float=torch.float32, device=None, verbose=0, use_late_symmetry=True)[source]

Bases: DeviceMixin, Module

Structure Factor calculator using FFT (Fast Fourier Transform).

This module encapsulates all FFT-related functionality for computing electron density maps and structure factors. It is initialized with a Cell and optionally a SpaceGroup, which are used for grid calculations.

Parameters:

  • cell (Cell) – Unit cell object containing cell parameters.

  • spacegroup (SpaceGroupLike, optional) – Space group specification (string, int, or gemmi.SpaceGroup). If None, defaults to P1.

  • max_res (float, optional) – Maximum resolution for grid spacing in Angstroms. Default is 1.0.

  • radius_angstrom (float, optional) – Radius in Angstroms for density calculation around each atom. Default is 4.0.

  • dtype_float (torch.dtype, optional) – Data type for floating point tensors. Default is dtypes.float.

  • device (torch.device, optional) – Computation device. Defaults to the configured device.current.

  • verbose (int, optional) – Verbosity level for logging. Default is 0.

cell

Unit cell object.

Type:

Cell

spacegroup

Space group object (SpaceGroup nn.Module with matrices and translations).

Type:

SpaceGroup

symmetry

Alias for spacegroup (backward compatibility).

Type:

SpaceGroup

max_res

Maximum resolution for grid spacing.

Type:

float

radius_angstrom

Radius for density calculation around each atom.

Type:

float

gridsize

Grid dimensions (nx, ny, nz) when grid is set up.

Type:

torch.Tensor or None

real_space_grid

Real-space coordinate grid with shape (nx, ny, nz, 3).

Type:

torch.Tensor or None

voxel_size

Voxel dimensions.

Type:

torch.Tensor or None

map_symmetry

Symmetry operator for map calculations.

Type:

MapSymmetry or None

Examples

Standalone usage:

from torchref.symmetry import Cell
cell = Cell([50, 60, 70, 90, 90, 90])
sf_fft = SfFFT(cell, spacegroup='P212121', max_res=1.5)
sf_fft.setup_grid()
density_map = sf_fft.build_density_map(xyz, b, occ, A, B, inv_frac, frac)
sf = sf_fft.map_to_structure_factors(density_map, hkl)

With ModelFT (composition):

model = ModelFT()
model.load_pdb('structure.pdb')
sf = model.get_structure_factor(hkl)  # Uses internal SfFFT instance
__init__(cell=None, spacegroup=None, max_res=1.5, radius_angstrom=3.0, dtype_float=torch.float32, device=None, verbose=0, use_late_symmetry=True)[source]

Initialize the SfFFT module with cell and spacegroup.

Parameters:

  • cell (Cell, optional) – Unit cell object. If None, must be set later via set_cell().

  • spacegroup (SpaceGroupLike, optional) – Space group specification. If None, defaults to P1.

  • max_res (float, optional) – Maximum resolution for grid spacing in Angstroms. Default is 1.5.

  • radius_angstrom (float, optional) – Radius in Angstroms for density calculation. Default is 3.0.

  • dtype_float (torch.dtype, optional) – Data type for floating point tensors. Default is dtypes.float.

  • device (torch.device, optional) – Computation device. Default is None (uses cell’s device). If Cell is also None, defaults to CPU.

  • verbose (int, optional) – Verbosity level for logging. Default is 0.

  • use_late_symmetry (bool, optional) – If True (default), apply symmetry in reciprocal space after FFT (“late symmetry”) for faster structure factor calculation (~5x speedup). If False, apply symmetry to density map before FFT (“early symmetry”).

property cell: Cell | None

Unit cell object.

property spacegroup: SpaceGroup | None

Space group object (SpaceGroup nn.Module).

property symmetry: SpaceGroup | None

Symmetry operations handler (alias for spacegroup).

property fractional_matrix: Tensor | None

Get fractionalization matrix from cell, on this module’s device/dtype.

property inv_fractional_matrix: Tensor | None

Get orthogonalization matrix from cell, on this module’s device/dtype.

set_cell_and_spacegroup(cell, spacegroup=None)[source]

Set cell and spacegroup for this SfFFT instance.

Parameters:

  • cell (Cell) – Unit cell object.

  • spacegroup (SpaceGroupLike, optional) – Space group specification.

compute_optimal_gridsize(max_res=None)[source]

Compute optimal grid dimensions using the stored cell and spacegroup.

Uses Cell.compute_grid_size() for base calculation and Symmetry.suggest_grid_size() for symmetry optimization.

Parameters:

max_res (float, optional) – Maximum resolution in Angstroms. If None, uses self.max_res.

Returns:

Optimal grid dimensions (nx, ny, nz).

Return type:

tuple of int

Raises:

RuntimeError – If cell has not been set.

static compute_real_space_grid(fractional_matrix, gridsize, device=device(type='cpu'))[source]

Generate the real-space coordinate grid.

Parameters:

  • cell_data (torch.Tensor) – Unit cell parameters [a, b, c, alpha, beta, gamma].

  • gridsize (torch.Tensor) – Grid dimensions (nx, ny, nz).

  • device (torch.device, optional) – Target device. Default is CPU.

Returns:

Real-space grid with shape (nx, ny, nz, 3).

Return type:

torch.Tensor

setup_grid(gridsize=None, max_res=None)[source]

Setup the real-space grid for electron density calculation.

This method initializes and stores the grid state for subsequent density map calculations. Uses the stored cell and spacegroup.

Parameters:

  • gridsize (tuple of int, optional) – Explicit grid size (nx, ny, nz). If None, computed automatically using Cell.compute_grid_size() and Symmetry.suggest_grid_size().

  • max_res (float, optional) – Maximum resolution in Angstroms. If None, uses self.max_res.

Raises:

RuntimeError – If cell has not been set.

build_density_map(xyz_iso, adp_iso, occ_iso, A_iso, B_iso, xyz_aniso=None, u_aniso=None, occ_aniso=None, A_aniso=None, B_aniso=None, apply_symmetry=True)[source]

Build electron density map from atomic parameters.

This method requires setup_grid() to have been called first.

Parameters:

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

  • adp_iso (torch.Tensor) – Isotropic ADPs (atomic displacement parameters) with shape (n_iso,).

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

  • A_iso (torch.Tensor) – ITC92 A parameters for isotropic atoms with shape (n_iso, 5).

  • B_iso (torch.Tensor) – ITC92 B parameters for isotropic atoms with shape (n_iso, 5).

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

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

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

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

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

  • apply_symmetry (bool, optional) – If True, apply crystallographic symmetry to the map. Default is True.

Returns:

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

Return type:

torch.Tensor

Raises:

RuntimeError – If setup_grid() has not been called.

map_to_structure_factors(density_map, hkl, apply_symmetry=True)[source]

Convert density map to structure factors via FFT.

Parameters:

  • density_map (torch.Tensor) – Electron density map with shape (nx, ny, nz). If apply_symmetry=True, this should be a P1 density map.

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

  • apply_symmetry (bool, optional) – If True and late symmetry is enabled/compatible, apply symmetry in reciprocal space. Default is False (assume map already has symmetry applied or use early symmetry path).

Returns:

Complex structure factors with shape (n_reflections,).

Return type:

torch.Tensor

compute_structure_factors(hkl, xyz_iso, adp_iso, occ_iso, A_iso, B_iso, xyz_aniso=None, u_aniso=None, occ_aniso=None, A_aniso=None, B_aniso=None, apply_symmetry=True)[source]

Compute structure factors from atomic parameters (end-to-end).

This is a convenience method that builds the density map and computes structure factors in one call.

When use_late_symmetry=True (default) and the grid is compatible, symmetry is applied in reciprocal space after FFT for ~5x speedup. Otherwise, symmetry is applied to the density map before FFT.

Parameters:

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

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

  • adp_iso (torch.Tensor) – Isotropic ADPs (atomic displacement parameters) with shape (n_iso,).

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

  • A_iso (torch.Tensor) – ITC92 A parameters for isotropic atoms.

  • B_iso (torch.Tensor) – ITC92 B parameters for isotropic atoms.

  • xyz_aniso (torch.Tensor, optional) – Anisotropic atom coordinates.

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

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

  • A_aniso (torch.Tensor, optional) – ITC92 A parameters for anisotropic atoms.

  • B_aniso (torch.Tensor, optional) – ITC92 B parameters for anisotropic atoms.

  • apply_symmetry (bool, optional) – If True, apply crystallographic symmetry. Default is True.

Returns:

  • sf (torch.Tensor) – Complex structure factors with shape (n_reflections,).

  • density_map (torch.Tensor) – Electron density map with shape (nx, ny, nz). Note: When using late symmetry, this is the P1 map (without symmetry).

Return type:

Tuple[Tensor, Tensor]

reset_cache()[source]

Drop the cached symmetry extractor; recomputed on next use.

copy()[source]

Create a deep copy of this SfFFT module.

Returns:

A new SfFFT instance with cloned cell, spacegroup, and buffers.

Return type:

SfFFT

torchref.model.sf_fft.FFT(*args, **kwargs)[source]

Deprecated: use SfFFT instead.