torchref.base.alignment package

Structure alignment and superposition functions.

This submodule provides functions for: - Superposition of coordinate sets (Kabsch algorithm) - Rotation operations and conversions - PDB alignment utilities

torchref.base.alignment.superpose_vectors_robust_torch(ref_coords, mov_coords, weights=None, max_iterations=10)

Perform weighted superposition of two coordinate sets using SVD (PyTorch version).

Parameters:

• ref_coords (torch.Tensor) – Reference coordinates of shape (N, 3).

• mov_coords (torch.Tensor) – Mobile coordinates of shape (N, 3) to be superposed onto reference.

• weights (torch.Tensor, optional) – Weights for each atom of shape (N, 1). Default is uniform weights.

• max_iterations (int, optional) – Maximum number of iterations for refinement. Default is 10.

Returns:

4x4 transformation matrix (shape (3, 4) returned).

Return type:

torch.Tensor

torchref.base.alignment.superpose_vectors_robust(target_coords, mobile_coords, weights=None, max_iterations=1)

Superpose mobile coordinates onto target coordinates using the Kabsch algorithm.

Computes the optimal rotation and translation to minimize the weighted RMSD between two sets of 3D coordinates, with robust handling of special cases such as reflection.

Parameters:

• target_coords (numpy.ndarray) – Target coordinates with shape (N, 3).

• mobile_coords (numpy.ndarray) – Mobile coordinates with shape (N, 3) to be superposed onto target.

• weights (numpy.ndarray, optional) – Per-atom weights for the superposition with shape (N,). Default is uniform weights.

• max_iterations (int, optional) – Number of iterations for refinement. Default is 1 (standard Kabsch).

Returns:

• transformation_matrix (numpy.ndarray) – 4x4 transformation matrix that maps mobile_coords onto target_coords.

• rmsd (float) – Weighted root-mean-square deviation after superposition.

Raises:

ValueError – If input coordinate arrays have different shapes.

Notes

The algorithm uses SVD decomposition of the covariance matrix and handles the reflection case by checking the determinant of the rotation matrix.

torchref.base.alignment.align_torch(xyz1, xyz2, idx_to_move=None)

Align two coordinate sets using superposition (PyTorch version).

Parameters:

• xyz1 (torch.Tensor) – Target coordinates of shape (N, 3).

• xyz2 (torch.Tensor) – Coordinates to be aligned of shape (N, 3).

• idx_to_move (torch.Tensor, optional) – Indices of atoms to use for alignment. If None, uses all atoms.

Returns:

Aligned coordinates of shape (N, 3).

Return type:

torch.Tensor

torchref.base.alignment.get_alignement_matrix(xyz1, xyz2, idx_to_move=None)

Get the alignment transformation matrix between two coordinate sets.

Parameters:

• xyz1 (torch.Tensor) – Target coordinates of shape (N, 3).

• xyz2 (torch.Tensor) – Coordinates to be aligned of shape (N, 3).

• idx_to_move (torch.Tensor, optional) – Indices of atoms to use for alignment. If None, uses all atoms.

Returns:

Transformation matrix of shape (3, 4).

Return type:

torch.Tensor

torchref.base.alignment.align_pdbs(pdb1, pdb2, Atoms=None)

Align two PDB structures using the Kabsch algorithm.

Superimposes pdb2 onto pdb1 by minimizing the RMSD between corresponding atoms. The transformation is applied in-place to pdb2.

Parameters:

• pdb1 (pandas.DataFrame) – Reference PDB structure with ‘x’, ‘y’, ‘z’, ‘name’, and ‘tempfactor’ columns.

• pdb2 (pandas.DataFrame) – Mobile PDB structure to be aligned onto pdb1.

• Atoms (list, optional) – List of atom names to use for alignment. If None, all atoms are used.

Returns:

• pdb2 (pandas.DataFrame) – Transformed pdb2 with updated coordinates.

• rmsd (float) – Root-mean-square deviation after alignment.

torchref.base.alignment.get_alignment_matrix(pdb1, pdb2, Atoms=None)

Calculate the transformation matrix to align two PDB structures.

Computes the 4x4 transformation matrix that would superimpose pdb2 onto pdb1 without actually applying the transformation.

Parameters:

• pdb1 (pandas.DataFrame) – Reference PDB structure with ‘x’, ‘y’, ‘z’, ‘name’, and ‘tempfactor’ columns.

• pdb2 (pandas.DataFrame) – Mobile PDB structure.

• Atoms (list, optional) – List of atom names to use for alignment. If None, all atoms are used.

Returns:

• transformation_matrix (numpy.ndarray) – 4x4 transformation matrix.

• rmsd (float) – Root-mean-square deviation that would result from the alignment.

torchref.base.alignment.apply_transformation(points, transformation_matrix)

Apply a 4x4 transformation matrix to 3D points (PyTorch version).

Parameters:

• points (torch.Tensor) – 3D points of shape (N, 3).

• transformation_matrix (torch.Tensor) – Transformation matrix of shape (3, 4) or (4, 4).

Returns:

Transformed 3D points of shape (N, 3).

Return type:

torch.Tensor

torchref.base.alignment.apply_transformation_numpy(points, transformation_matrix)

Apply a 4x4 transformation matrix to 3D points (NumPy version).

Converts points to homogeneous coordinates, applies the transformation, and returns the transformed 3D coordinates.

Parameters:

• points (numpy.ndarray) – 3D coordinates with shape (N, 3).

• transformation_matrix (numpy.ndarray) – 4x4 transformation matrix containing rotation and translation.

