import torch
import torch.nn as nn
import torch.distributions as D
import sys
import numpy as np
sys.path.append("/databases/mpnn/github_repo/sc_utils")
import openfold.np.residue_constants as rc
from openfold.utils.rigid_utils import Rigid
from openfold.utils import feats
from openfold.data.data_transforms import make_atom14_masks, atom37_to_torsion_angles
from openfold.np.residue_constants import (
restype_rigid_group_default_frame,
restype_atom14_to_rigid_group,
restype_atom14_mask,
restype_atom14_rigid_group_positions,
restype_atom14_to_rigid_group,
)
from openfold.np import residue_constants
from openfold.utils.loss import between_residue_clash_loss
torch_pi = torch.tensor(np.pi, device="cpu")
map_mpnn_to_af2_seq = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]], device="cpu")
def pack_side_chains(feature_dict, model_sc, num_denoising_steps, num_samples=10, repack_everything=True):
device = feature_dict["X"].device
torsion_dict = make_torsion_features(feature_dict, repack_everything)
feature_dict["X"] = torsion_dict["xyz14_noised"]
feature_dict["X_m"] = torsion_dict["xyz14_m"]
if "Y" not in list(feature_dict):
feature_dict["Y"] = torch.zeros([feature_dict["X"].shape[0], feature_dict["X"].shape[1],16,3], device=device)
feature_dict["Y_t"] = torch.zeros([feature_dict["X"].shape[0], feature_dict["X"].shape[1],16], device=device)
feature_dict["Y_m"] = torch.zeros([feature_dict["X"].shape[0], feature_dict["X"].shape[1],16], device=device)
h_V, h_E, E_idx = model_sc.encode(feature_dict)
feature_dict["h_V"] = h_V
feature_dict["h_E"] = h_E
feature_dict["E_idx"] = E_idx
for step in range(num_denoising_steps):
mean, concentration, mix_logits = model_sc.decode(feature_dict)
mix = D.Categorical(logits=mix_logits)
comp = D.VonMises(mean, concentration)
pred_dist = D.MixtureSameFamily(mix, comp)
predicted_samples = pred_dist.sample([num_samples])
log_probs_of_samples = pred_dist.log_prob(predicted_samples)
sample = torch.gather(predicted_samples, dim=0, index=torch.argmax(log_probs_of_samples,0)[None,])[0,]
torsions_pred_unit = torch.cat([torch.sin(sample[:,:,:,None]), torch.cos(sample[:,:,:,None])], -1)
torsion_dict["torsions_noised"][:,:,3:] = torsions_pred_unit*torsion_dict["mask_fix_sc"] + torsion_dict["torsions_true"]*(1-torsion_dict["mask_fix_sc"])
pred_frames = feats.torsion_angles_to_frames(torsion_dict["rigids"],
torsion_dict["torsions_noised"],
torsion_dict["aatype"],
torch.tensor(restype_rigid_group_default_frame, device=device))
xyz14_noised = feats.frames_and_literature_positions_to_atom14_pos(pred_frames,
torsion_dict["aatype"],
torch.tensor(restype_rigid_group_default_frame, device=device),
torch.tensor(restype_atom14_to_rigid_group, device=device),
torch.tensor(restype_atom14_mask, device=device),
torch.tensor(restype_atom14_rigid_group_positions, device=device))
xyz14_noised = xyz14_noised*feature_dict["X_m"][:,:,:,None]
feature_dict["X"] = xyz14_noised
S_af2 = torsion_dict["S_af2"]
feature_dict["X"] = xyz14_noised
log_prob = pred_dist.log_prob(sample)*torsion_dict["mask_fix_sc"][...,0] + 2.0*(1-torsion_dict["mask_fix_sc"][...,0])
tmp_types = torch.tensor(restype_atom14_to_rigid_group, device=device)[S_af2]
tmp_types[tmp_types<4]=4
tmp_types -= 4
atom_types_for_b_factor = torch.nn.functional.one_hot(tmp_types,4)
uncertainty = log_prob[:,:,None,:]*atom_types_for_b_factor
b_factor_pred = uncertainty.sum(-1)
feature_dict["b_factors"] = b_factor_pred
feature_dict["mean"] = mean
feature_dict["concentration"] = concentration
feature_dict["mix_logits"] = mix_logits
feature_dict["log_prob"] = log_prob
feature_dict["sample"] = sample
feature_dict["true_torsion_sin_cos"] = torsion_dict["torsions_true"]
return feature_dict
def make_torsion_features(feature_dict, repack_everything=True):
device = feature_dict["mask"].device
mask = feature_dict["mask"]
B, L = mask.shape
xyz37 = torch.zeros([B, L, 37, 3], device=device, dtype=torch.float32)
xyz37[:,:,:3] = feature_dict["X"][:,:,:3]
xyz37[:,:,4] = feature_dict["X"][:,:,3]
S_af2 = torch.argmax(torch.nn.functional.one_hot(feature_dict["S"], 21).float() @ map_mpnn_to_af2_seq.to(device).float(), -1)
masks14_37 = make_atom14_masks({"aatype": S_af2})
temp_dict = {"aatype": S_af2,
"all_atom_positions": xyz37,
"all_atom_mask": masks14_37['atom37_atom_exists']}
torsion_dict = atom37_to_torsion_angles("")(temp_dict)
rigids = Rigid.make_transform_from_reference(
n_xyz=xyz37[:,:,0,:],
ca_xyz=xyz37[:,:,1,:],
c_xyz=xyz37[:,:,2,:],
eps=1e-9)
if not repack_everything:
xyz37_true = feature_dict["xyz_37"]
temp_dict_true = {"aatype": S_af2,
"all_atom_positions": xyz37_true,
"all_atom_mask": masks14_37['atom37_atom_exists']}
torsion_dict_true = atom37_to_torsion_angles("")(temp_dict_true)
torsions_true = torch.clone(torsion_dict_true['torsion_angles_sin_cos'])[:,:,3:]
mask_fix_sc = feature_dict["chain_mask"][:,:,None,None]
else:
torsions_true = torch.zeros([B, L, 4, 2], device=device)
mask_fix_sc = torch.ones([B,L,1,1], device=device)
random_angle = 2*torch_pi*torch.rand([S_af2.shape[0], S_af2.shape[1], 4], device=device)
random_sin_cos = torch.cat([torch.sin(random_angle)[...,None], torch.cos(random_angle)[...,None]], -1)
torsions_noised = torch.clone(torsion_dict['torsion_angles_sin_cos'])
torsions_noised[:,:,3:] = random_sin_cos*mask_fix_sc + torsions_true*(1-mask_fix_sc)
pred_frames = feats.torsion_angles_to_frames(rigids,
torsions_noised,
S_af2,
torch.tensor(restype_rigid_group_default_frame,
device=device))
xyz14_noised = feats.frames_and_literature_positions_to_atom14_pos(pred_frames,
S_af2,
torch.tensor(restype_rigid_group_default_frame, device=device),
torch.tensor(restype_atom14_to_rigid_group, device=device),
torch.tensor(restype_atom14_mask, device=device),
torch.tensor(restype_atom14_rigid_group_positions, device=device))
xyz14_m = masks14_37['atom14_atom_exists']*mask[:,:,None]
xyz14_noised = xyz14_noised*xyz14_m[:,:,:,None]
torsion_dict["xyz14_m"] = xyz14_m
torsion_dict["xyz14_noised"] = xyz14_noised
torsion_dict["mask_for_loss"] = mask
torsion_dict["rigids"] = rigids
torsion_dict["torsions_noised"] = torsions_noised
torsion_dict["mask_fix_sc"] = mask_fix_sc
torsion_dict["torsions_true"] = torsions_true
torsion_dict["S_af2"] = S_af2
return torsion_dict
class Packer(nn.Module):
def __init__(self,
edge_features=128,
node_features=128,
num_positional_embeddings=16,
num_chain_embeddings=16,
num_rbf=16,
top_k=30,
augment_eps=0.,
atom37_order=False,
device=None,
atom_context_num=16,
lower_bound=0.0,
upper_bound=20.0,
hidden_dim=128,
num_encoder_layers=3,
num_decoder_layers=3,
dropout=0.1,
num_mix=3):
super(Packer, self).__init__()
self.edge_features = edge_features
self.node_features = node_features
self.num_positional_embeddings = num_positional_embeddings
self.num_chain_embeddings = num_chain_embeddings
self.num_rbf = num_rbf
self.top_k = top_k
self.augment_eps = augment_eps
self.atom37_order = atom37_order
self.device = device
self.atom_context_num = atom_context_num
self.lower_bound = lower_bound
self.upper_bound = upper_bound
self.hidden_dim = hidden_dim
self.num_encoder_layers = num_encoder_layers
self.num_decoder_layers = num_decoder_layers
self.dropout = dropout
self.softplus = nn.Softplus(beta=1, threshold=20)
self.features = ProteinFeatures(
edge_features=edge_features,
node_features=node_features,
num_positional_embeddings=num_positional_embeddings,
num_chain_embeddings=num_chain_embeddings,
num_rbf=num_rbf,
top_k=top_k,
augment_eps=augment_eps,
atom37_order=atom37_order,
device=device,
atom_context_num=atom_context_num,
lower_bound=lower_bound,
upper_bound=upper_bound)
self.W_e = nn.Linear(edge_features, hidden_dim, bias=True)
self.W_v = nn.Linear(node_features, hidden_dim, bias=True)
self.W_f = nn.Linear(edge_features, hidden_dim, bias=True)
self.W_v_sc = nn.Linear(node_features, hidden_dim, bias=True)
self.linear_down = nn.Linear(2*hidden_dim, hidden_dim, bias=True)
self.W_torsions = nn.Linear(hidden_dim, 4*3*num_mix, bias=True)
self.num_mix = num_mix
self.dropout = nn.Dropout(dropout)
self.encoder_layers = nn.ModuleList([
EncLayer(hidden_dim, hidden_dim*2, dropout=dropout)
for _ in range(num_encoder_layers)
])
self.W_c = nn.Linear(hidden_dim, hidden_dim, bias=True)
self.W_e_context = nn.Linear(hidden_dim, hidden_dim, bias=True)
self.W_nodes_y = nn.Linear(hidden_dim, hidden_dim, bias=True)
self.W_edges_y = nn.Linear(hidden_dim, hidden_dim, bias=True)
self.context_encoder_layers = nn.ModuleList([
DecLayer(hidden_dim, hidden_dim*2, dropout=dropout)
for _ in range(2)
])
self.V_C = nn.Linear(hidden_dim, hidden_dim, bias=False)
self.V_C_norm = nn.LayerNorm(hidden_dim)
self.y_context_encoder_layers = nn.ModuleList([
DecLayerJ(hidden_dim, hidden_dim, dropout=dropout)
for _ in range(2)
])
self.h_V_C_dropout = nn.Dropout(dropout)
self.decoder_layers = nn.ModuleList([
DecLayer(hidden_dim, hidden_dim*3, dropout=dropout)
for _ in range(num_decoder_layers)
])
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def encode(self, feature_dict):
mask = feature_dict["mask"]
V, E, E_idx, Y_nodes, Y_edges, E_context, Y_m = self.features.features_encode(feature_dict)
h_E_context = self.W_e_context(E_context)
h_V = self.W_v(V)
h_E = self.W_e(E)
mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1)
mask_attend = mask.unsqueeze(-1) * mask_attend
for layer in self.encoder_layers:
h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend)
h_V_C = self.W_c(h_V)
Y_m_edges = Y_m[:,:,:,None]*Y_m[:,:,None,:]
Y_nodes = self.W_nodes_y(Y_nodes)
Y_edges = self.W_edges_y(Y_edges)
for i in range(len(self.context_encoder_layers)):
Y_nodes = self.y_context_encoder_layers[i](Y_nodes, Y_edges, Y_m, Y_m_edges)
h_E_context_cat = torch.cat([h_E_context, Y_nodes], -1)
h_V_C = self.context_encoder_layers[i](h_V_C, h_E_context_cat, mask, Y_m)
h_V_C = self.V_C(h_V_C)
h_V = h_V + self.V_C_norm(self.h_V_C_dropout(h_V_C))
return h_V, h_E, E_idx
def decode(self, feature_dict):
h_V = feature_dict["h_V"]
h_E = feature_dict["h_E"]
E_idx = feature_dict["E_idx"]
mask = feature_dict["mask"]
device=h_V.device
V, F = self.features.features_decode(feature_dict)
h_F = self.W_f(F)
h_EF = torch.cat([h_E, h_F], -1)
h_V_sc = self.W_v_sc(V)
h_V_combined = torch.cat([h_V, h_V_sc], -1)
h_V = self.linear_down(h_V_combined)
for layer in self.decoder_layers:
h_EV = cat_neighbors_nodes(h_V, h_EF, E_idx)
h_V = layer(h_V, h_EV, mask)
torsions = self.W_torsions(h_V)
torsions = torsions.reshape(h_V.shape[0],h_V.shape[1], 4, self.num_mix, 3)
mean = torsions[:,:,:,:,0].float()
concentration = 0.1 + self.softplus(torsions[:,:,:,:,1]).float()
mix_logits = torsions[:,:,:,:,2].float()
return mean, concentration, mix_logits
class DecLayerJ(nn.Module):
def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30):
super(DecLayerJ, self).__init__()
self.num_hidden = num_hidden
self.num_in = num_in
self.scale = scale
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.norm1 = nn.LayerNorm(num_hidden)
self.norm2 = nn.LayerNorm(num_hidden)
self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
self.W2 = nn.Linear(num_hidden, num_hidden, bias=True)
self.W3 = nn.Linear(num_hidden, num_hidden, bias=True)
self.act = torch.nn.GELU()
self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4)
def forward(self, h_V, h_E, mask_V=None, mask_attend=None):
""" Parallel computation of full transformer layer """
h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,-1,h_E.size(-2),-1)
h_EV = torch.cat([h_V_expand, h_E], -1)
h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV)))))
if mask_attend is not None:
h_message = mask_attend.unsqueeze(-1) * h_message
dh = torch.sum(h_message, -2) / self.scale
h_V = self.norm1(h_V + self.dropout1(dh))
dh = self.dense(h_V)
h_V = self.norm2(h_V + self.dropout2(dh))
if mask_V is not None:
mask_V = mask_V.unsqueeze(-1)
h_V = mask_V * h_V
return h_V
class ProteinFeatures(nn.Module):
def __init__(self,
edge_features=128,
node_features=128,
num_positional_embeddings=16,
num_chain_embeddings=16,
num_rbf=16,
top_k=30,
augment_eps=0.,
atom37_order=False,
device=None,
atom_context_num=16,
lower_bound=0.0,
upper_bound=20.0):
""" Extract protein features """
super(ProteinFeatures, self).__init__()
self.edge_features = edge_features
self.node_features = node_features
self.num_positional_embeddings = num_positional_embeddings
self.num_chain_embeddings = num_chain_embeddings
self.num_rbf = num_rbf
self.top_k = top_k
self.augment_eps = augment_eps
self.atom37_order = atom37_order
self.device = device
self.atom_context_num = atom_context_num
self.lower_bound = lower_bound
self.upper_bound = upper_bound
self.N_idx = 0
self.CA_idx = 1
self.C_idx = 2
if atom37_order:
self.O_idx = 4
else:
self.O_idx = 3
self.positional_embeddings = PositionalEncodings(num_positional_embeddings)
enc_node_in = 21
enc_edge_in = num_positional_embeddings + num_rbf*25
self.enc_node_in = enc_node_in
self.enc_edge_in = enc_edge_in
self.enc_edge_embedding = nn.Linear(enc_edge_in, edge_features, bias=False)
self.enc_norm_edges = nn.LayerNorm(edge_features)
self.enc_node_embedding = nn.Linear(enc_node_in, node_features, bias=False)
self.enc_norm_nodes = nn.LayerNorm(node_features)
dec_node_in = 14*atom_context_num*num_rbf
dec_edge_in = num_rbf*14*14+42
self.dec_node_in = dec_node_in
self.dec_edge_in = dec_edge_in
self.W_XY_project_down1 = nn.Linear(num_rbf+120, num_rbf, bias=True)
self.dec_edge_embedding1 = nn.Linear(dec_edge_in, edge_features, bias=False)
self.dec_norm_edges1 = nn.LayerNorm(edge_features)
self.dec_node_embedding1 = nn.Linear(dec_node_in, node_features, bias=False)
self.dec_norm_nodes1 = nn.LayerNorm(node_features)
self.node_project_down = nn.Linear(5*num_rbf+64+4, node_features, bias=True)
self.norm_nodes = nn.LayerNorm(node_features)
self.type_linear = nn.Linear(147, 64)
self.y_nodes = nn.Linear(147, node_features, bias=False)
self.y_edges = nn.Linear(num_rbf, node_features, bias=False)
self.norm_y_edges = nn.LayerNorm(node_features)
self.norm_y_nodes = nn.LayerNorm(node_features)
self.periodic_table_features = torch.tensor([[ 0, 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],
[ 0, 1, 18, 1, 2, 13, 14, 15, 16, 17, 18, 1, 2, 13,
14, 15, 16, 17, 18, 1, 2, 3, 4, 5, 6, 7, 8, 9,
10, 11, 12, 13, 14, 15, 16, 17, 18, 1, 2, 3, 4, 5,
6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 1,
2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
16, 17, 18, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3,
3, 3, 3, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10, 11,
12, 13, 14, 15, 16, 17, 18],
[ 0, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3,
3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6,
6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
7, 7, 7, 7, 7, 7, 7]], dtype=torch.long, device=device)
def _dist(self, X, mask, eps=1E-6):
mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
D = mask_2D * torch.sqrt(torch.sum(dX**2, 3) + eps)
D_max, _ = torch.max(D, -1, keepdim=True)
D_adjust = D + (1. - mask_2D) * D_max
sampled_top_k = self.top_k
D_neighbors, E_idx = torch.topk(D_adjust, np.minimum(self.top_k, X.shape[1]), dim=-1, largest=False)
return D_neighbors, E_idx
def _make_angle_features(self, A, B, C, Y):
v1 = A - B
v2 = C - B
e1 = torch.nn.functional.normalize(v1, dim=-1)
e1_v2_dot = torch.einsum("bli, bli -> bl", e1, v2)[...,None]
u2 = v2 - e1*e1_v2_dot
e2 = torch.nn.functional.normalize(u2, dim=-1)
e3 = torch.cross(e1, e2, dim=-1)
R_residue = torch.cat((e1[:,:,:,None],e2[:,:,:,None],e3[:,:,:,None]), dim=-1)
local_vectors = torch.einsum("blqp, blyq -> blyp", R_residue, Y-B[:,:,None,:])
rxy = torch.sqrt(local_vectors[...,0]**2 + local_vectors[...,1]**2 + 1e-8)
f1 = local_vectors[...,0]/rxy
f2 = local_vectors[...,1]/rxy
rxyz = torch.norm(local_vectors, dim=-1)+1e-8
f3 = rxy/rxyz
f4 = local_vectors[...,2]/rxyz
f = torch.cat([f1[...,None],f2[...,None],f3[...,None],f4[...,None]],-1)
return f
def _rbf(self, D, D_mu_shape=[1,1,1,-1], lower_bound=0.0, upper_bound=20.0, num_bins=16):
device = D.device
D_min, D_max, D_count = lower_bound, upper_bound, num_bins
D_mu = torch.linspace(D_min, D_max, D_count, device=device)
D_mu = D_mu.view(D_mu_shape)
D_sigma = (D_max - D_min) / D_count
D_expand = torch.unsqueeze(D, -1)
RBF = torch.exp(-((D_expand - D_mu) / D_sigma)**2)
return RBF
def _get_rbf(self, A, B, E_idx, D_mu_shape=[1,1,1,-1], lower_bound=2.0, upper_bound=22.0, num_bins=16):
D_A_B = torch.sqrt(torch.sum((A[:,:,None,:] - B[:,None,:,:])**2,-1) + 1e-6)
D_A_B_neighbors = gather_edges(D_A_B[:,:,:,None], E_idx)[:,:,:,0]
RBF_A_B = self._rbf(D_A_B_neighbors, D_mu_shape=D_mu_shape, lower_bound=lower_bound, upper_bound=upper_bound, num_bins=num_bins)
return RBF_A_B
def features_encode(self, features):
"""
make protein graph and encode backbone
"""
S = features["S"]
X = features["X"]
Y = features["Y"]
Y_m = features["Y_m"]
Y_t = features["Y_t"]
mask = features["mask"]
R_idx = features["R_idx"]
chain_labels = features["chain_labels"]
if self.training and self.augment_eps > 0:
X = X + self.augment_eps * torch.randn_like(X)
Ca = X[:,:,self.CA_idx,:]
N = X[:,:,self.N_idx,:]
C = X[:,:,self.C_idx,:]
O = X[:,:,self.O_idx,:]
b = Ca - N
c = C - Ca
a = torch.cross(b, c, dim=-1)
Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca
_, E_idx = self._dist(Ca, mask)
backbone_coords_list = [N, Ca, C, O, Cb]
RBF_all = []
for atom_1 in backbone_coords_list:
for atom_2 in backbone_coords_list:
RBF_all.append(self._get_rbf(atom_1,
atom_2,
E_idx,
D_mu_shape=[1,1,1,-1],
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
num_bins=self.num_rbf))
RBF_all = torch.cat(tuple(RBF_all), dim=-1)
offset = R_idx[:,:,None]-R_idx[:,None,:]
offset = gather_edges(offset[:,:,:,None], E_idx)[:,:,:,0]
d_chains = ((chain_labels[:, :, None] - chain_labels[:,None,:])==0).long()
E_chains = gather_edges(d_chains[:,:,:,None], E_idx)[:,:,:,0]
E_positional = self.positional_embeddings(offset.long(), E_chains)
E = torch.cat((E_positional, RBF_all), -1)
E = self.enc_edge_embedding(E)
E = self.enc_norm_edges(E)
V = torch.nn.functional.one_hot(S, self.enc_node_in).float()
V = self.enc_node_embedding(V)
V = self.enc_norm_nodes(V)
Y_t = Y_t.long()
Y_t_g = self.periodic_table_features[1][Y_t]
Y_t_p = self.periodic_table_features[2][Y_t]
Y_t_g_1hot_ = torch.nn.functional.one_hot(Y_t_g, 19)
Y_t_p_1hot_ = torch.nn.functional.one_hot(Y_t_p, 8)
Y_t_1hot_ = torch.nn.functional.one_hot(Y_t, 120)
Y_t_1hot_ = torch.cat([Y_t_1hot_, Y_t_g_1hot_, Y_t_p_1hot_], -1)
Y_t_1hot = self.type_linear(Y_t_1hot_.float())
D_N_Y = torch.sqrt(torch.sum((N[:,:,None,:] - Y)**2,-1) + 1e-6)
D_N_Y = self._rbf(D_N_Y, D_mu_shape=[1,1,1,-1], lower_bound=self.lower_bound, upper_bound=self.upper_bound, num_bins=self.num_rbf)
D_Ca_Y = torch.sqrt(torch.sum((Ca[:,:,None,:] - Y)**2,-1) + 1e-6)
D_Ca_Y = self._rbf(D_Ca_Y, D_mu_shape=[1,1,1,-1], lower_bound=self.lower_bound, upper_bound=self.upper_bound, num_bins=self.num_rbf)
D_C_Y = torch.sqrt(torch.sum((C[:,:,None,:] - Y)**2,-1) + 1e-6)
D_C_Y = self._rbf(D_C_Y, D_mu_shape=[1,1,1,-1], lower_bound=self.lower_bound, upper_bound=self.upper_bound, num_bins=self.num_rbf)
D_O_Y = torch.sqrt(torch.sum((O[:,:,None,:] - Y)**2,-1) + 1e-6)
D_O_Y = self._rbf(D_O_Y, D_mu_shape=[1,1,1,-1], lower_bound=self.lower_bound, upper_bound=self.upper_bound, num_bins=self.num_rbf)
D_Cb_Y = torch.sqrt(torch.sum((Cb[:,:,None,:] - Y)**2,-1) + 1e-6)
D_Cb_Y = self._rbf(D_Cb_Y, D_mu_shape=[1,1,1,-1], lower_bound=self.lower_bound, upper_bound=self.upper_bound, num_bins=self.num_rbf)
f_angles = self._make_angle_features(N, Ca, C, Y)
D_all = torch.cat((D_N_Y, D_Ca_Y, D_C_Y, D_O_Y, D_Cb_Y, Y_t_1hot, f_angles), dim=-1)
E_context = self.node_project_down(D_all)
E_context = self.norm_nodes(E_context)
Y_edges = self._rbf(torch.sqrt(torch.sum((Y[:,:,:,None,:] - Y[:,:,None,:,:])**2,-1) + 1e-6))
Y_edges = self.y_edges(Y_edges)
Y_nodes = self.y_nodes(Y_t_1hot_.float())
Y_edges = self.norm_y_edges(Y_edges)
Y_nodes = self.norm_y_nodes(Y_nodes)
return V, E, E_idx, Y_nodes, Y_edges, E_context, Y_m
def features_decode(self, features):
"""
Make features for decoding. Explicit side chain atom and other atom distances.
"""
S = features["S"]
X = features["X"]
X_m = features["X_m"]
mask = features["mask"]
E_idx = features["E_idx"]
Y = features["Y"][:,:,:self.atom_context_num]
Y_m = features["Y_m"][:,:,:self.atom_context_num]
Y_t = features["Y_t"][:,:,:self.atom_context_num]
X_m = X_m*mask[:,:,None]
device = S.device
B, L, _, _ = X.shape
RBF_sidechain = []
X_m_gathered = gather_nodes(X_m, E_idx)
for i in range(14):
for j in range(14):
rbf_features = self._get_rbf(X[:,:,i,:],
X[:,:,j,:],
E_idx,
D_mu_shape=[1,1,1,-1],
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
num_bins=self.num_rbf)
rbf_features = rbf_features*X_m[:,:,i,None,None]*X_m_gathered[:,:,:,j,None]
RBF_sidechain.append(rbf_features)
D_XY = torch.sqrt(torch.sum((X[:,:,:,None,:] - Y[:,:,None,:,:])**2,-1) + 1e-6)
XY_features = self._rbf(D_XY,
D_mu_shape=[1,1,1,1,-1],
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
num_bins=self.num_rbf)
XY_features = XY_features*X_m[:,:,:,None,None]*Y_m[:,:,None,:,None]
Y_t_1hot = torch.nn.functional.one_hot(Y_t.long(), 120).float()
XY_Y_t = torch.cat([XY_features, Y_t_1hot[:,:,None,:,:].repeat(1,1,14,1,1)],-1)
XY_Y_t = self.W_XY_project_down1(XY_Y_t)
XY_features = XY_Y_t.view([B, L, -1])
V = self.dec_node_embedding1(XY_features)
V = self.dec_norm_nodes1(V)
S_1h = torch.nn.functional.one_hot(S, self.enc_node_in).float()
S_1h_gathered = gather_nodes(S_1h, E_idx)
S_features = torch.cat([S_1h[:,:,None,:].repeat(1,1,E_idx.shape[2],1), S_1h_gathered], -1)
F = torch.cat(tuple(RBF_sidechain), dim=-1)
F = torch.cat([F, S_features],-1)
F = self.dec_edge_embedding1(F)
F = self.dec_norm_edges1(F)
return V, F
class PositionWiseFeedForward(nn.Module):
def __init__(self, num_hidden, num_ff):
super(PositionWiseFeedForward, self).__init__()
self.W_in = nn.Linear(num_hidden, num_ff, bias=True)
self.W_out = nn.Linear(num_ff, num_hidden, bias=True)
self.act = torch.nn.GELU()
def forward(self, h_V):
h = self.act(self.W_in(h_V))
h = self.W_out(h)
return h
class PositionalEncodings(nn.Module):
def __init__(self, num_embeddings, max_relative_feature=32):
super(PositionalEncodings, self).__init__()
self.num_embeddings = num_embeddings
self.max_relative_feature = max_relative_feature
self.linear = nn.Linear(2*max_relative_feature+1+1, num_embeddings)
def forward(self, offset, mask):
d = torch.clip(offset + self.max_relative_feature, 0, 2*self.max_relative_feature)*mask + (1-mask)*(2*self.max_relative_feature+1)
d_onehot = torch.nn.functional.one_hot(d, 2*self.max_relative_feature+1+1)
E = self.linear(d_onehot.float())
return E
class DecLayer(nn.Module):
def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30):
super(DecLayer, self).__init__()
self.num_hidden = num_hidden
self.num_in = num_in
self.scale = scale
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.norm1 = nn.LayerNorm(num_hidden)
self.norm2 = nn.LayerNorm(num_hidden)
self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
self.W2 = nn.Linear(num_hidden, num_hidden, bias=True)
self.W3 = nn.Linear(num_hidden, num_hidden, bias=True)
self.act = torch.nn.GELU()
self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4)
def forward(self, h_V, h_E, mask_V=None, mask_attend=None):
""" Parallel computation of full transformer layer """
h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_E.size(-2),-1)
h_EV = torch.cat([h_V_expand, h_E], -1)
h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV)))))
if mask_attend is not None:
h_message = mask_attend.unsqueeze(-1) * h_message
dh = torch.sum(h_message, -2) / self.scale
h_V = self.norm1(h_V + self.dropout1(dh))
dh = self.dense(h_V)
h_V = self.norm2(h_V + self.dropout2(dh))
if mask_V is not None:
mask_V = mask_V.unsqueeze(-1)
h_V = mask_V * h_V
return h_V
class EncLayer(nn.Module):
def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30):
super(EncLayer, self).__init__()
self.num_hidden = num_hidden
self.num_in = num_in
self.scale = scale
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.dropout3 = nn.Dropout(dropout)
self.norm1 = nn.LayerNorm(num_hidden)
self.norm2 = nn.LayerNorm(num_hidden)
self.norm3 = nn.LayerNorm(num_hidden)
self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
self.W2 = nn.Linear(num_hidden, num_hidden, bias=True)
self.W3 = nn.Linear(num_hidden, num_hidden, bias=True)
self.W11 = nn.Linear(num_hidden + num_in, num_hidden, bias=True)
self.W12 = nn.Linear(num_hidden, num_hidden, bias=True)
self.W13 = nn.Linear(num_hidden, num_hidden, bias=True)
self.act = torch.nn.GELU()
self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4)
def forward(self, h_V, h_E, E_idx, mask_V=None, mask_attend=None):
""" Parallel computation of full transformer layer """
h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1)
h_EV = torch.cat([h_V_expand, h_EV], -1)
h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV)))))
if mask_attend is not None:
h_message = mask_attend.unsqueeze(-1) * h_message
dh = torch.sum(h_message, -2) / self.scale
h_V = self.norm1(h_V + self.dropout1(dh))
dh = self.dense(h_V)
h_V = self.norm2(h_V + self.dropout2(dh))
if mask_V is not None:
mask_V = mask_V.unsqueeze(-1)
h_V = mask_V * h_V
h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1)
h_EV = torch.cat([h_V_expand, h_EV], -1)
h_message = self.W13(self.act(self.W12(self.act(self.W11(h_EV)))))
h_E = self.norm3(h_E + self.dropout3(h_message))
return h_V, h_E
def gather_edges(edges, neighbor_idx):
neighbors = neighbor_idx.unsqueeze(-1).expand(-1, -1, -1, edges.size(-1))
edge_features = torch.gather(edges, 2, neighbors)
return edge_features
def gather_nodes(nodes, neighbor_idx):
neighbors_flat = neighbor_idx.reshape((neighbor_idx.shape[0], -1))
neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2))
neighbor_features = torch.gather(nodes, 1, neighbors_flat)
neighbor_features = neighbor_features.view(list(neighbor_idx.shape)[:3] + [-1])
return neighbor_features
def gather_nodes_t(nodes, neighbor_idx):
idx_flat = neighbor_idx.unsqueeze(-1).expand(-1, -1, nodes.size(2))
neighbor_features = torch.gather(nodes, 1, idx_flat)
return neighbor_features
def cat_neighbors_nodes(h_nodes, h_neighbors, E_idx):
h_nodes = gather_nodes(h_nodes, E_idx)
h_nn = torch.cat([h_neighbors, h_nodes], -1)
return h_nn