torchref.base.alignment

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)

Weighted SVD (Kabsch) superposition of two coordinate sets (PyTorch).

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) – Weighted Kabsch steps, best weighted RMSD wins. Weights are never reweighted, so past the first step this normally changes nothing.

Returns:

Transformation matrix of shape (3, 4): rotation block plus translation column. Unlike the NumPy sibling superpose_vectors_robust(), no RMSD is returned.

Return type:

torch.Tensor

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

Superpose mobile onto target coordinates with the weighted Kabsch algorithm.

Reflections are rejected by determinant correction, so the result is always a proper rotation.

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 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 matrix mapping mobile_coords onto target_coords. If the SVD fails this degrades to the identity (printed, not raised).

• rmsd (float) – Weighted RMSD after superposition – unweighted on the SVD-failure path.

Raises:

ValueError – If input coordinate arrays have different shapes.

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 pdb2 onto pdb1 by weighted Kabsch, rewriting pdb2’s coordinates in place.

Weighting is 1 / tempfactor, so every atom used must have a nonzero ‘tempfactor’ or the call divides by zero.

Parameters:

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

• pdb2 (pandas.DataFrame) – Mobile structure; its coordinate columns are overwritten.

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

Returns:

• pdb2 (pandas.DataFrame) – The same object, with updated coordinates.

• rmsd (float) – Unweighted all-atom RMSD after alignment.

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

Transformation matrix that would superimpose pdb2 onto pdb1, without applying it.

Same 1 / tempfactor weighting as align_pdbs(), so ‘tempfactor’ must be nonzero for every atom used.

Parameters:

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

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

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

Returns:

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

• rmsd (float) – Weighted RMSD 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).

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)

Invert a rigid-body 4x4 transformation matrix.

Uses R^-1 = R^T, so a matrix that is not a pure rotation plus translation (any scale or shear) is inverted silently wrongly.

Parameters:

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

Returns:

Inverse 4x4 transformation matrix.

Return type:

numpy.ndarray

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

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

Convention: rotate by phi about z, then by rho about the resulting x axis.

Parameters:

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

• phi (torch.Tensor) – Angle in degrees; must be a scalar tensor (torch.cos is applied to it), so a Python float is rejected.

• rho (torch.Tensor) – Angle in degrees (scalar tensor).

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

Same convention as rotate_coords_torch(); the matrix is built in float64 regardless of the input dtype.

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=None, dtype=None)

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 or torch.device, optional) – Device for output tensor. Defaults to the configured default device.

• 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:

Rotation matrix of shape (3, 3), or (B, 3, 3) for batched input.

Return type:

torch.Tensor

torchref.base.alignment.rotation_matrix_euler_xyz(angles)

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

R = Rz(gamma) @ Ry(beta) @ Rx(alpha). Distinct world axes, so unlike ZYZ there is no gimbal-lock singularity at beta=0.

Parameters:

angles (torch.Tensor) – Tensor of three rotation angles (alpha, beta, gamma) in radians, applied as Rx(alpha), Ry(beta), Rz(gamma) — outer product Rz·Ry·Rx. Shape (3,) or (B, 3); returns (3, 3) or (B, 3, 3) respectively.

Returns:

3x3 rotation matrix (or batched (B, 3, 3)).

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 U so that corrected F² has minimal CV within each resolution shell.

Flattens the anisotropic falloff before E-value conversion.

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) – Resolution limits in Angstroms; reflections outside are dropped.

• d_max (float) – Resolution limits in Angstroms; reflections outside are dropped.

• n_iterations (int) – Approximate total inner steps: LBFGS is stepped n_iterations // 20 times at max_iter=20, so values below 20 optimize nothing at all.

• lr (float) – LBFGS step scale – not an SGD learning rate.

• 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,). sqrt(E²), so non-negative and unsigned – never treat these as signed normalized amplitudes.

• 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]

Modules

normalization	Normalization functions for E-value conversion and anisotropy correction.
rotation	Rotation utility functions.
scoring	Correlation scoring utilities for reciprocal-space alignment.
superposition	Superposition functions for coordinate alignment.