Returns:

Transformed 3D coordinates with shape (N, 3).

Return type:

numpy.ndarray

torchref.base.alignment.invert_transformation_matrix(transformation_matrix)

Compute the inverse of a 4x4 transformation matrix.

Efficiently inverts a rigid-body transformation matrix by transposing the rotation component and computing the inverse translation.

Parameters:

transformation_matrix (numpy.ndarray) – 4x4 transformation matrix containing rotation (top-left 3x3) and translation (top-right 3x1).

Returns:

Inverse 4x4 transformation matrix.

Return type:

numpy.ndarray

Notes

This function assumes the input is a valid rigid-body transformation (rotation + translation). For such matrices, the inverse rotation is simply the transpose, and the inverse translation is computed as -R^T @ t.

torchref.base.alignment.rotate_coords_torch(coords, phi, rho)

Rotate coordinates using phi and rho angles (PyTorch version).

Parameters:

• coords (torch.Tensor) – Coordinates of shape (N, 3) to rotate.

• phi (float) – Rotation angle phi in degrees.

• rho (float) – Rotation angle rho in degrees.

Returns:

Rotated coordinates of shape (N, 3).

Return type:

torch.Tensor

torchref.base.alignment.rotate_coords_numpy(coords, phi, rho)

Rotate 3D coordinates by phi and rho angles (NumPy version).

Applies a rotation transformation to a set of 3D coordinates using two rotation angles (phi and rho) in degrees.

Parameters:

• coords (numpy.ndarray) – Array of 3D coordinates with shape (N, 3).

• phi (float) – First rotation angle in degrees.

• rho (float) – Second rotation angle in degrees.

Returns:

Rotated coordinates with the same shape as input (N, 3).

Return type:

numpy.ndarray

torchref.base.alignment.axis_angle_to_rotation_matrix(axis_angle)

Convert axis-angle representation to 3x3 rotation matrix.

Uses Rodrigues’ formula. Supports batched input and gradients.

Parameters:

axis_angle (torch.Tensor) – Axis-angle representation with shape (3,) or (N, 3). Direction is rotation axis, magnitude is angle in radians.

Returns:

Rotation matrix with shape (3, 3) or (N, 3, 3).

Return type:

torch.Tensor

torchref.base.alignment.rotation_matrix_to_axis_angle(R)

Convert 3x3 rotation matrix to axis-angle representation.

Parameters:

R (torch.Tensor) – Rotation matrix with shape (3, 3) or (N, 3, 3).

Returns:

Axis-angle representation with shape (3,) or (N, 3).

Return type:

torch.Tensor

torchref.base.alignment.quaternion_to_rotation_matrix(q)

Convert quaternion to rotation matrix.

Parameters:

q (torch.Tensor) – Quaternion with shape (4,) or (N, 4). Format: [w, x, y, z].

Returns:

Rotation matrix with shape (3, 3) or (N, 3, 3).

Return type:

torch.Tensor

torchref.base.alignment.random_rotation_uniform(n=1, device=device(type='cpu'), dtype=torch.float32)

Generate uniform random rotations over SO(3).

Uses Shoemake’s quaternion-based uniform sampling.

Parameters:

• n (int, optional) – Number of rotations to generate. Default is 1.

• device (str, optional) – Device for output tensor. Defaults to the configured device.current.

• dtype (torch.dtype, optional) – Data type for output tensor. Default is dtypes.float.

Returns:

Rotation matrices with shape (n, 3, 3) or (3, 3) if n=1.

Return type:

torch.Tensor

torchref.base.alignment.rotation_matrix_euler_zyz(angles)

Create rotation matrix from ZYZ Euler angles (differentiable PyTorch version).

R = Rz(alpha) @ Ry(beta) @ Rz(gamma)

Parameters:

angles (torch.Tensor) – Tensor of three rotation angles (alpha, beta, gamma) in radians. Or shape (B, 3) for batched input. The function will return (B, 3, 3) in that case.

Returns:

3x3 rotation matrix.

Return type:

torch.Tensor

torchref.base.alignment.compute_radial_shells(d_min, d_max, n_shells, device=None)

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.assign_to_shells(s_mag, shell_edges)

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.compute_anisotropy_correction(s_vectors, U)

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.compute_shell_cv(F_squared, shell_idx, n_shells, min_count=10)

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.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)

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.apply_anisotropy_correction(F_squared, s_vectors, U)

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.F_squared_to_E_values(F_squared, s_vectors, n_shells=20, d_min=None, d_max=None)

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]

Submodules

• torchref.base.alignment.normalization module
  • compute_radial_shells()
  • assign_to_shells()
  • compute_anisotropy_correction()
  • compute_shell_cv()
  • fit_anisotropy_correction()
  • apply_anisotropy_correction()
  • F_squared_to_E_values()
• torchref.base.alignment.rotation module
  • rotate_coords_torch()
  • rotate_coords_numpy()
  • axis_angle_to_rotation_matrix()
  • rotation_matrix_to_axis_angle()
  • quaternion_to_rotation_matrix()
  • random_rotation_uniform()
  • rotation_matrix_euler_zyz()
• torchref.base.alignment.scoring module
  • binned_correlation()
• torchref.base.alignment.superposition module
  • superpose_vectors_robust_torch()
  • superpose_vectors_robust()
  • align_torch()
  • get_alignement_matrix()
  • align_pdbs()
  • get_alignment_matrix()
  • apply_transformation()
  • apply_transformation_numpy()
  • invert_transformation_matrix()