File size: 2,170 Bytes
3a8e4d3 | 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 | #
# For licensing see accompanying LICENSE file.
# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
#
import torch
from torch import nn
def compute_aggregated_metric(logits, end=1.0):
"""Compute the metric from the logits.
Parameters
----------
logits : torch.Tensor
The logits of the metric
end : float
Max value of the metric, by default 1.0
Returns
-------
Tensor
The metric value
"""
num_bins = logits.shape[-1]
bin_width = end / num_bins
bounds = torch.arange(
start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
)
probs = nn.functional.softmax(logits, dim=-1)
plddt = torch.sum(
probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
dim=-1,
)
return plddt
class ConfidenceModule(nn.Module):
def __init__(
self,
hidden_size,
transformer_blocks,
num_plddt_bins=50,
):
super().__init__()
self.transformer_blocks = transformer_blocks
self.to_plddt_logits = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.LayerNorm(hidden_size),
nn.SiLU(),
nn.Linear(hidden_size, num_plddt_bins),
)
def forward(
self,
latent,
feats,
):
token_pe_pos = torch.cat(
[
feats["residue_index"].unsqueeze(-1).float(), # (B, M, 1)
feats["entity_id"].unsqueeze(-1).float(), # (B, M, 1)
feats["asym_id"].unsqueeze(-1).float(), # (B, M, 1)
feats["sym_id"].unsqueeze(-1).float(), # (B, M, 1)
],
dim=-1,
)
latent = self.transformer_blocks(
latents=latent,
c=None,
pos=token_pe_pos,
)
# Compute the pLDDT
plddt_logits = self.to_plddt_logits(latent)
# Compute the aggregated pLDDT
plddt = compute_aggregated_metric(plddt_logits)
return dict(
plddt=plddt,
plddt_logits=plddt_logits,
)
|