File size: 4,220 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
"""Alignment based metrics."""
import numpy as np
def transform_ls(
x: np.ndarray,
b: np.ndarray,
*,
allow_reflection: bool = False,
) -> np.ndarray:
"""Find the least squares best fit rotation between two sets of N points.
Solve Ax = b for A. Where A is the transform rotating x^T into b^T.
Args:
x: NxD numpy array of coordinates. Usually dimension D is 3.
b: NxD numpy array of coordinates. Usually dimension D is 3.
allow_reflection: Whether the returned transformation can reflect as well as
rotate.
Returns:
Matrix A transforming x into b, i.e. s.t. Ax^T = b^T.
"""
assert x.shape[1] >= b.shape[1]
assert b.shape[0] == x.shape[0], '%d, %d' % (b.shape[0], x.shape[0])
# First postmultiply by x.;
# Axx^t = b x^t
bxt = np.dot(b.transpose(), x) / b.shape[0]
u, _, v = np.linalg.svd(bxt)
r = np.dot(u, v)
if not allow_reflection:
flip = np.ones((v.shape[1], 1))
flip[v.shape[1] - 1, 0] = np.sign(np.linalg.det(r))
r = np.dot(u, v * flip)
return r
def align(
*,
x: np.ndarray,
y: np.ndarray,
x_indices: np.ndarray,
y_indices: np.ndarray,
) -> np.ndarray:
"""Align x to y considering only included_idxs.
Args:
x: NxD np array of coordinates.
y: NxD np array of coordinates.
x_indices: An np array of indices for `x` that will be used in the
alignment. Must be of the same length as `y_included_idxs`.
y_indices: An np array of indices for `y` that will be used in the
alignment. Must be of the same length as `x_included_idxs`.
Returns:
NxD np array of points obtained by applying a rigid transformation to x.
These points are aligned to y and the alignment is the optimal alignment
over the points in included_idxs.
Raises:
ValueError: If the number of included indices is not the same for both
input arrays.
"""
if len(x_indices) != len(y_indices):
raise ValueError(
'Number of included indices must be the same for both input arrays,'
f' but got for x: {len(x_indices)}, and for y: {len(y_indices)}.'
)
x_mean = np.mean(x[x_indices, :], axis=0)
y_mean = np.mean(y[y_indices, :], axis=0)
centered_x = x - x_mean
centered_y = y - y_mean
t = transform_ls(centered_x[x_indices, :], centered_y[y_indices, :])
transformed_x = np.dot(centered_x, t.transpose()) + y_mean
return transformed_x
def deviations_from_coords(
decoy_coords: np.ndarray,
gt_coords: np.ndarray,
align_idxs: np.ndarray | None = None,
include_idxs: np.ndarray | None = None,
) -> np.ndarray:
"""Returns the raw per-atom deviations used in RMSD computation."""
if decoy_coords.shape != gt_coords.shape:
raise ValueError(
'decoy_coords.shape and gt_coords.shape must match.Found: %s and %s.'
% (decoy_coords.shape, gt_coords.shape)
)
# Include and align all residues unless specified otherwise.
if include_idxs is None:
include_idxs = np.arange(decoy_coords.shape[0])
if align_idxs is None:
align_idxs = include_idxs
aligned_decoy_coords = align(
x=decoy_coords,
y=gt_coords,
x_indices=align_idxs,
y_indices=align_idxs,
)
deviations = np.linalg.norm(
aligned_decoy_coords[include_idxs] - gt_coords[include_idxs], axis=1
)
return deviations
def rmsd_from_coords(
decoy_coords: np.ndarray,
gt_coords: np.ndarray,
align_idxs: np.ndarray | None = None,
include_idxs: np.ndarray | None = None,
) -> float:
"""Computes the *aligned* RMSD of two Mx3 np arrays of coordinates.
Args:
decoy_coords: [M, 3] np array of decoy atom coordinates.
gt_coords: [M, 3] np array of gt atom coordinates.
align_idxs: [M] np array of indices specifying coordinates to align on.
Defaults to None, in which case all the include_idx (see after) are used.
include_idxs: [M] np array of indices specifying coordinates to score.
Defaults to None, in which case all indices are used for scoring.
Returns:
rmsd value of the aligned decoy and gt coordinates.
"""
deviations = deviations_from_coords(
decoy_coords, gt_coords, align_idxs, include_idxs
)
return np.sqrt(np.mean(np.square(deviations)))
|