File size: 7,595 Bytes
4691fb1 | 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 135 136 137 138 139 140 141 142 143 144 145 146 | """Shared RoPE stability decoder and symmetric pair residual."""
from __future__ import annotations
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
import torch
from torch import nn
from torch.nn import functional as F
AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
AMINO_ACID_TO_INDEX = {residue: index for index, residue in enumerate(AMINO_ACIDS)}
ESMC_EMBEDDING_DIM = 1152
@dataclass
class MutationTask:
wildtype_embedding: torch.Tensor
positions: torch.Tensor
target_residues: torch.Tensor
class RotarySelfAttention(nn.Module):
def __init__(self, model_dim: int, num_heads: int, dropout: float) -> None:
super().__init__()
self.num_heads = num_heads
self.head_dim = model_dim // num_heads
self.dropout = dropout
self.query_key_value = nn.Linear(model_dim, model_dim * 3)
self.output = nn.Linear(model_dim, model_dim)
def _rotate(self, values: torch.Tensor) -> torch.Tensor:
first, second = values.chunk(2, dim=-1)
return torch.cat([-second, first], dim=-1)
def _apply_rope(self, values: torch.Tensor) -> torch.Tensor:
length = values.size(-2)
positions = torch.arange(length, device=values.device, dtype=torch.float32)
frequencies = torch.exp(-math.log(10000) * torch.arange(0, self.head_dim, 2, device=values.device, dtype=torch.float32) / self.head_dim)
angles = positions[:, None] * frequencies[None, :]
cosine = torch.cat([angles.cos(), angles.cos()], dim=-1).to(values.dtype)[None, None, :, :]
sine = torch.cat([angles.sin(), angles.sin()], dim=-1).to(values.dtype)[None, None, :, :]
return values * cosine + self._rotate(values) * sine
def forward(self, tokens: torch.Tensor) -> torch.Tensor:
batch, length, _ = tokens.shape
query, key, value = self.query_key_value(tokens).chunk(3, dim=-1)
shape = (batch, length, self.num_heads, self.head_dim)
query = self._apply_rope(query.reshape(shape).transpose(1, 2))
key = self._apply_rope(key.reshape(shape).transpose(1, 2))
value = value.reshape(shape).transpose(1, 2)
attended = F.scaled_dot_product_attention(query, key, value, dropout_p=self.dropout if self.training else 0.0)
return self.output(attended.transpose(1, 2).reshape(batch, length, -1))
class RotaryTransformerEncoderLayer(nn.Module):
def __init__(self, model_dim: int, num_heads: int, dropout: float) -> None:
super().__init__()
self.attention_norm = nn.LayerNorm(model_dim)
self.attention = RotarySelfAttention(model_dim, num_heads, dropout)
self.feedforward_norm = nn.LayerNorm(model_dim)
self.feedforward = nn.Sequential(
nn.Linear(model_dim, model_dim * 4),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(model_dim * 4, model_dim),
)
self.dropout = nn.Dropout(dropout)
def forward(self, tokens: torch.Tensor) -> torch.Tensor:
tokens = tokens + self.dropout(self.attention(self.attention_norm(tokens)))
return tokens + self.dropout(self.feedforward(self.feedforward_norm(tokens)))
class StabilityBackbone(nn.Module):
def __init__(self, model_dim: int, depth: int, num_heads: int) -> None:
super().__init__()
self.input_proj = nn.Linear(ESMC_EMBEDDING_DIM, model_dim)
self.encoder = nn.ModuleList(RotaryTransformerEncoderLayer(model_dim, num_heads, 0.1) for _ in range(depth))
self.norm = nn.LayerNorm(model_dim)
self.score_head = nn.Linear(model_dim, len(AMINO_ACIDS))
def forward(self, embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
hidden = self.input_proj(embedding)
for layer in self.encoder:
hidden = layer(hidden)
hidden = self.norm(hidden)
return hidden, self.score_head(hidden)
class StabilityArc(nn.Module):
def __init__(self, model_dim: int = 256, depth: int = 4, num_heads: int = 8) -> None:
super().__init__()
self.backbone = StabilityBackbone(model_dim, depth, num_heads)
self.site_projection = nn.Sequential(nn.Linear(model_dim, model_dim), nn.GELU(), nn.Linear(model_dim, model_dim))
self.pair_residue = nn.Embedding(len(AMINO_ACIDS), model_dim)
self.pair_head = nn.Sequential(nn.Linear(model_dim * 3 + 2, model_dim), nn.GELU(), nn.Linear(model_dim, 1))
def score_task_components(self, task: MutationTask, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
hidden, scores = self.backbone(task.wildtype_embedding.to(device).unsqueeze(0))
site = self.site_projection(hidden.squeeze(0))
return scores.squeeze(0), site
def score_edit_sets(self, scores: torch.Tensor, site: torch.Tensor, contact_map: torch.Tensor, edit_sets: Sequence[Sequence[tuple[int, int]]], device: torch.device) -> torch.Tensor:
direct = []
pair_positions_left: list[int] = []
pair_positions_right: list[int] = []
pair_residues_left: list[int] = []
pair_residues_right: list[int] = []
pair_groups: list[int] = []
for group, edits in enumerate(edit_sets):
direct.append(scores[[position for position, _ in edits], [residue for _, residue in edits]].sum())
for left, (left_position, left_residue) in enumerate(edits):
for right_position, right_residue in edits[left + 1:]:
pair_positions_left.append(left_position)
pair_positions_right.append(right_position)
pair_residues_left.append(left_residue)
pair_residues_right.append(right_residue)
pair_groups.append(group)
predictions = torch.stack(direct)
if not pair_groups:
return predictions
left = site[torch.tensor(pair_positions_left, dtype=torch.long, device=device)] + self.pair_residue(torch.tensor(pair_residues_left, dtype=torch.long, device=device))
right = site[torch.tensor(pair_positions_right, dtype=torch.long, device=device)] + self.pair_residue(torch.tensor(pair_residues_right, dtype=torch.long, device=device))
distance = (torch.tensor(pair_positions_left, dtype=torch.float32, device=device) - torch.tensor(pair_positions_right, dtype=torch.float32, device=device)).abs().unsqueeze(-1) / max(site.size(0) - 1, 1)
contact = contact_map.to(device)[
torch.tensor(pair_positions_left, dtype=torch.long, device=device),
torch.tensor(pair_positions_right, dtype=torch.long, device=device),
].unsqueeze(-1)
features = torch.cat([left + right, (left - right).abs(), left * right, distance, contact], dim=-1)
residuals = self.pair_head(features).squeeze(-1)
return predictions + predictions.new_zeros(len(edit_sets)).scatter_add_(0, torch.tensor(pair_groups, dtype=torch.long, device=device), residuals)
def load_stabilityarc(checkpoint: Path, device: str = "cpu") -> tuple[StabilityArc, list[str]]:
payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
config = payload["config"]
if config["position_encoding"] != "rope" or config["pair_residual"] != "contact-aware":
raise ValueError("Checkpoint must be a direct RoPE stability scorer with a contact-aware pair residual")
model = StabilityArc(config["model_dim"], config["depth"], config["num_heads"])
model.load_state_dict(payload["state_dict"], strict=True)
return model.to(device).eval(), list(payload["training_proteins"]) |