Synthyra/ESMFold2

1from __future__ import annotations
2
3from dataclasses import Field, replace
4from typing import Any, ClassVar, Protocol, TypeVar
5
6import numpy as np
7import torch
8
9from .esmfold2_protein_structure import compute_affine_and_rmsd
10
11
12class Alignable(Protocol):
13    # Trick to detect whether an object is a dataclass
14    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
15
16    @property
17    def atom37_positions(self) -> np.ndarray:  # type: ignore
18        pass
19
20    @property
21    def atom37_mask(self) -> np.ndarray:  # type: ignore
22        pass
23
24    def __len__(self) -> int: ...
25
26
27T = TypeVar("T", bound=Alignable)
28
29
30class Aligner:
31    def __init__(
32        self,
33        mobile: Alignable,
34        target: Alignable,
35        only_use_backbone: bool = False,
36        use_reflection: bool = False,
37    ):
38        """
39        Aligns a mobile protein chain against a target protein chain.
40
41        Args:
42            mobile (ProteinChain): Protein chain to be aligned.
43            target (ProteinChain): Protein chain target.
44            only_use_backbone (bool): Whether to only use backbone atoms.
45            use_reflection (bool): Whether to align to target reflection.
46        """
47        # Check proteins must have same number of residues
48        assert len(mobile) == len(target)
49
50        # Determine overlapping atoms
51        joint_atom37_mask = mobile.atom37_mask.astype(bool) & target.atom37_mask.astype(
52            bool
53        )
54
55        # Backbone atoms are first sites in atom37 representation
56        if only_use_backbone:
57            joint_atom37_mask[:, 3:] = False
58
59        # Extract matching atom positions and convert to batched tensors
60        mobile_atom_tensor = (
61            torch.from_numpy(mobile.atom37_positions).type(torch.double).unsqueeze(0)
62        )
63        target_atom_tensor = (
64            torch.from_numpy(target.atom37_positions).type(torch.double).unsqueeze(0)
65        )
66        joint_atom37_mask = (
67            torch.from_numpy(joint_atom37_mask).type(torch.bool).unsqueeze(0)
68        )
69
70        # If using reflection flip target
71        if use_reflection:
72            target_atom_tensor = -target_atom_tensor
73
74        # Compute alignment and rmsd
75        affine3D, rmsd = compute_affine_and_rmsd(
76            mobile_atom_tensor, target_atom_tensor, atom_exists_mask=joint_atom37_mask
77        )
78        self._affine3D = affine3D
79        self._rmsd = rmsd.item()
80
81    @property
82    def rmsd(self):
83        return self._rmsd
84
85    def apply(self, mobile: T) -> T:
86        """Apply alignment to a protein chain"""
87        # Extract atom positions and convert to batched tensors
88        mobile_atom_tensor = (
89            torch.from_numpy(mobile.atom37_positions[mobile.atom37_mask])
90            .type(torch.float32)
91            .unsqueeze(0)
92        )
93
94        # Transform atom arrays
95        aligned_atom_tensor = self._affine3D.apply(mobile_atom_tensor).squeeze(0)
96
97        # Rebuild atom37 positions
98        aligned_atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
99        aligned_atom37_positions[mobile.atom37_mask] = aligned_atom_tensor
100
101        return replace(mobile, atom37_positions=aligned_atom37_positions)
102