torchref.base.reciprocal package
Reciprocal space functions for crystallography.
This submodule provides functions for working with reciprocal space: - Reciprocal basis matrix calculations - Miller index operations (HKL generation, d-spacing) - Symmetry operations in reciprocal space - Grid operations for structure factor placement/extraction - Interpolation functions for rotation/translation searches
- torchref.base.reciprocal.reciprocal_basis_matrix(cell)[source]
Compute the reciprocal space basis matrix from unit cell parameters.
- Parameters:
cell (torch.Tensor) – Cell parameters [a, b, c, alpha, beta, gamma].
- Returns:
Reciprocal basis matrix of shape (3, 3) with a*, b*, c* as rows.
- Return type:
- torchref.base.reciprocal.reciprocal_basis_matrix_numpy(cell)[source]
Calculate the reciprocal basis matrix from unit cell parameters (NumPy version).
Computes the reciprocal space basis vectors (a*, b*, c*) that define the transformation from Miller indices to scattering vectors.
- 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.
- Returns:
3x3 matrix containing reciprocal basis vectors as rows [a*, b*, c*].
- Return type:
- torchref.base.reciprocal.get_scattering_vectors(hkl, cell, recB=None)[source]
Calculate scattering vectors from Miller indices.
- Parameters:
hkl (torch.Tensor) – Miller indices of shape (N, 3).
cell (torch.Tensor) – Cell parameters [a, b, c, alpha, beta, gamma].
recB (torch.Tensor, optional) – Pre-computed reciprocal basis matrix of shape (3, 3).
- Returns:
Scattering vectors of shape (N, 3).
- Return type:
- torchref.base.reciprocal.get_scattering_vectors_numpy(hkl, cell)[source]
Calculate scattering vectors from Miller indices and unit cell (NumPy version).
Transforms Miller indices to reciprocal space scattering vectors using the reciprocal basis matrix.
- Parameters:
hkl (numpy.ndarray or list) – Miller indices with shape (N, 3).
cell (numpy.ndarray or list) – Unit cell parameters [a, b, c, alpha, beta, gamma] where lengths are in Angstroms and angles are in degrees.
- Returns:
Scattering vectors in reciprocal space with shape (N, 3).
- Return type:
- torchref.base.reciprocal.get_s(hkl, cell)[source]
Calculate the magnitude of scattering vectors for given Miller indices.
Computes |s| = 1/d where d is the interplanar spacing for each reflection.
- Parameters:
hkl (numpy.ndarray) – Miller indices with shape (N, 3).
cell (numpy.ndarray or list) – Unit cell parameters [a, b, c, alpha, beta, gamma] where lengths are in Angstroms and angles are in degrees.
- Returns:
Magnitude of scattering vectors with shape (N,).
- Return type:
- torchref.base.reciprocal.generate_possible_hkl(cell, d_min, device=None)[source]
Generate all possible Miller indices within a resolution limit.
Creates a complete set of (h, k, l) indices where the d-spacing is greater than or equal to d_min.
- Parameters:
cell (torch.Tensor, shape (6,)) – Unit cell parameters [a, b, c, alpha, beta, gamma] in Angstroms and degrees.
d_min (float) – High resolution limit in Angstroms (minimum d-spacing).
device (torch.device, optional) – Device for computation. If None, uses cell’s device.
- Returns:
All Miller indices with d-spacing >= d_min.
- Return type:
torch.Tensor, shape (M, 3), dtype int32
Examples
import torch cell = torch.tensor([50.0, 60.0, 70.0, 90.0, 90.0, 90.0]) hkl = generate_possible_hkl(cell, d_min=2.0) print(f"Generated {len(hkl)} reflections")
- torchref.base.reciprocal.get_d_spacing(hkl, cell, recB=None)[source]
Calculate d-spacing from Miller indices.
- Parameters:
hkl (torch.Tensor) – Miller indices of shape (N, 3).
cell (torch.Tensor) – Cell parameters [a, b, c, alpha, beta, gamma].
recB (torch.Tensor, optional) – Pre-computed reciprocal basis matrix of shape (3, 3).
- Returns:
D-spacing values of shape (N,) in Angstroms.
- Return type:
- torchref.base.reciprocal.compute_d_spacing_batch(hkl, cell, recB=None)[source]
Compute d-spacing for a batch of Miller indices.
Wrapper around get_d_spacing for convenience.
- Parameters:
hkl (torch.Tensor, shape (N, 3)) – Miller indices.
cell (torch.Tensor, shape (6,)) – Unit cell parameters.
recB (torch.Tensor, optional) – Pre-computed reciprocal basis matrix.
- Returns:
D-spacing values in Angstroms.
- Return type:
torch.Tensor, shape (N,)
- torchref.base.reciprocal.place_on_grid(hkls, structure_factor, grid_size, enforce_hermitian=True)[source]
Place structure factors on a reciprocal-space grid.
Vectorized placement of batched structure factors on reciprocal-space grid.
- Parameters:
hkls (torch.Tensor) – Miller indices of shape (N, 3).
structure_factor (torch.Tensor) – Structure factors of shape (N,) or (B, N) for batched input.
grid_size (tuple or torch.Tensor) – Grid dimensions (Nx, Ny, Nz).
enforce_hermitian (bool, optional) – Whether to enforce Hermitian symmetry. Default is True.
- Returns:
Complex tensor grid of structure factors of shape (Nx, Ny, Nz) or (B, Nx, Ny, Nz) for batched input.
- Return type:
- torchref.base.reciprocal.extract_structure_factor_from_grid(reciprocal_grid, hkls)[source]
Extract structure factors from reciprocal space grid at given Miller indices.
- Parameters:
reciprocal_grid (torch.Tensor) – Complex tensor of shape (Nx, Ny, Nz) or (B, Nx, Ny, Nz).
hkls (torch.Tensor) – Miller indices of shape (N, 3).
- Returns:
Structure factors of shape (N,) or (B, N) for batched input.
- Return type:
- torchref.base.reciprocal.apply_translation_phase(F_calc, hkl, translation_frac)[source]
Apply translation phase shift to structure factors.
For a translation t in fractional coordinates, the structure factor transforms as: F’(hkl) = F(hkl) * exp(2πi * hkl · t)
- Parameters:
F_calc (torch.Tensor) – Complex structure factors of shape (N,).
hkl (torch.Tensor) – Miller indices of shape (N, 3).
translation_frac (torch.Tensor) – Translation vector in fractional coordinates of shape (3,).
- Returns:
Phase-shifted structure factors of shape (N,).
- Return type:
- torchref.base.reciprocal.interpolate_structure_factor_from_grid(reciprocal_grid, hkl_float, interpolate_amplitude=True)[source]
Interpolate structure factors from reciprocal space grid at non-integer positions.
- Parameters:
reciprocal_grid (torch.Tensor) – Complex tensor of shape (Nx, Ny, Nz).
hkl_float (torch.Tensor) – Non-integer HKL positions of shape (N, 3).
interpolate_amplitude (bool, optional) – If True (default), interpolate amplitudes instead of complex values. This avoids phase cancellation issues where linear interpolation of complex numbers with different phases can give incorrect magnitudes (e.g., interpolating F=1 and F=-1 gives 0 instead of 1).
- Returns:
Interpolated structure factors of shape (N,). If interpolate_amplitude=True, returns real-valued amplitudes. If interpolate_amplitude=False, returns complex values (use with caution).
- Return type:
Notes
For a rotation R applied to the model, the structure factor at hkl becomes F(R^T @ hkl), so you would call this with hkl_float = hkl @ R.
WARNING: Complex interpolation (interpolate_amplitude=False) can give incorrect results when neighboring voxels have very different phases. For example, if F1 = A*exp(i*0) and F2 = A*exp(i*π), linear interpolation gives magnitude 0 at the midpoint instead of A. Use interpolate_amplitude=True for rotation functions where only magnitudes matter.
- torchref.base.reciprocal.interpolate_complex_from_grid(reciprocal_grid, hkl_float)[source]
Interpolate complex structure factors from reciprocal space grid.
Unlike amplitude interpolation, this preserves phase information, which is essential for translation searches where phases are used to compute correlation functions.
- Parameters:
reciprocal_grid (torch.Tensor) – Complex tensor of shape (Nx, Ny, Nz).
hkl_float (torch.Tensor) – Non-integer HKL positions of shape (N, 3).
- Returns:
Interpolated complex structure factors of shape (N,).
- Return type:
Notes
This function performs trilinear interpolation of complex values. For rotation-only searches (where only magnitudes matter), use interpolate_structure_factor_from_grid(interpolate_amplitude=True) instead.
For translation searches, complex interpolation is needed because the translation function depends on the phase relationship between F_obs and F_calc.
WARNING: Complex interpolation can give reduced magnitudes when neighboring voxels have very different phases. This is acceptable for translation searches where we’re computing correlation functions, but not for rotation searches.
- torchref.base.reciprocal.trilinear_interpolate_patterson(grid, points, chunk_size=10000000)[source]
Memory-efficient trilinear interpolation on a 3D grid.
Pure torch implementation for GPU acceleration and gradient flow. Replaces scipy.ndimage.map_coordinates for torch tensors.
- Parameters:
grid (torch.Tensor) – 3D grid of values with shape (nx, ny, nz).
points (torch.Tensor) – Coordinates to sample with shape (n_points, 3). Should be in fractional coordinates [0, 1) for ‘wrap’ mode. Or batch K, n_points, 3 for multiple batches.
chunk_size (int, optional) – Number of points to process at once. Default is 1M.
- Returns:
Interpolated values with shape (n_points,) or (batch, n_points).
- Return type:
Notes
Supports automatic differentiation for gradient-based optimization. Uses chunked processing and in-place accumulation to reduce memory.
- torchref.base.reciprocal.interpolate_for_rotation(hkl, R, cell, reciprocal_space_grid)[source]
Interpolate structure factors for rotated HKL positions. This works for all cells. :param hkl: HKL positions, shape (N, 3). :type hkl: torch.Tensor :param R: Rotation matrix, shape (B, 3, 3). :type R: torch.Tensor :param cell: Unit cell object with reciprocal_basis_matrix. :type cell: Cell :param reciprocal_space_grid: Reciprocal space grid of structure factors, shape (Nx, Ny, Nz). :type reciprocal_space_grid: torch.Tensor
- Returns:
Interpolated structure factors, shape (B, N) / (N)
- Return type:
- torchref.base.reciprocal.smooth_reciprocal_grid(reciprocal_grid, sigma, mode='amplitude_phase')[source]
Apply Gaussian smoothing to a reciprocal space grid using native PyTorch.
- Parameters:
reciprocal_grid (torch.Tensor) – Complex tensor of shape (Nx, Ny, Nz).
sigma (float) – Standard deviation of the Gaussian kernel in voxel units.
mode (str, optional) –
Smoothing mode. Options: - “amplitude_phase”: Smooth amplitude and phase separately using
circular mean for phases. This avoids phase cancellation while preserving the relationship between amplitude and phase. Default.
”amplitude_only”: Smooth only amplitudes, preserve original phases. Useful when phase accuracy is critical.
”complex”: Smooth real and imaginary parts separately. WARNING: This can cause phase cancellation when neighboring voxels have different phases (e.g., averaging exp(i*0) and exp(i*π) gives 0).
- Returns:
Smoothed reciprocal space grid of shape (Nx, Ny, Nz).
- Return type:
Notes
Uses FFT-based convolution for efficient 3D Gaussian smoothing with periodic (wrap) boundary conditions.
- torchref.base.reciprocal.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.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:
hkl (torch.Tensor, shape (N, 3)) – Miller indices.
translations (torch.Tensor, shape (n_ops, 3)) – Translation vectors in fractional coordinates.
- Returns:
Complex phase factors exp(2*pi*i * h.t).
- Return type:
torch.Tensor, shape (n_ops, N)
- torchref.base.reciprocal.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:
reciprocal_grid (torch.Tensor, shape (Nx, Ny, Nz)) – Complex reciprocal space grid from FFT of P1 density map.
hkl (torch.Tensor, shape (N, 3)) – Target Miller indices.
rotation_matrices (torch.Tensor, shape (n_ops, 3, 3)) – Real-space rotation matrices from symmetry operations.
translations (torch.Tensor, shape (n_ops, 3)) – Translation vectors from symmetry operations.
- Returns:
Complex structure factors with symmetry applied.
- Return type:
torch.Tensor, shape (N,)
- class torchref.base.reciprocal.ReciprocalSymmetryExtractor(hkl, symmetry, grid_shape, device=None)[source]
Bases:
DeviceMixinClass-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)
- __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,)
Submodules
- torchref.base.reciprocal.basis module
- torchref.base.reciprocal.grid_operations module
- torchref.base.reciprocal.hkl module
- torchref.base.reciprocal.interpolation module
- torchref.base.reciprocal.symmetry module