torchref.base.alignment.normalization module

Normalization functions for E-value conversion and anisotropy correction.

This module provides GPU-accelerated PyTorch implementations of: - Radial shell computation and assignment - Anisotropy correction for F² values - E-value normalization within resolution shells

These functions are used for molecular replacement and related analyses.

torchref.base.alignment.normalization.compute_radial_shells(d_min, d_max, n_shells, device=None)[source]

Compute uniform radial shell boundaries in reciprocal space.

Shells are spaced uniformly in 1/d (|s|) for even coverage of resolution.

Parameters:

  • d_min (float) – High resolution limit in Angstroms.

  • d_max (float) – Low resolution limit in Angstroms.

  • n_shells (int) – Number of radial shells.

  • device (torch.device, optional) – Device for output tensors. Default is CPU.

Returns:

  • shell_edges (torch.Tensor) – Shell boundaries in Angstroms^-1, shape (n_shells+1,).

  • shell_centers (torch.Tensor) – Shell centers in Angstroms^-1, shape (n_shells,).

Return type:

Tuple[Tensor, Tensor]

torchref.base.alignment.normalization.assign_to_shells(s_mag, shell_edges)[source]

Assign reflections to radial shells.

Parameters:

  • s_mag (torch.Tensor) – |s| values in Angstroms^-1, shape (N,).

  • shell_edges (torch.Tensor) – Shell boundaries in Angstroms^-1, shape (n_shells+1,).

Returns:

shell_idx – Shell index for each reflection, shape (N,). Values 0 to n_shells-1, or -1 for out-of-range.

Return type:

torch.Tensor

torchref.base.alignment.normalization.compute_anisotropy_correction(s_vectors, U)[source]

Compute anisotropic correction factor for F² values.

The correction is: exp(-2*pi^2 * s^T U s)

This scales F² values to correct for anisotropic diffraction before normalizing to E-values.

Parameters:

  • s_vectors (torch.Tensor) – Reciprocal space vectors in Angstroms^-1, shape (N, 3).

  • U (torch.Tensor) – Anisotropic parameters [u11, u22, u33, u12, u13, u23], shape (6,).

Returns:

correction – Correction factors, shape (N,).

Return type:

torch.Tensor

torchref.base.alignment.normalization.compute_shell_cv(F_squared, shell_idx, n_shells, min_count=10)[source]

Compute mean coefficient of variation of F² values within resolution shells.

After proper anisotropy correction, F² should have similar CV in all directions within each resolution shell.

Parameters:

  • F_squared (torch.Tensor) – F² values, shape (N,).

  • shell_idx (torch.Tensor) – Shell assignments, shape (N,).

  • n_shells (int) – Number of shells.

  • min_count (int) – Minimum reflections per shell to include in calculation.

Returns:

mean_cv – Mean coefficient of variation across shells.

Return type:

float

torchref.base.alignment.normalization.fit_anisotropy_correction(F_squared, s_vectors, n_shells=20, d_min=4.0, d_max=50.0, n_iterations=100, lr=0.01, verbose=True)[source]

Fit anisotropy correction parameters to minimize variance within shells.

Optimizes U parameters so that corrected F² values have minimal coefficient of variation within each resolution shell, making the distribution more isotropic before E-value conversion.

Uses PyTorch’s LBFGS optimizer for efficient optimization.

Parameters:

  • F_squared (torch.Tensor) – F² values, shape (N,).

  • s_vectors (torch.Tensor) – Reciprocal space vectors in Angstroms^-1, shape (N, 3).

  • n_shells (int) – Number of resolution shells for variance calculation.

  • d_min (float) – High resolution limit in Angstroms.

  • d_max (float) – Low resolution limit in Angstroms.

  • n_iterations (int) – Number of optimization iterations.

  • lr (float) – Learning rate for optimizer.

  • verbose (bool) – Print progress.

Returns:

  • U_optimal (torch.Tensor) – Optimal anisotropic parameters, shape (6,).

  • final_cv (float) – Final mean coefficient of variation after correction.

Return type:

Tuple[Tensor, float]

torchref.base.alignment.normalization.apply_anisotropy_correction(F_squared, s_vectors, U)[source]

Apply anisotropic correction to F² values.

Should be called BEFORE converting F² to E-values.

Parameters:

  • F_squared (torch.Tensor) – F² values, shape (N,).

  • s_vectors (torch.Tensor) – Reciprocal space vectors in Angstroms^-1, shape (N, 3).

  • U (torch.Tensor) – Anisotropic parameters, shape (6,).

Returns:

F2_corrected – Corrected F² values, shape (N,).

Return type:

torch.Tensor

torchref.base.alignment.normalization.F_squared_to_E_values(F_squared, s_vectors, n_shells=20, d_min=None, d_max=None)[source]

Convert F² values to E-values by normalizing within resolution shells.

E-values are normalized structure factors where <E²> = 1 within each resolution shell.

Parameters:

  • F_squared (torch.Tensor) – F² values, shape (N,).

  • s_vectors (torch.Tensor) – Reciprocal space vectors in Angstroms^-1, shape (N, 3).

  • n_shells (int) – Number of resolution shells for normalization.

  • d_min (float, optional) – High resolution limit in Angstroms. If None, derived from s_vectors.

  • d_max (float, optional) – Low resolution limit in Angstroms. If None, derived from s_vectors.

Returns:

  • E_values (torch.Tensor) – Normalized E-values, shape (N,).

  • E_squared (torch.Tensor) – E² values (for correlation calculations), shape (N,).

  • shell_idx (torch.Tensor) – Shell assignments for each reflection, shape (N,).

Return type:

Tuple[Tensor, Tensor, Tensor]