torchref.base.scattering.scattering_table module

Table-based ITC92 scattering factor lookup.

This module provides fast, vectorized lookup of scattering factor parameters using a pre-computed table stored as a .pt file. The table eliminates runtime gemmi dependency for scattering parameter access.

The table supports: - Neutral atoms for Z=1 to 103 (H to Lr) - Common ions with various charge states - Element symbol to atomic number mapping

Example

from torchref.base.scattering.scattering_table import (
    load_scattering_table,
    get_scattering_params_by_z,
    get_element_to_z_mapping,
)

# Load the table
table = load_scattering_table(device='cuda')

# Get parameters for multiple atoms at once
z_tensor = torch.tensor([6, 7, 8])  # C, N, O
A, B = get_scattering_params_by_z(z_tensor)
torchref.base.scattering.scattering_table.load_scattering_table(device=None, dtype=None, force_reload=False)[source]

Load pre-computed ITC92 scattering factors from .pt file.

The table is cached globally after first load for efficiency. Use force_reload=True to reload from disk.

Parameters:

  • device (torch.device, optional) – Device to place tensors on. Default is None (keeps original device).

  • dtype (torch.dtype, optional) – Data type for floating point tensors. Default is None (keeps original dtype).

  • force_reload (bool, optional) – Force reload from disk even if cached. Default is False.

Returns:

Dictionary containing: - ‘A’: Tensor(max_z + 1, 5) - neutral A coefficients indexed by Z - ‘B’: Tensor(max_z + 1, 5) - neutral B coefficients indexed by Z - ‘element_to_z’: dict mapping element symbols to atomic numbers - ‘z_to_element’: dict mapping atomic numbers to element symbols - ‘ions’: dict mapping ion keys to (A, B) tuples - ‘metadata’: dict with source information

Return type:

dict

Raises:

FileNotFoundError – If the pre-computed table file does not exist.

Examples

table = load_scattering_table(device='cuda', dtype=torch.float32)
A = table['A']  # Shape (104, 5)
z_to_elem = table['z_to_element']  # {1: 'H', 6: 'C', ...}
torchref.base.scattering.scattering_table.get_element_to_z_mapping()[source]

Return element symbol to atomic number mapping.

This function loads the scattering table if not already cached and returns the element_to_z dictionary.

Returns:

Mapping of element symbols to atomic numbers. Example: {‘H’: 1, ‘C’: 6, ‘N’: 7, ‘O’: 8, …}

Return type:

dict

Examples

element_to_z = get_element_to_z_mapping()
z_carbon = element_to_z['C']  # Returns 6
torchref.base.scattering.scattering_table.get_z_to_element_mapping()[source]

Return atomic number to element symbol mapping.

Returns:

Mapping of atomic numbers to element symbols. Example: {1: ‘H’, 6: ‘C’, 7: ‘N’, 8: ‘O’, …}

Return type:

dict

torchref.base.scattering.scattering_table.get_scattering_params_by_z(z_tensor, device=None, dtype=None)[source]

Fast vectorized lookup of scattering parameters by atomic number.

This function provides efficient batch lookup of ITC92 scattering parameters using tensor indexing. All atoms are looked up in a single operation.

Parameters:

  • z_tensor (torch.Tensor) – Atomic numbers for all atoms, shape (n_atoms,). Values should be in range 1-103.

  • device (torch.device, optional) – Device to place output tensors on. Default uses z_tensor’s device.

  • dtype (torch.dtype, optional) – Data type for output tensors. Default is torch.float32.

Returns:

  • A (torch.Tensor) – ITC92 A parameters (amplitudes) with shape (n_atoms, 5).

  • B (torch.Tensor) – ITC92 B parameters (widths) with shape (n_atoms, 5).

Return type:

Tuple[Tensor, Tensor]

Examples

# Get parameters for C, N, O atoms
z = torch.tensor([6, 7, 8, 6, 6, 7])
A, B = get_scattering_params_by_z(z)
# A.shape == (6, 5), B.shape == (6, 5)

# Use on GPU
z_gpu = z.cuda()
A, B = get_scattering_params_by_z(z_gpu, device='cuda')
torchref.base.scattering.scattering_table.get_scattering_params_for_ion(element, charge, device=None, dtype=None)[source]

Get scattering parameters for a specific ion.

Parameters:

  • element (str) – Element symbol (e.g., ‘Fe’, ‘O’).

  • charge (int) – Ionic charge (positive or negative integer).

  • device (torch.device, optional) – Device to place tensors on.

  • dtype (torch.dtype, optional) – Data type for tensors.

Returns:

(A, B) tensors of shape (5,), or None if ion not found.

Return type:

tuple or None

Examples

# Get Fe2+ parameters
A, B = get_scattering_params_for_ion('Fe', 2)

# Get O2- parameters
A, B = get_scattering_params_for_ion('O', -2)
torchref.base.scattering.scattering_table.elements_to_z(elements, normalize=True)[source]

Convert a list of element symbols to atomic numbers.

Parameters:

  • elements (list of str) – Element symbols (e.g., [‘C’, ‘N’, ‘O’, ‘C’, ‘C’]).

  • normalize (bool, optional) – If True, normalize element names (strip whitespace, capitalize). Default is True.

Returns:

Tensor of atomic numbers with shape (n_atoms,). Unknown elements are assigned Z=0.

Return type:

torch.Tensor

Examples

z = elements_to_z(['C', 'N', 'O'])
# Returns tensor([6, 7, 8])

z = elements_to_z([' c ', 'N', 'O'], normalize=True)
# Returns tensor([6, 7, 8])