torchref.base.fourier package

Fourier transform functions for crystallography.

This submodule provides functions for: - FFT and inverse FFT operations for crystallographic conventions - Real-space grid generation

torchref.base.fourier.fft(reciprocal_grid, volume=None)[source]

Perform FFT to obtain real space electron density.

Uses fftn with norm=”forward” to match crystallographic sign convention directly, avoiding expensive flip/roll operations.

Crystallographic convention: ρ(r) = (1/V) Σ F(h) exp(-2πi h·r)

PyTorch fftn with norm=”forward” gives:

fftn(x)[n] = (1/N) Σ_k x[k] exp(-2πi k·n/N)

When input structure factors F are correctly scaled (with V/N factor from ifft), we need to multiply by N/V to recover the original electron density:

ρ = fftn(F) * (N / V)

Parameters:

  • reciprocal_grid (torch.Tensor) – Reciprocal space grid of shape (Nx, Ny, Nz) or (B, Nx, Ny, Nz). Expected to contain correctly scaled structure factors (from ifft with volume).

  • volume (float, optional) – Unit cell volume in ų. If provided, result is scaled by N/V to give correctly normalized electron density.

Returns:

Real-valued tensor of electron density with same shape as input.

Return type:

torch.Tensor

torchref.base.fourier.ifft(real_space_map, volume=None)[source]

Perform inverse FFT to obtain reciprocal space structure factors.

Crystallographic convention: F(h) = Σ ρ(r) exp(+2πi h·r) * ΔV where ΔV = V_cell / N is the voxel volume.

PyTorch ifftn with norm=”forward” gives unnormalized DFT:

DFT[k] = Σ x[n] exp(+2πi k·n/N)

To obtain correctly scaled structure factors, we multiply by voxel volume:

F(h) = DFT(ρ) * (V_cell / N)

Parameters:

  • real_space_map (torch.Tensor) – Real space electron density map of shape (Nx, Ny, Nz) or (B, Nx, Ny, Nz).

  • volume (float, optional) – Unit cell volume in ų. If provided, result is scaled by voxel volume (V_cell / N_total) to give correctly normalized structure factors.

Returns:

Complex-valued tensor of structure factors with same shape as input.

Return type:

torch.Tensor

torchref.base.fourier.get_real_grid(cell=None, fractional_matrix=None, max_res=0.8, gridsize=None, device=None)[source]

Generate a real space grid for electron density calculations.

Parameters:

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

  • fractional_matrix (torch.Tensor, optional) – Pre-computed fractionalization matrix.

  • max_res (float, optional) – Maximum resolution for automatic grid sizing. Default is 0.8.

  • gridsize (torch.Tensor or array-like, optional) – Explicit grid dimensions [nx, ny, nz]. If None, calculated from max_res.

  • device (torch.device or str, optional) – Device for tensor placement. If None, inferred from fractional_matrix or cell (whichever tensor is provided); falls back to CPU.

Returns:

Real space grid of shape (nx, ny, nz, 3) containing Cartesian coordinates.

Return type:

torch.Tensor

torchref.base.fourier.find_grid_size(cell, max_res)[source]

Calculate grid size based on unit cell and resolution.

Parameters:

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

  • max_res (float) – Maximum resolution in Angstroms.

Returns:

Grid dimensions [nx, ny, nz] as int32.

Return type:

torch.Tensor

torchref.base.fourier.get_real_grid_numpy(cell, max_res=0.8, gridsize=None)[source]

Generate a real-space grid of Cartesian coordinates (NumPy version).

Creates a 3D grid in fractional coordinates and converts it to Cartesian coordinates. Grid points are placed at cell edges following CCTBX convention.

Parameters:

  • cell (numpy.ndarray or list) – Unit cell parameters [a, b, c, alpha, beta, gamma] where lengths are in Angstroms and angles are in degrees.

  • max_res (float, optional) – Maximum resolution in Angstroms for grid spacing. Default is 0.8. Ignored if gridsize is provided.

  • gridsize (list or numpy.ndarray, optional) – Explicit grid dimensions [nx, ny, nz]. If provided, overrides max_res.

Returns:

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

Return type:

numpy.ndarray

torchref.base.fourier.get_grids(cell, max_res=0.8)[source]

Generate real-space and reciprocal-space grids for Fourier transforms.

Creates a 3D grid in fractional coordinates and converts it to Cartesian coordinates, along with an empty reciprocal space grid.

Parameters:

  • cell (numpy.ndarray or list) – Unit cell parameters [a, b, c, alpha, beta, gamma] where lengths are in Angstroms and angles are in degrees.

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

Returns:

  • recgrid (numpy.ndarray) – Empty reciprocal space grid with shape determined by resolution.

  • xyz_real_grid (numpy.ndarray) – Real-space grid coordinates with shape (nx, ny, nz, 3).

torchref.base.fourier.put_hkl_on_grid(real_space_grid, diff, hkl)[source]

Place structure factors on a reciprocal space grid.

Maps structure factor values to their corresponding positions on a 3D reciprocal space grid based on Miller indices.

Parameters:

  • real_space_grid (numpy.ndarray) – Real-space grid used to determine the reciprocal grid dimensions. Shape should be (nx, ny, nz, 3) or similar.

  • diff (numpy.ndarray) – Structure factor values (complex) to place on the grid.

  • hkl (numpy.ndarray) – Miller indices with shape (N, 3), used as grid indices.

Returns:

Complex reciprocal space grid with shape (nx, ny, nz).

Return type:

numpy.ndarray

Submodules