Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from Embeddings import MSA_emb, Extra_emb, Templ_emb, Recycling | |
| from Track_module import IterativeSimulator | |
| from AuxiliaryPredictor import DistanceNetwork, MaskedTokenNetwork, ExpResolvedNetwork, LDDTNetwork | |
| from util import INIT_CRDS | |
| from opt_einsum import contract as einsum | |
| from icecream import ic | |
| class RoseTTAFoldModule(nn.Module): | |
| def __init__(self, n_extra_block=4, n_main_block=8, n_ref_block=4,\ | |
| d_msa=256, d_msa_full=64, d_pair=128, d_templ=64, | |
| n_head_msa=8, n_head_pair=4, n_head_templ=4, | |
| d_hidden=32, d_hidden_templ=64, | |
| p_drop=0.15, d_t1d=24, d_t2d=44, | |
| SE3_param_full={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, | |
| SE3_param_topk={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, | |
| ): | |
| super(RoseTTAFoldModule, self).__init__() | |
| # | |
| # Input Embeddings | |
| d_state = SE3_param_topk['l0_out_features'] | |
| self.latent_emb = MSA_emb(d_msa=d_msa, d_pair=d_pair, d_state=d_state, p_drop=p_drop) | |
| self.full_emb = Extra_emb(d_msa=d_msa_full, d_init=25, p_drop=p_drop) | |
| self.templ_emb = Templ_emb(d_pair=d_pair, d_templ=d_templ, d_state=d_state, | |
| n_head=n_head_templ, | |
| d_hidden=d_hidden_templ, p_drop=0.25, d_t1d=d_t1d, d_t2d=d_t2d) | |
| # Update inputs with outputs from previous round | |
| self.recycle = Recycling(d_msa=d_msa, d_pair=d_pair, d_state=d_state) | |
| # | |
| self.simulator = IterativeSimulator(n_extra_block=n_extra_block, | |
| n_main_block=n_main_block, | |
| n_ref_block=n_ref_block, | |
| d_msa=d_msa, d_msa_full=d_msa_full, | |
| d_pair=d_pair, d_hidden=d_hidden, | |
| n_head_msa=n_head_msa, | |
| n_head_pair=n_head_pair, | |
| SE3_param_full=SE3_param_full, | |
| SE3_param_topk=SE3_param_topk, | |
| p_drop=p_drop) | |
| ## | |
| self.c6d_pred = DistanceNetwork(d_pair, p_drop=p_drop) | |
| self.aa_pred = MaskedTokenNetwork(d_msa, p_drop=p_drop) | |
| self.lddt_pred = LDDTNetwork(d_state) | |
| self.exp_pred = ExpResolvedNetwork(d_msa, d_state) | |
| def forward(self, msa_latent, msa_full, seq, xyz, idx, | |
| seq1hot=None, t1d=None, t2d=None, xyz_t=None, alpha_t=None, | |
| msa_prev=None, pair_prev=None, state_prev=None, | |
| return_raw=False, return_full=False, | |
| use_checkpoint=False, return_infer=False): | |
| B, N, L = msa_latent.shape[:3] | |
| # Get embeddings | |
| #ic(seq.shape) | |
| #ic(msa_latent.shape) | |
| #ic(seq1hot.shape) | |
| #ic(idx.shape) | |
| #ic(xyz.shape) | |
| #ic(seq1hot.shape) | |
| #ic(t1d.shape) | |
| #ic(t2d.shape) | |
| idx = idx.long() | |
| msa_latent, pair, state = self.latent_emb(msa_latent, seq, idx, seq1hot=seq1hot) | |
| msa_full = self.full_emb(msa_full, seq, idx, seq1hot=seq1hot) | |
| # | |
| # Do recycling | |
| if msa_prev == None: | |
| msa_prev = torch.zeros_like(msa_latent[:,0]) | |
| if pair_prev == None: | |
| pair_prev = torch.zeros_like(pair) | |
| if state_prev == None: | |
| state_prev = torch.zeros_like(state) | |
| #ic(seq.shape) | |
| #ic(msa_prev.shape) | |
| #ic(pair_prev.shape) | |
| #ic(xyz.shape) | |
| #ic(state_prev.shape) | |
| msa_recycle, pair_recycle, state_recycle = self.recycle(seq, msa_prev, pair_prev, xyz, state_prev) | |
| msa_latent[:,0] = msa_latent[:,0] + msa_recycle.reshape(B,L,-1) | |
| pair = pair + pair_recycle | |
| state = state + state_recycle | |
| # | |
| #ic(t1d.dtype) | |
| #ic(t2d.dtype) | |
| #ic(alpha_t.dtype) | |
| #ic(xyz_t.dtype) | |
| #ic(pair.dtype) | |
| #ic(state.dtype) | |
| #import pdb; pdb.set_trace() | |
| # add template embedding | |
| pair, state = self.templ_emb(t1d, t2d, alpha_t, xyz_t, pair, state, use_checkpoint=use_checkpoint) | |
| #ic(seq.dtype) | |
| #ic(msa_latent.dtype) | |
| #ic(msa_full.dtype) | |
| #ic(pair.dtype) | |
| #ic(xyz.dtype) | |
| #ic(state.dtype) | |
| #ic(idx.dtype) | |
| # Predict coordinates from given inputs | |
| msa, pair, R, T, alpha_s, state = self.simulator(seq, msa_latent, msa_full.type(torch.float32), pair, xyz[:,:,:3], | |
| state, idx, use_checkpoint=use_checkpoint) | |
| if return_raw: | |
| # get last structure | |
| xyz = einsum('bnij,bnaj->bnai', R[-1], xyz[:,:,:3]-xyz[:,:,1].unsqueeze(-2)) + T[-1].unsqueeze(-2) | |
| return msa[:,0], pair, xyz, state, alpha_s[-1] | |
| # predict masked amino acids | |
| logits_aa = self.aa_pred(msa) | |
| # | |
| # predict distogram & orientograms | |
| logits = self.c6d_pred(pair) | |
| # Predict LDDT | |
| lddt = self.lddt_pred(state) | |
| # predict experimentally resolved or not | |
| logits_exp = self.exp_pred(msa[:,0], state) | |
| if return_infer: | |
| #get last structure | |
| xyz = einsum('bnij,bnaj->bnai', R[-1], xyz[:,:,:3]-xyz[:,:,1].unsqueeze(-2)) + T[-1].unsqueeze(-2) | |
| return logits, logits_aa, logits_exp, xyz, lddt, msa[:,0], pair, state, alpha_s[-1] | |
| # get all intermediate bb structures | |
| xyz = einsum('rbnij,bnaj->rbnai', R, xyz[:,:,:3]-xyz[:,:,1].unsqueeze(-2)) + T.unsqueeze(-2) | |
| return logits, logits_aa, logits_exp, xyz, alpha_s, lddt | |