Spaces:
Sleeping
Sleeping
File size: 2,246 Bytes
b72d311 | 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 | """
ProtStab CNN — protein thermodynamic stability predictor.
Input : one-hot encoded protein sequence (VOCAB_SIZE × MAX_LEN)
Output: predicted ΔG (kcal/mol) — positive = stable, negative = unstable
"""
import torch
import torch.nn as nn
AA_VOCAB = list("ACDEFGHIKLMNPQRSTVWYX") # X = unknown / non-standard
AA_TO_IDX = {aa: i for i, aa in enumerate(AA_VOCAB)}
VOCAB_SIZE = len(AA_VOCAB) # 21
MAX_LEN = 256 # sequences truncated/padded to this length
def encode_sequence(seq: str) -> torch.Tensor:
"""Return one-hot tensor of shape (VOCAB_SIZE, MAX_LEN)."""
seq = seq.upper().strip()[:MAX_LEN]
one_hot = torch.zeros(VOCAB_SIZE, MAX_LEN)
for i, aa in enumerate(seq):
idx = AA_TO_IDX.get(aa, AA_TO_IDX["X"])
one_hot[idx, i] = 1.0
return one_hot
class ConvBlock(nn.Module):
def __init__(self, in_ch: int, out_ch: int, kernel: int):
super().__init__()
self.net = nn.Sequential(
nn.Conv1d(in_ch, out_ch, kernel_size=kernel, padding=kernel // 2),
nn.BatchNorm1d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.net(x)
class ProtStabCNN(nn.Module):
"""
3-layer 1D CNN with global average pooling → MLP head.
~500k parameters, pre-trained on DMSv4 (455k sequences).
"""
def __init__(self, dropout: float = 0.3):
super().__init__()
self.encoder = nn.Sequential(
ConvBlock(VOCAB_SIZE, 64, kernel=5),
ConvBlock(64, 128, kernel=5),
ConvBlock(128, 256, kernel=3),
)
self.head = nn.Sequential(
nn.Linear(256, 128),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(128, 32),
nn.ReLU(inplace=True),
nn.Linear(32, 1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.encoder(x) # (B, 256, MAX_LEN)
x = x.mean(dim=2) # global average pool → (B, 256)
return self.head(x).squeeze(-1) # (B,)
def count_parameters(self) -> int:
return sum(p.numel() for p in self.parameters() if p.requires_grad)
|