Synthyra/ESMFold2-Experimental-Cutoff2025

1import torch
2import torch.nn.functional as F
3
4from .esmfold2_affine3d import Affine3D
5
6
7def masked_mean(
8    mask: torch.Tensor,
9    value: torch.Tensor,
10    dim: int | None | tuple[int, ...] = None,
11    eps=1e-10,
12) -> torch.Tensor:
13    """Compute the mean of `value` where only positions where `mask == true` are
14    counted.
15    """
16    mask = mask.expand(*value.shape)
17    return torch.sum(mask * value, dim=dim) / (eps + torch.sum(mask, dim=dim))
18
19
20def _pae_bins(
21    max_bin: float = 31, num_bins: int = 64, device: torch.device = torch.device("cpu")
22):
23    bins = torch.linspace(0, max_bin, steps=(num_bins - 1), device=device)
24    step = max_bin / (num_bins - 2)
25    bin_centers = bins + step / 2
26    bin_centers = torch.cat(
27        [bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0
28    )
29    return bin_centers
30
31
32def _compute_pae_masks(mask: torch.Tensor):
33    square_mask = (mask.unsqueeze(-1) * mask.unsqueeze(-2)).bool()
34    return square_mask
35
36
37def compute_predicted_aligned_error(
38    logits: torch.Tensor,
39    aa_mask: torch.Tensor,
40    sequence_id: torch.Tensor | None = None,
41    max_bin: float = 31,
42) -> torch.Tensor:
43    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
44    square_mask = _compute_pae_masks(aa_mask)
45    min_v = torch.finfo(logits.dtype).min
46    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
47
48    return (probs * bins).sum(dim=-1)
49
50
51@torch.no_grad
52def compute_tm(logits: torch.Tensor, aa_mask: torch.Tensor, max_bin: float = 31.0):
53    square_mask = _compute_pae_masks(aa_mask)
54    seqlens = aa_mask.sum(-1, keepdim=True)
55    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
56    d0 = 1.24 * (seqlens.clamp_min(19) - 15) ** (1 / 3) - 1.8
57    f_d = 1.0 / (1 + (bins / d0.unsqueeze(-1)) ** 2)
58
59    min_v = torch.finfo(logits.dtype).min
60    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
61    # This is the sum over bins
62    ptm = (probs * f_d.unsqueeze(-2)).sum(dim=-1)
63    # This is the mean over residues j
64    ptm = masked_mean(square_mask, ptm, dim=-1)
65    # The we do a max over residues i
66    return ptm.max(dim=-1).values
67
68
69def tm_loss(
70    logits: torch.Tensor,
71    pred_affine: torch.Tensor,
72    targ_affine: torch.Tensor,
73    targ_mask: torch.Tensor,
74    tm_mask: torch.Tensor | None = None,
75    sequence_id: torch.Tensor | None = None,
76    max_bin: float = 31,
77):
78    pred = Affine3D.from_tensor(pred_affine)
79    targ = Affine3D.from_tensor(targ_affine)
80
81    def transform(affine: Affine3D):
82        pts = affine.trans[..., None, :, :]
83        return affine.invert()[..., None].apply(pts)
84
85    with torch.no_grad():
86        sq_diff = (transform(pred) - transform(targ)).square().sum(dim=-1)
87
88        num_bins = logits.shape[-1]
89        sq_bins = torch.linspace(
90            0, max_bin, num_bins - 1, device=logits.device
91        ).square()
92        # Gets the bin id by using a sum.
93        true_bins = (sq_diff[..., None] > sq_bins).sum(dim=-1).long()
94
95    errors = F.cross_entropy(logits.movedim(3, 1), true_bins, reduction="none")
96    square_mask = _compute_pae_masks(targ_mask)
97    loss = masked_mean(square_mask, errors, dim=(-1, -2))
98
99    if tm_mask is not None:
100        loss = masked_mean(tm_mask, loss, dim=None)
101    else:
102        loss = loss.mean()
103
104    return loss
105