Spaces:
Running
Running
| from typing import TypeVar | |
| import numpy as np | |
| import torch | |
| from torch import Tensor | |
| from esm.utils import residue_constants as RC | |
| from esm.utils.structure.affine3d import Affine3D | |
| ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) | |
| def atom3_to_backbone_frames(bb_positions: torch.Tensor) -> Affine3D: | |
| N, CA, C = bb_positions.unbind(dim=-2) | |
| return Affine3D.from_graham_schmidt(C, CA, N) | |
| def index_by_atom_name( | |
| atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2 | |
| ) -> ArrayOrTensor: | |
| squeeze = False | |
| if isinstance(atom_names, str): | |
| atom_names = [atom_names] | |
| squeeze = True | |
| indices = [RC.atom_order[atom_name] for atom_name in atom_names] | |
| dim = dim % atom37.ndim | |
| index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim)) | |
| result = atom37[index] # type: ignore | |
| if squeeze: | |
| result = result.squeeze(dim) | |
| return result | |
| def get_protein_normalization_frame(coords: Tensor) -> Affine3D: | |
| """Given a set of coordinates for a protein, compute a single frame that can be used to normalize the coordinates. | |
| Specifically, we compute the average position of the N, CA, and C atoms use those 3 points to construct a frame | |
| using the Gram-Schmidt algorithm. The average CA position is used as the origin of the frame. | |
| Args: | |
| coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates | |
| Returns: | |
| Affine3D: tensor of Affine3D frame | |
| """ | |
| bb_coords = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2) | |
| coord_mask = torch.all( | |
| torch.all(torch.isfinite(bb_coords), dim=-1), | |
| dim=-1, | |
| ) | |
| average_position_per_n_ca_c = bb_coords.masked_fill( | |
| ~coord_mask[..., None, None], 0 | |
| ).sum(-3) / (coord_mask.sum(-1)[..., None, None] + 1e-8) | |
| frame = atom3_to_backbone_frames(average_position_per_n_ca_c.float()) | |
| return frame | |
| def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor: | |
| """Given a set of coordinates and a single frame, apply the frame to the coordinates. | |
| Args: | |
| coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates | |
| frame (Affine3D): Affine3D frame | |
| Returns: | |
| torch.FloatTensor: [L, 37, 3] tensor of transformed coordinates | |
| """ | |
| coords_trans_rot = frame[..., None, None].invert().apply(coords) | |
| # only transform coordinates with frame that have a valid rotation | |
| valid_frame = frame.trans.norm(dim=-1) > 0 | |
| is_inf = torch.isinf(coords) | |
| coords = coords_trans_rot.where(valid_frame[..., None, None, None], coords) | |
| coords.masked_fill_(is_inf, torch.inf) | |
| return coords | |
| def normalize_coordinates(coords: Tensor) -> Tensor: | |
| return apply_frame_to_coords(coords, get_protein_normalization_frame(coords)) | |