Spaces:
Runtime error
Runtime error
File size: 4,233 Bytes
753e275 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffab.modules.common.geometry import angstrom_to_nm, pairwise_dihedrals
from diffab.modules.common.layers import AngularEncoding
from diffab.utils.protein.constants import BBHeavyAtom, AA
class PairEmbedding(nn.Module):
def __init__(self, feat_dim, max_num_atoms, max_aa_types=22, max_relpos=32):
super().__init__()
self.max_num_atoms = max_num_atoms
self.max_aa_types = max_aa_types
self.max_relpos = max_relpos
self.aa_pair_embed = nn.Embedding(self.max_aa_types*self.max_aa_types, feat_dim)
self.relpos_embed = nn.Embedding(2*max_relpos+1, feat_dim)
self.aapair_to_distcoef = nn.Embedding(self.max_aa_types*self.max_aa_types, max_num_atoms*max_num_atoms)
nn.init.zeros_(self.aapair_to_distcoef.weight)
self.distance_embed = nn.Sequential(
nn.Linear(max_num_atoms*max_num_atoms, feat_dim), nn.ReLU(),
nn.Linear(feat_dim, feat_dim), nn.ReLU(),
)
self.dihedral_embed = AngularEncoding()
feat_dihed_dim = self.dihedral_embed.get_out_dim(2) # Phi and Psi
infeat_dim = feat_dim+feat_dim+feat_dim+feat_dihed_dim
self.out_mlp = nn.Sequential(
nn.Linear(infeat_dim, feat_dim), nn.ReLU(),
nn.Linear(feat_dim, feat_dim), nn.ReLU(),
nn.Linear(feat_dim, feat_dim),
)
def forward(self, aa, res_nb, chain_nb, pos_atoms, mask_atoms, structure_mask=None, sequence_mask=None):
"""
Args:
aa: (N, L).
res_nb: (N, L).
chain_nb: (N, L).
pos_atoms: (N, L, A, 3)
mask_atoms: (N, L, A)
structure_mask: (N, L)
sequence_mask: (N, L), mask out unknown amino acids to generate.
Returns:
(N, L, L, feat_dim)
"""
N, L = aa.size()
# Remove other atoms
pos_atoms = pos_atoms[:, :, :self.max_num_atoms]
mask_atoms = mask_atoms[:, :, :self.max_num_atoms]
mask_residue = mask_atoms[:, :, BBHeavyAtom.CA] # (N, L)
mask_pair = mask_residue[:, :, None] * mask_residue[:, None, :]
pair_structure_mask = structure_mask[:, :, None] * structure_mask[:, None, :] if structure_mask is not None else None
# Pair identities
if sequence_mask is not None:
# Avoid data leakage at training time
aa = torch.where(sequence_mask, aa, torch.full_like(aa, fill_value=AA.UNK))
aa_pair = aa[:,:,None]*self.max_aa_types + aa[:,None,:] # (N, L, L)
feat_aapair = self.aa_pair_embed(aa_pair)
# Relative sequential positions
same_chain = (chain_nb[:, :, None] == chain_nb[:, None, :])
relpos = torch.clamp(
res_nb[:,:,None] - res_nb[:,None,:],
min=-self.max_relpos, max=self.max_relpos,
) # (N, L, L)
feat_relpos = self.relpos_embed(relpos + self.max_relpos) * same_chain[:,:,:,None]
# Distances
d = angstrom_to_nm(torch.linalg.norm(
pos_atoms[:,:,None,:,None] - pos_atoms[:,None,:,None,:],
dim = -1, ord = 2,
)).reshape(N, L, L, -1) # (N, L, L, A*A)
c = F.softplus(self.aapair_to_distcoef(aa_pair)) # (N, L, L, A*A)
d_gauss = torch.exp(-1 * c * d**2)
mask_atom_pair = (mask_atoms[:,:,None,:,None] * mask_atoms[:,None,:,None,:]).reshape(N, L, L, -1)
feat_dist = self.distance_embed(d_gauss * mask_atom_pair)
if pair_structure_mask is not None:
# Avoid data leakage at training time
feat_dist = feat_dist * pair_structure_mask[:, :, :, None]
# Orientations
dihed = pairwise_dihedrals(pos_atoms) # (N, L, L, 2)
feat_dihed = self.dihedral_embed(dihed)
if pair_structure_mask is not None:
# Avoid data leakage at training time
feat_dihed = feat_dihed * pair_structure_mask[:, :, :, None]
# All
feat_all = torch.cat([feat_aapair, feat_relpos, feat_dist, feat_dihed], dim=-1)
feat_all = self.out_mlp(feat_all) # (N, L, L, F)
feat_all = feat_all * mask_pair[:, :, :, None]
return feat_all
|