File size: 4,889 Bytes
74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 74b8b2a e0ac7d5 | 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 | """Predicted-aligned-error scores and training loss."""
from __future__ import annotations
import torch
import torch.nn.functional as F
from torch import Tensor
from .esmfold2_affine3d import Affine3D
_CPU_DEVICE = torch.device("cpu")
def _compute_pae_masks(mask: Tensor) -> Tensor:
# mask: (..., l), where l is the residue count.
residue_mask = mask.bool() # (..., l)
return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2) # (..., l, l)
def _pae_bins(
max_bin: float = 31,
num_bins: int = 64,
device: torch.device = _CPU_DEVICE,
) -> Tensor:
"""Return the representative distance for each PAE probability bin."""
boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device) # (n_bins - 1,)
width = max_bin / (num_bins - 2)
centers = boundaries + width / 2 # (n_bins - 1,)
overflow_center = centers[-1:] + width # (1,)
return torch.cat((centers, overflow_center)) # (n_bins,)
def _masked_probabilities(logits: Tensor, pair_mask: Tensor) -> Tensor:
# logits: (..., l, l, n_bins); pair_mask: (..., l, l).
masked_logits = logits.masked_fill(
~pair_mask.unsqueeze(-1),
torch.finfo(logits.dtype).min,
) # (..., l, l, n_bins)
return masked_logits.softmax(dim=-1) # (..., l, l, n_bins)
def masked_mean(
mask: Tensor,
value: Tensor,
dim: int | tuple[int, ...] | None = None,
eps: float = 1e-10,
) -> Tensor:
"""Average values over true entries of a broadcast-compatible mask."""
# value has arbitrary shape; reduced_shape removes the axes named by dim.
weights = mask.expand_as(value) # value.shape
weighted_sum = torch.sum(weights * value, dim=dim) # reduced_shape
weight_sum = torch.sum(weights, dim=dim) # reduced_shape
return weighted_sum / (weight_sum + eps) # reduced_shape
def compute_predicted_aligned_error(
logits: Tensor,
aa_mask: Tensor,
sequence_id: Tensor | None = None,
max_bin: float = 31,
) -> Tensor:
"""Convert PAE logits ``X`` with shape (..., l, l, n) to distances."""
del sequence_id
pair_mask = _compute_pae_masks(aa_mask) # (..., l, l)
probabilities = _masked_probabilities(logits, pair_mask) # (..., l, l, n_bins)
centers = _pae_bins(max_bin, logits.shape[-1], logits.device) # (n_bins,)
return torch.sum(probabilities * centers, dim=-1) # (..., l, l)
@torch.no_grad()
def compute_tm(logits: Tensor, aa_mask: Tensor, max_bin: float = 31.0) -> Tensor:
"""Estimate TM score from logits (..., l, l, n_bins) and residue mask (..., l)."""
pair_mask = _compute_pae_masks(aa_mask) # (..., l, l)
sequence_lengths = aa_mask.sum(dim=-1, keepdim=True) # (..., 1)
centers = _pae_bins(max_bin, logits.shape[-1], logits.device) # (n_bins,)
distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8 # (..., 1)
tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2) # (..., 1, n_bins)
probabilities = _masked_probabilities(logits, pair_mask) # (..., l, l, n_bins)
score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1) # (..., l, l)
score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1) # (..., l)
return score_per_anchor.max(dim=-1).values # (...,)
def _local_coordinates(frames: Affine3D) -> Tensor:
# frames.shape: (..., l); trans: (..., l, 3).
origins = frames.trans[..., None, :, :] # (..., 1, l, 3)
return frames.invert()[..., None].apply(origins) # (..., l, l, 3)
def tm_loss(
logits: Tensor,
pred_affine: Tensor,
targ_affine: Tensor,
targ_mask: Tensor,
tm_mask: Tensor | None = None,
sequence_id: Tensor | None = None,
max_bin: float = 31,
) -> Tensor:
"""Cross-entropy loss for logits (b, l, l, n_bins) and residue frames (b, l)."""
del sequence_id
predicted_frames = Affine3D.from_tensor(pred_affine) # frame shape: (b, l)
target_frames = Affine3D.from_tensor(targ_affine) # frame shape: (b, l)
with torch.no_grad():
squared_error = (
(_local_coordinates(predicted_frames) - _local_coordinates(target_frames))
.square()
.sum(dim=-1)
) # (b, l, l)
boundaries = torch.linspace(
0,
max_bin,
logits.shape[-1] - 1,
device=logits.device,
).square() # (n_bins - 1,)
target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long() # (b, l, l)
cross_entropy = F.cross_entropy(
logits.movedim(3, 1),
target_bins,
reduction="none",
) # (b, l, l)
pair_mask = _compute_pae_masks(targ_mask) # (b, l, l)
loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2)) # (b,)
if tm_mask is None:
return loss_per_sample.mean() # ()
return masked_mean(tm_mask, loss_per_sample) # ()
|