Upload 95 files

rfdiffusion.egg-info/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: rfdiffusion
+Version: 1.1.0
+Summary: RFdiffusion is an open source method for protein structure generation.
+Home-page: https://github.com/RosettaCommons/RFdiffusion
+Author: Rosetta Commons
+License-File: LICENSE

rfdiffusion.egg-info/SOURCES.txt

@@ -0,0 +1,34 @@
+LICENSE
+README.md
+setup.py
+rfdiffusion/Attention_module.py
+rfdiffusion/AuxiliaryPredictor.py
+rfdiffusion/Embeddings.py
+rfdiffusion/RoseTTAFoldModel.py
+rfdiffusion/SE3_network.py
+rfdiffusion/Track_module.py
+rfdiffusion/__init__.py
+rfdiffusion/chemical.py
+rfdiffusion/contigs.py
+rfdiffusion/coords6d.py
+rfdiffusion/diffusion.py
+rfdiffusion/igso3.py
+rfdiffusion/kinematics.py
+rfdiffusion/model_input_logger.py
+rfdiffusion/scoring.py
+rfdiffusion/util.py
+rfdiffusion/util_module.py
+rfdiffusion.egg-info/PKG-INFO
+rfdiffusion.egg-info/SOURCES.txt
+rfdiffusion.egg-info/dependency_links.txt
+rfdiffusion.egg-info/requires.txt
+rfdiffusion.egg-info/top_level.txt
+rfdiffusion/inference/__init__.py
+rfdiffusion/inference/model_runners.py
+rfdiffusion/inference/symmetry.py
+rfdiffusion/inference/utils.py
+rfdiffusion/potentials/__init__.py
+rfdiffusion/potentials/manager.py
+rfdiffusion/potentials/potentials.py
+scripts/run_inference.py
+tests/test_diffusion.py

rfdiffusion.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

rfdiffusion.egg-info/requires.txt

@@ -0,0 +1,2 @@
+torch
+se3-transformer

rfdiffusion.egg-info/top_level.txt

@@ -0,0 +1 @@
+rfdiffusion

rfdiffusion/Attention_module.py

@@ -0,0 +1,404 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from opt_einsum import contract as einsum
+from rfdiffusion.util_module import init_lecun_normal
+
+class FeedForwardLayer(nn.Module):
+    def __init__(self, d_model, r_ff, p_drop=0.1):
+        super(FeedForwardLayer, self).__init__()
+        self.norm = nn.LayerNorm(d_model)
+        self.linear1 = nn.Linear(d_model, d_model*r_ff)
+        self.dropout = nn.Dropout(p_drop)
+        self.linear2 = nn.Linear(d_model*r_ff, d_model)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # initialize linear layer right before ReLu: He initializer (kaiming normal)
+        nn.init.kaiming_normal_(self.linear1.weight, nonlinearity='relu')
+        nn.init.zeros_(self.linear1.bias)
+
+        # initialize linear layer right before residual connection: zero initialize
+        nn.init.zeros_(self.linear2.weight)
+        nn.init.zeros_(self.linear2.bias)
+    
+    def forward(self, src):
+        src = self.norm(src)
+        src = self.linear2(self.dropout(F.relu_(self.linear1(src))))
+        return src
+
+class Attention(nn.Module):
+    # calculate multi-head attention
+    def __init__(self, d_query, d_key, n_head, d_hidden, d_out):
+        super(Attention, self).__init__()
+        self.h = n_head
+        self.dim = d_hidden
+        #
+        self.to_q = nn.Linear(d_query, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_key, n_head*d_hidden, bias=False)
+        self.to_v = nn.Linear(d_key, n_head*d_hidden, bias=False)
+        #
+        self.to_out = nn.Linear(n_head*d_hidden, d_out)
+        self.scaling = 1/math.sqrt(d_hidden)
+        #
+        # initialize all parameters properly
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, query, key, value):
+        B, Q = query.shape[:2]
+        B, K = key.shape[:2]
+        #
+        query = self.to_q(query).reshape(B, Q, self.h, self.dim)
+        key = self.to_k(key).reshape(B, K, self.h, self.dim)
+        value = self.to_v(value).reshape(B, K, self.h, self.dim)
+        #
+        query = query * self.scaling
+        attn = einsum('bqhd,bkhd->bhqk', query, key)
+        attn = F.softmax(attn, dim=-1)
+        #
+        out = einsum('bhqk,bkhd->bqhd', attn, value)
+        out = out.reshape(B, Q, self.h*self.dim)
+        #
+        out = self.to_out(out)
+
+        return out
+
+class AttentionWithBias(nn.Module):
+    def __init__(self, d_in=256, d_bias=128, n_head=8, d_hidden=32):
+        super(AttentionWithBias, self).__init__()
+        self.norm_in = nn.LayerNorm(d_in)
+        self.norm_bias = nn.LayerNorm(d_bias)
+        #
+        self.to_q = nn.Linear(d_in, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_in, n_head*d_hidden, bias=False)
+        self.to_v = nn.Linear(d_in, n_head*d_hidden, bias=False)
+        self.to_b = nn.Linear(d_bias, n_head, bias=False)
+        self.to_g = nn.Linear(d_in, n_head*d_hidden)
+        self.to_out = nn.Linear(n_head*d_hidden, d_in)
+
+        self.scaling = 1/math.sqrt(d_hidden)
+        self.h = n_head
+        self.dim = d_hidden
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+        
+        # bias: normal distribution
+        self.to_b = init_lecun_normal(self.to_b)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, x, bias):
+        B, L = x.shape[:2]
+        #
+        x = self.norm_in(x)
+        bias = self.norm_bias(bias)
+        #
+        query = self.to_q(x).reshape(B, L, self.h, self.dim)
+        key = self.to_k(x).reshape(B, L, self.h, self.dim)
+        value = self.to_v(x).reshape(B, L, self.h, self.dim)
+        bias = self.to_b(bias) # (B, L, L, h)
+        gate = torch.sigmoid(self.to_g(x))
+        #
+        key = key * self.scaling
+        attn = einsum('bqhd,bkhd->bqkh', query, key)
+        attn = attn + bias
+        attn = F.softmax(attn, dim=-2)
+        #
+        out = einsum('bqkh,bkhd->bqhd', attn, value).reshape(B, L, -1)
+        out = gate * out
+        #
+        out = self.to_out(out)
+        return out
+
+# MSA Attention (row/column) from AlphaFold architecture
+class SequenceWeight(nn.Module):
+    def __init__(self, d_msa, n_head, d_hidden, p_drop=0.1):
+        super(SequenceWeight, self).__init__()
+        self.h = n_head
+        self.dim = d_hidden
+        self.scale = 1.0 / math.sqrt(self.dim)
+
+        self.to_query = nn.Linear(d_msa, n_head*d_hidden)
+        self.to_key = nn.Linear(d_msa, n_head*d_hidden)
+        self.dropout = nn.Dropout(p_drop)
+
+        self.reset_parameter()
+    
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_query.weight)
+        nn.init.xavier_uniform_(self.to_key.weight)
+
+    def forward(self, msa):
+        B, N, L = msa.shape[:3]
+       
+        tar_seq = msa[:,0]
+        
+        q = self.to_query(tar_seq).view(B, 1, L, self.h, self.dim)
+        k = self.to_key(msa).view(B, N, L, self.h, self.dim)
+        
+        q = q * self.scale
+        attn = einsum('bqihd,bkihd->bkihq', q, k)
+        attn = F.softmax(attn, dim=1)
+        return self.dropout(attn)
+
+class MSARowAttentionWithBias(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, n_head=8, d_hidden=32):
+        super(MSARowAttentionWithBias, self).__init__()
+        self.norm_msa = nn.LayerNorm(d_msa)
+        self.norm_pair = nn.LayerNorm(d_pair)
+        #
+        self.seq_weight = SequenceWeight(d_msa, n_head, d_hidden, p_drop=0.1)
+        self.to_q = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_v = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_b = nn.Linear(d_pair, n_head, bias=False)
+        self.to_g = nn.Linear(d_msa, n_head*d_hidden)
+        self.to_out = nn.Linear(n_head*d_hidden, d_msa)
+
+        self.scaling = 1/math.sqrt(d_hidden)
+        self.h = n_head
+        self.dim = d_hidden
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+        
+        # bias: normal distribution
+        self.to_b = init_lecun_normal(self.to_b)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, msa, pair): # TODO: make this as tied-attention
+        B, N, L = msa.shape[:3]
+        #
+        msa = self.norm_msa(msa)
+        pair = self.norm_pair(pair)
+        #
+        seq_weight = self.seq_weight(msa) # (B, N, L, h, 1)
+        query = self.to_q(msa).reshape(B, N, L, self.h, self.dim)
+        key = self.to_k(msa).reshape(B, N, L, self.h, self.dim)
+        value = self.to_v(msa).reshape(B, N, L, self.h, self.dim)
+        bias = self.to_b(pair) # (B, L, L, h)
+        gate = torch.sigmoid(self.to_g(msa))
+        #
+        query = query * seq_weight.expand(-1, -1, -1, -1, self.dim)
+        key = key * self.scaling
+        attn = einsum('bsqhd,bskhd->bqkh', query, key)
+        attn = attn + bias
+        attn = F.softmax(attn, dim=-2)
+        #
+        out = einsum('bqkh,bskhd->bsqhd', attn, value).reshape(B, N, L, -1)
+        out = gate * out
+        #
+        out = self.to_out(out)
+        return out
+
+class MSAColAttention(nn.Module):
+    def __init__(self, d_msa=256, n_head=8, d_hidden=32):
+        super(MSAColAttention, self).__init__()
+        self.norm_msa = nn.LayerNorm(d_msa)
+        #
+        self.to_q = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_v = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_g = nn.Linear(d_msa, n_head*d_hidden)
+        self.to_out = nn.Linear(n_head*d_hidden, d_msa)
+
+        self.scaling = 1/math.sqrt(d_hidden)
+        self.h = n_head
+        self.dim = d_hidden
+        
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, msa):
+        B, N, L = msa.shape[:3]
+        #
+        msa = self.norm_msa(msa)
+        #
+        query = self.to_q(msa).reshape(B, N, L, self.h, self.dim)
+        key = self.to_k(msa).reshape(B, N, L, self.h, self.dim)
+        value = self.to_v(msa).reshape(B, N, L, self.h, self.dim)
+        gate = torch.sigmoid(self.to_g(msa))
+        #
+        query = query * self.scaling
+        attn = einsum('bqihd,bkihd->bihqk', query, key)
+        attn = F.softmax(attn, dim=-1)
+        #
+        out = einsum('bihqk,bkihd->bqihd', attn, value).reshape(B, N, L, -1)
+        out = gate * out
+        #
+        out = self.to_out(out)
+        return out
+
+class MSAColGlobalAttention(nn.Module):
+    def __init__(self, d_msa=64, n_head=8, d_hidden=8):
+        super(MSAColGlobalAttention, self).__init__()
+        self.norm_msa = nn.LayerNorm(d_msa)
+        #
+        self.to_q = nn.Linear(d_msa, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_msa, d_hidden, bias=False)
+        self.to_v = nn.Linear(d_msa, d_hidden, bias=False)
+        self.to_g = nn.Linear(d_msa, n_head*d_hidden)
+        self.to_out = nn.Linear(n_head*d_hidden, d_msa)
+
+        self.scaling = 1/math.sqrt(d_hidden)
+        self.h = n_head
+        self.dim = d_hidden
+        
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, msa):
+        B, N, L = msa.shape[:3]
+        #
+        msa = self.norm_msa(msa)
+        #
+        query = self.to_q(msa).reshape(B, N, L, self.h, self.dim)
+        query = query.mean(dim=1) # (B, L, h, dim)
+        key = self.to_k(msa) # (B, N, L, dim)
+        value = self.to_v(msa) # (B, N, L, dim)
+        gate = torch.sigmoid(self.to_g(msa)) # (B, N, L, h*dim)
+        #
+        query = query * self.scaling
+        attn = einsum('bihd,bkid->bihk', query, key) # (B, L, h, N)
+        attn = F.softmax(attn, dim=-1)
+        #
+        out = einsum('bihk,bkid->bihd', attn, value).reshape(B, 1, L, -1) # (B, 1, L, h*dim)
+        out = gate * out # (B, N, L, h*dim)
+        #
+        out = self.to_out(out)
+        return out
+
+# Instead of triangle attention, use Tied axail attention with bias from coordinates..?
+class BiasedAxialAttention(nn.Module):
+    def __init__(self, d_pair, d_bias, n_head, d_hidden, p_drop=0.1, is_row=True):
+        super(BiasedAxialAttention, self).__init__()
+        #
+        self.is_row = is_row
+        self.norm_pair = nn.LayerNorm(d_pair)
+        self.norm_bias = nn.LayerNorm(d_bias)
+
+        self.to_q = nn.Linear(d_pair, n_head*d_hidden, bias=False)
+        self.to_k = nn.Linear(d_pair, n_head*d_hidden, bias=False)
+        self.to_v = nn.Linear(d_pair, n_head*d_hidden, bias=False)
+        self.to_b = nn.Linear(d_bias, n_head, bias=False) 
+        self.to_g = nn.Linear(d_pair, n_head*d_hidden)
+        self.to_out = nn.Linear(n_head*d_hidden, d_pair)
+        
+        self.scaling = 1/math.sqrt(d_hidden)
+        self.h = n_head
+        self.dim = d_hidden
+        
+        # initialize all parameters properly
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+
+        # bias: normal distribution
+        self.to_b = init_lecun_normal(self.to_b)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    def forward(self, pair, bias):
+        # pair: (B, L, L, d_pair)
+        B, L = pair.shape[:2]
+        
+        if self.is_row:
+            pair = pair.permute(0,2,1,3)
+            bias = bias.permute(0,2,1,3)
+
+        pair = self.norm_pair(pair)
+        bias = self.norm_bias(bias)
+        
+        query = self.to_q(pair).reshape(B, L, L, self.h, self.dim)
+        key = self.to_k(pair).reshape(B, L, L, self.h, self.dim)
+        value = self.to_v(pair).reshape(B, L, L, self.h, self.dim)
+        bias = self.to_b(bias) # (B, L, L, h)
+        gate = torch.sigmoid(self.to_g(pair)) # (B, L, L, h*dim) 
+        
+        query = query * self.scaling
+        key = key / math.sqrt(L) # normalize for tied attention
+        attn = einsum('bnihk,bnjhk->bijh', query, key) # tied attention
+        attn = attn + bias # apply bias
+        attn = F.softmax(attn, dim=-2) # (B, L, L, h)
+        
+        out = einsum('bijh,bkjhd->bikhd', attn, value).reshape(B, L, L, -1)
+        out = gate * out
+        
+        out = self.to_out(out)
+        if self.is_row:
+            out = out.permute(0,2,1,3)
+        return out
+

rfdiffusion/AuxiliaryPredictor.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+
+class DistanceNetwork(nn.Module):
+    def __init__(self, n_feat, p_drop=0.1):
+        super(DistanceNetwork, self).__init__()
+        #
+        self.proj_symm = nn.Linear(n_feat, 37*2)
+        self.proj_asymm = nn.Linear(n_feat, 37+19)
+    
+        self.reset_parameter()
+    
+    def reset_parameter(self):
+        # initialize linear layer for final logit prediction
+        nn.init.zeros_(self.proj_symm.weight)
+        nn.init.zeros_(self.proj_asymm.weight)
+        nn.init.zeros_(self.proj_symm.bias)
+        nn.init.zeros_(self.proj_asymm.bias)
+
+    def forward(self, x):
+        # input: pair info (B, L, L, C)
+
+        # predict theta, phi (non-symmetric)
+        logits_asymm = self.proj_asymm(x)
+        logits_theta = logits_asymm[:,:,:,:37].permute(0,3,1,2)
+        logits_phi = logits_asymm[:,:,:,37:].permute(0,3,1,2)
+
+        # predict dist, omega
+        logits_symm = self.proj_symm(x)
+        logits_symm = logits_symm + logits_symm.permute(0,2,1,3)
+        logits_dist = logits_symm[:,:,:,:37].permute(0,3,1,2)
+        logits_omega = logits_symm[:,:,:,37:].permute(0,3,1,2)
+
+        return logits_dist, logits_omega, logits_theta, logits_phi
+
+class MaskedTokenNetwork(nn.Module):
+    def __init__(self, n_feat):
+        super(MaskedTokenNetwork, self).__init__()
+        self.proj = nn.Linear(n_feat, 21)
+        
+        self.reset_parameter()
+    
+    def reset_parameter(self):
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x):
+        B, N, L = x.shape[:3]
+        logits = self.proj(x).permute(0,3,1,2).reshape(B, -1, N*L)
+
+        return logits
+
+class LDDTNetwork(nn.Module):
+    def __init__(self, n_feat, n_bin_lddt=50):
+        super(LDDTNetwork, self).__init__()
+        self.proj = nn.Linear(n_feat, n_bin_lddt)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x):
+        logits = self.proj(x) # (B, L, 50)
+
+        return logits.permute(0,2,1)
+
+class ExpResolvedNetwork(nn.Module):
+    def __init__(self, d_msa, d_state, p_drop=0.1):
+        super(ExpResolvedNetwork, self).__init__()
+        self.norm_msa = nn.LayerNorm(d_msa)
+        self.norm_state = nn.LayerNorm(d_state)
+        self.proj = nn.Linear(d_msa+d_state, 1)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, seq, state):
+        B, L = seq.shape[:2]
+        
+        seq = self.norm_msa(seq)
+        state = self.norm_state(state)
+        feat = torch.cat((seq, state), dim=-1)
+        logits = self.proj(feat)
+        return logits.reshape(B, L)
+
+
+

rfdiffusion/Embeddings.py

@@ -0,0 +1,303 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from opt_einsum import contract as einsum
+import torch.utils.checkpoint as checkpoint
+from rfdiffusion.util import get_tips
+from rfdiffusion.util_module import Dropout, create_custom_forward, rbf, init_lecun_normal
+from rfdiffusion.Attention_module import Attention, FeedForwardLayer, AttentionWithBias
+from rfdiffusion.Track_module import PairStr2Pair
+import math
+
+# Module contains classes and functions to generate initial embeddings
+
+class PositionalEncoding2D(nn.Module):
+    # Add relative positional encoding to pair features
+    def __init__(self, d_model, minpos=-32, maxpos=32, p_drop=0.1):
+        super(PositionalEncoding2D, self).__init__()
+        self.minpos = minpos
+        self.maxpos = maxpos
+        self.nbin = abs(minpos)+maxpos+1
+        self.emb = nn.Embedding(self.nbin, d_model)
+        self.drop = nn.Dropout(p_drop)
+
+    def forward(self, x, idx):
+        bins = torch.arange(self.minpos, self.maxpos, device=x.device)
+        seqsep = idx[:,None,:] - idx[:,:,None] # (B, L, L)
+        #
+        ib = torch.bucketize(seqsep, bins).long() # (B, L, L)
+        emb = self.emb(ib) #(B, L, L, d_model)
+        x = x + emb # add relative positional encoding
+        return self.drop(x)
+
+class MSA_emb(nn.Module):
+    # Get initial seed MSA embedding
+    def __init__(self, d_msa=256, d_pair=128, d_state=32, d_init=22+22+2+2,
+                 minpos=-32, maxpos=32, p_drop=0.1, input_seq_onehot=False):
+        super(MSA_emb, self).__init__()
+        self.emb = nn.Linear(d_init, d_msa) # embedding for general MSA
+        self.emb_q = nn.Embedding(22, d_msa) # embedding for query sequence -- used for MSA embedding
+        self.emb_left = nn.Embedding(22, d_pair) # embedding for query sequence -- used for pair embedding
+        self.emb_right = nn.Embedding(22, d_pair) # embedding for query sequence -- used for pair embedding
+        self.emb_state = nn.Embedding(22, d_state)
+        self.drop = nn.Dropout(p_drop)
+        self.pos = PositionalEncoding2D(d_pair, minpos=minpos, maxpos=maxpos, p_drop=p_drop)
+
+        self.input_seq_onehot=input_seq_onehot
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        self.emb = init_lecun_normal(self.emb)
+        self.emb_q = init_lecun_normal(self.emb_q)
+        self.emb_left = init_lecun_normal(self.emb_left)
+        self.emb_right = init_lecun_normal(self.emb_right)
+        self.emb_state = init_lecun_normal(self.emb_state)
+
+        nn.init.zeros_(self.emb.bias)
+
+    def forward(self, msa, seq, idx):
+        # Inputs:
+        #   - msa: Input MSA (B, N, L, d_init)
+        #   - seq: Input Sequence (B, L)
+        #   - idx: Residue index
+        # Outputs:
+        #   - msa: Initial MSA embedding (B, N, L, d_msa)
+        #   - pair: Initial Pair embedding (B, L, L, d_pair)
+
+        N = msa.shape[1] # number of sequenes in MSA
+
+        # msa embedding
+        msa = self.emb(msa) # (B, N, L, d_model) # MSA embedding
+
+        # Sergey's one hot trick
+        tmp = (seq @ self.emb_q.weight).unsqueeze(1) # (B, 1, L, d_model) -- query embedding
+
+        msa = msa + tmp.expand(-1, N, -1, -1) # adding query embedding to MSA
+        msa = self.drop(msa)
+
+        # pair embedding
+        # Sergey's one hot trick
+        left  = (seq @ self.emb_left.weight)[:,None] # (B, 1, L, d_pair)
+        right = (seq @ self.emb_right.weight)[:,:,None] # (B, L, 1, d_pair)
+
+        pair = left + right # (B, L, L, d_pair)
+        pair = self.pos(pair, idx) # add relative position
+
+        # state embedding
+        # Sergey's one hot trick
+        state = self.drop(seq @ self.emb_state.weight)
+        return msa, pair, state
+
+class Extra_emb(nn.Module):
+    # Get initial seed MSA embedding
+    def __init__(self, d_msa=256, d_init=22+1+2, p_drop=0.1, input_seq_onehot=False):
+        super(Extra_emb, self).__init__()
+        self.emb = nn.Linear(d_init, d_msa) # embedding for general MSA
+        self.emb_q = nn.Embedding(22, d_msa) # embedding for query sequence
+        self.drop = nn.Dropout(p_drop)
+
+        self.input_seq_onehot=input_seq_onehot
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        self.emb = init_lecun_normal(self.emb)
+        nn.init.zeros_(self.emb.bias)
+
+    def forward(self, msa, seq, idx):
+        # Inputs:
+        #   - msa: Input MSA (B, N, L, d_init)
+        #   - seq: Input Sequence (B, L)
+        #   - idx: Residue index
+        # Outputs:
+        #   - msa: Initial MSA embedding (B, N, L, d_msa)
+        N = msa.shape[1] # number of sequenes in MSA
+        msa = self.emb(msa) # (B, N, L, d_model) # MSA embedding
+
+        # Sergey's one hot trick
+        seq = (seq @ self.emb_q.weight).unsqueeze(1) # (B, 1, L, d_model) -- query embedding
+        msa = msa + seq.expand(-1, N, -1, -1) # adding query embedding to MSA
+        return self.drop(msa)
+
+class TemplatePairStack(nn.Module):
+    # process template pairwise features
+    # use structure-biased attention
+    def __init__(self, n_block=2, d_templ=64, n_head=4, d_hidden=16, p_drop=0.25):
+        super(TemplatePairStack, self).__init__()
+        self.n_block = n_block
+        proc_s = [PairStr2Pair(d_pair=d_templ, n_head=n_head, d_hidden=d_hidden, p_drop=p_drop) for i in range(n_block)]
+        self.block = nn.ModuleList(proc_s)
+        self.norm = nn.LayerNorm(d_templ)
+    def forward(self, templ, rbf_feat, use_checkpoint=False):
+        B, T, L = templ.shape[:3]
+        templ = templ.reshape(B*T, L, L, -1)
+
+        for i_block in range(self.n_block):
+            if use_checkpoint:
+                templ = checkpoint.checkpoint(create_custom_forward(self.block[i_block]), templ, rbf_feat)
+            else:
+                templ = self.block[i_block](templ, rbf_feat)
+        return self.norm(templ).reshape(B, T, L, L, -1)
+
+class TemplateTorsionStack(nn.Module):
+    def __init__(self, n_block=2, d_templ=64, n_head=4, d_hidden=16, p_drop=0.15):
+        super(TemplateTorsionStack, self).__init__()
+        self.n_block=n_block
+        self.proj_pair = nn.Linear(d_templ+36, d_templ)
+        proc_s = [AttentionWithBias(d_in=d_templ, d_bias=d_templ,
+                                    n_head=n_head, d_hidden=d_hidden) for i in range(n_block)]
+        self.row_attn = nn.ModuleList(proc_s)
+        proc_s = [FeedForwardLayer(d_templ, 4, p_drop=p_drop) for i in range(n_block)]
+        self.ff = nn.ModuleList(proc_s)
+        self.norm = nn.LayerNorm(d_templ)
+
+    def reset_parameter(self):
+        self.proj_pair = init_lecun_normal(self.proj_pair)
+        nn.init.zeros_(self.proj_pair.bias)
+
+    def forward(self, tors, pair, rbf_feat, use_checkpoint=False):
+        B, T, L = tors.shape[:3]
+        tors = tors.reshape(B*T, L, -1)
+        pair = pair.reshape(B*T, L, L, -1)
+        pair = torch.cat((pair, rbf_feat), dim=-1)
+        pair = self.proj_pair(pair)
+
+        for i_block in range(self.n_block):
+            if use_checkpoint:
+                tors = tors + checkpoint.checkpoint(create_custom_forward(self.row_attn[i_block]), tors, pair)
+            else:
+                tors = tors + self.row_attn[i_block](tors, pair)
+            tors = tors + self.ff[i_block](tors)
+        return self.norm(tors).reshape(B, T, L, -1)
+
+class Templ_emb(nn.Module):
+    # Get template embedding
+    # Features are
+    #   t2d:
+    #   - 37 distogram bins + 6 orientations (43)
+    #   - Mask (missing/unaligned) (1)
+    #   t1d:
+    #   - tiled AA sequence (20 standard aa + gap)
+    #   - confidence (1)
+    #   - contacting or note (1). NB this is added for diffusion model. Used only in complex training examples - 1 signifies that a residue in the non-diffused chain\
+    #     i.e. the context, is in contact with the diffused chain.
+    #
+    #Added extra t1d dimension for contacting or not
+    def __init__(self, d_t1d=21+1+1, d_t2d=43+1, d_tor=30, d_pair=128, d_state=32,
+                 n_block=2, d_templ=64,
+                 n_head=4, d_hidden=16, p_drop=0.25):
+        super(Templ_emb, self).__init__()
+        # process 2D features
+        self.emb = nn.Linear(d_t1d*2+d_t2d, d_templ)
+        self.templ_stack = TemplatePairStack(n_block=n_block, d_templ=d_templ, n_head=n_head,
+                                             d_hidden=d_hidden, p_drop=p_drop)
+
+        self.attn = Attention(d_pair, d_templ, n_head, d_hidden, d_pair)
+
+        # process torsion angles
+        self.emb_t1d = nn.Linear(d_t1d+d_tor, d_templ)
+        self.proj_t1d = nn.Linear(d_templ, d_templ)
+        #self.tor_stack = TemplateTorsionStack(n_block=n_block, d_templ=d_templ, n_head=n_head,
+        #                                      d_hidden=d_hidden, p_drop=p_drop)
+        self.attn_tor = Attention(d_state, d_templ, n_head, d_hidden, d_state)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        self.emb = init_lecun_normal(self.emb)
+        nn.init.zeros_(self.emb.bias)
+
+        nn.init.kaiming_normal_(self.emb_t1d.weight, nonlinearity='relu')
+        nn.init.zeros_(self.emb_t1d.bias)
+
+        self.proj_t1d = init_lecun_normal(self.proj_t1d)
+        nn.init.zeros_(self.proj_t1d.bias)
+
+    def forward(self, t1d, t2d, alpha_t, xyz_t, pair, state, use_checkpoint=False):
+        # Input
+        #   - t1d: 1D template info (B, T, L, 23)
+        #   - t2d: 2D template info (B, T, L, L, 44)
+        B, T, L, _ = t1d.shape
+
+        # Prepare 2D template features
+        left = t1d.unsqueeze(3).expand(-1,-1,-1,L,-1)
+        right = t1d.unsqueeze(2).expand(-1,-1,L,-1,-1)
+        #
+        templ = torch.cat((t2d, left, right), -1) # (B, T, L, L, 90)
+        templ = self.emb(templ) # Template templures (B, T, L, L, d_templ)
+        # process each template features
+        xyz_t = xyz_t.reshape(B*T, L, -1, 3)
+        rbf_feat = rbf(torch.cdist(xyz_t[:,:,1], xyz_t[:,:,1]))
+        templ = self.templ_stack(templ, rbf_feat, use_checkpoint=use_checkpoint) # (B, T, L,L, d_templ)
+
+        # Prepare 1D template torsion angle features
+        t1d = torch.cat((t1d, alpha_t), dim=-1) # (B, T, L, 23+30)
+
+        # process each template features
+        t1d = self.proj_t1d(F.relu_(self.emb_t1d(t1d)))
+
+        # mixing query state features to template state features
+        state = state.reshape(B*L, 1, -1)
+        t1d = t1d.permute(0,2,1,3).reshape(B*L, T, -1)
+        if use_checkpoint:
+            out = checkpoint.checkpoint(create_custom_forward(self.attn_tor), state, t1d, t1d)
+            out = out.reshape(B, L, -1)
+        else:
+            out = self.attn_tor(state, t1d, t1d).reshape(B, L, -1)
+        state = state.reshape(B, L, -1)
+        state = state + out
+
+        # mixing query pair features to template information (Template pointwise attention)
+        pair = pair.reshape(B*L*L, 1, -1)
+        templ = templ.permute(0, 2, 3, 1, 4).reshape(B*L*L, T, -1)
+        if use_checkpoint:
+            out = checkpoint.checkpoint(create_custom_forward(self.attn), pair, templ, templ)
+            out = out.reshape(B, L, L, -1)
+        else:
+            out = self.attn(pair, templ, templ).reshape(B, L, L, -1)
+        #
+        pair = pair.reshape(B, L, L, -1)
+        pair = pair + out
+
+        return pair, state
+
+class Recycling(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, d_state=32):
+        super(Recycling, self).__init__()
+        self.proj_dist = nn.Linear(36+d_state*2, d_pair)
+        self.norm_state = nn.LayerNorm(d_state)
+        self.norm_pair = nn.LayerNorm(d_pair)
+        self.norm_msa = nn.LayerNorm(d_msa)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        self.proj_dist = init_lecun_normal(self.proj_dist)
+        nn.init.zeros_(self.proj_dist.bias)
+
+    def forward(self, seq, msa, pair, xyz, state):
+        B, L = pair.shape[:2]
+        state = self.norm_state(state)
+        #
+        left = state.unsqueeze(2).expand(-1,-1,L,-1)
+        right = state.unsqueeze(1).expand(-1,L,-1,-1)
+
+        # three anchor atoms
+        N  = xyz[:,:,0]
+        Ca = xyz[:,:,1]
+        C  = xyz[:,:,2]
+
+        # recreate Cb given N,Ca,C
+        b = Ca - N
+        c = C - Ca
+        a = torch.cross(b, c, dim=-1)
+        Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca
+
+        dist = rbf(torch.cdist(Cb, Cb))
+        dist = torch.cat((dist, left, right), dim=-1)
+        dist = self.proj_dist(dist)
+        pair = dist + self.norm_pair(pair)
+        msa = self.norm_msa(msa)
+        return msa, pair, state
+

rfdiffusion/RoseTTAFoldModel.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+from rfdiffusion.Embeddings import MSA_emb, Extra_emb, Templ_emb, Recycling
+from rfdiffusion.Track_module import IterativeSimulator
+from rfdiffusion.AuxiliaryPredictor import DistanceNetwork, MaskedTokenNetwork, ExpResolvedNetwork, LDDTNetwork
+from opt_einsum import contract as einsum
+
+class RoseTTAFoldModule(nn.Module):
+    def __init__(self, 
+                 n_extra_block, 
+                 n_main_block, 
+                 n_ref_block,
+                 d_msa,
+                 d_msa_full,
+                 d_pair,
+                 d_templ,
+                 n_head_msa,
+                 n_head_pair,
+                 n_head_templ,
+                 d_hidden,
+                 d_hidden_templ,
+                 p_drop,
+                 d_t1d,
+                 d_t2d,
+                 T, # total timesteps (used in timestep emb
+                 use_motif_timestep, # Whether to have a distinct emb for motif
+                 freeze_track_motif, # Whether to freeze updates to motif in track
+                 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},
+                 input_seq_onehot=False,     # For continuous vs. discrete sequence
+                 ):
+
+        super(RoseTTAFoldModule, self).__init__()
+
+        self.freeze_track_motif = freeze_track_motif
+
+        # 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, input_seq_onehot=input_seq_onehot) # Allowed to take onehotseq
+        self.full_emb = Extra_emb(d_msa=d_msa_full, d_init=25,
+                p_drop=p_drop, input_seq_onehot=input_seq_onehot) # Allowed to take onehotseq
+        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)
+        self.lddt_pred = LDDTNetwork(d_state)
+       
+        self.exp_pred = ExpResolvedNetwork(d_msa, d_state)
+
+    def forward(self, msa_latent, msa_full, seq, xyz, idx, t,
+                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, return_infer=False,
+                use_checkpoint=False, motif_mask=None, i_cycle=None, n_cycle=None):
+
+        B, N, L = msa_latent.shape[:3]
+        # Get embeddings
+        msa_latent, pair, state = self.latent_emb(msa_latent, seq, idx)
+        msa_full = self.full_emb(msa_full, seq, idx)
+
+        # Do recycling
+        if msa_prev == None:
+            msa_prev = torch.zeros_like(msa_latent[:,0])
+            pair_prev = torch.zeros_like(pair)
+            state_prev = torch.zeros_like(state)
+        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
+
+
+        # Get timestep embedding (if using)
+        if hasattr(self, 'timestep_embedder'):
+            assert t is not None
+            time_emb = self.timestep_embedder(L,t,motif_mask)
+            n_tmpl = t1d.shape[1]
+            t1d = torch.cat([t1d, time_emb[None,None,...].repeat(1,n_tmpl,1,1)], dim=-1)
+
+        # add template embedding
+        pair, state = self.templ_emb(t1d, t2d, alpha_t, xyz_t, pair, state, use_checkpoint=use_checkpoint)
+        
+        # Predict coordinates from given inputs
+        is_frozen_residue = motif_mask if self.freeze_track_motif else torch.zeros_like(motif_mask).bool()
+        msa, pair, R, T, alpha_s, state = self.simulator(seq, msa_latent, msa_full, pair, xyz[:,:,:3],
+                                                         state, idx, use_checkpoint=use_checkpoint,
+                                                         motif_mask=is_frozen_residue)
+        
+        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 LDDT
+        lddt = self.lddt_pred(state)
+
+        if return_infer:
+            # get last structure
+            xyz = einsum('bnij,bnaj->bnai', R[-1], xyz[:,:,:3]-xyz[:,:,1].unsqueeze(-2)) + T[-1].unsqueeze(-2)
+            
+            # get scalar plddt
+            nbin = lddt.shape[1]
+            bin_step = 1.0 / nbin
+            lddt_bins = torch.linspace(bin_step, 1.0, nbin, dtype=lddt.dtype, device=lddt.device)
+            pred_lddt = nn.Softmax(dim=1)(lddt)
+            pred_lddt = torch.sum(lddt_bins[None,:,None]*pred_lddt, dim=1)
+
+            return msa[:,0], pair, xyz, state, alpha_s[-1], logits_aa.permute(0,2,1), pred_lddt
+
+        #
+        # predict distogram & orientograms
+        logits = self.c6d_pred(pair)
+        
+        # predict experimentally resolved or not
+        logits_exp = self.exp_pred(msa[:,0], state)
+        
+        # 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

rfdiffusion/SE3_network.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+
+#from equivariant_attention.modules import get_basis_and_r, GSE3Res, GNormBias
+#from equivariant_attention.modules import GConvSE3, GNormSE3
+#from equivariant_attention.fibers import Fiber
+
+from rfdiffusion.util_module import init_lecun_normal_param
+from se3_transformer.model import SE3Transformer
+from se3_transformer.model.fiber import Fiber
+
+class SE3TransformerWrapper(nn.Module):
+    """SE(3) equivariant GCN with attention"""
+    def __init__(self, num_layers=2, num_channels=32, num_degrees=3, n_heads=4, div=4,
+                 l0_in_features=32, l0_out_features=32,
+                 l1_in_features=3, l1_out_features=2,
+                 num_edge_features=32):
+        super().__init__()
+        # Build the network
+        self.l1_in = l1_in_features
+        #
+        fiber_edge = Fiber({0: num_edge_features})
+        if l1_out_features > 0:
+            if l1_in_features > 0:
+                fiber_in = Fiber({0: l0_in_features, 1: l1_in_features})
+                fiber_hidden = Fiber.create(num_degrees, num_channels)
+                fiber_out = Fiber({0: l0_out_features, 1: l1_out_features})
+            else:
+                fiber_in = Fiber({0: l0_in_features})
+                fiber_hidden = Fiber.create(num_degrees, num_channels)
+                fiber_out = Fiber({0: l0_out_features, 1: l1_out_features})
+        else:
+            if l1_in_features > 0:
+                fiber_in = Fiber({0: l0_in_features, 1: l1_in_features})
+                fiber_hidden = Fiber.create(num_degrees, num_channels)
+                fiber_out = Fiber({0: l0_out_features})
+            else:
+                fiber_in = Fiber({0: l0_in_features})
+                fiber_hidden = Fiber.create(num_degrees, num_channels)
+                fiber_out = Fiber({0: l0_out_features})
+        
+        self.se3 = SE3Transformer(num_layers=num_layers,
+                                  fiber_in=fiber_in,
+                                  fiber_hidden=fiber_hidden,
+                                  fiber_out = fiber_out,
+                                  num_heads=n_heads,
+                                  channels_div=div,
+                                  fiber_edge=fiber_edge,
+                                  use_layer_norm=True)
+                                  #use_layer_norm=False)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+
+        # make sure linear layer before ReLu are initialized with kaiming_normal_
+        for n, p in self.se3.named_parameters():
+            if "bias" in n:
+                nn.init.zeros_(p)
+            elif len(p.shape) == 1:
+                continue
+            else:
+                if "radial_func" not in n:
+                    p = init_lecun_normal_param(p) 
+                else:
+                    if "net.6" in n:
+                        nn.init.zeros_(p)
+                    else:
+                        nn.init.kaiming_normal_(p, nonlinearity='relu')
+        
+        # make last layers to be zero-initialized
+        #self.se3.graph_modules[-1].to_kernel_self['0'] = init_lecun_normal_param(self.se3.graph_modules[-1].to_kernel_self['0'])
+        #self.se3.graph_modules[-1].to_kernel_self['1'] = init_lecun_normal_param(self.se3.graph_modules[-1].to_kernel_self['1'])
+        nn.init.zeros_(self.se3.graph_modules[-1].to_kernel_self['0'])
+        nn.init.zeros_(self.se3.graph_modules[-1].to_kernel_self['1'])
+
+    def forward(self, G, type_0_features, type_1_features=None, edge_features=None):
+        if self.l1_in > 0:
+            node_features = {'0': type_0_features, '1': type_1_features}
+        else:
+            node_features = {'0': type_0_features}
+        edge_features = {'0': edge_features}
+        return self.se3(G, node_features, edge_features)

rfdiffusion/Track_module.py

@@ -0,0 +1,474 @@
+import torch.utils.checkpoint as checkpoint
+from rfdiffusion.util_module import *
+from rfdiffusion.Attention_module import *
+from rfdiffusion.SE3_network import SE3TransformerWrapper
+
+# Components for three-track blocks
+# 1. MSA -> MSA update (biased attention. bias from pair & structure)
+# 2. Pair -> Pair update (biased attention. bias from structure)
+# 3. MSA -> Pair update (extract coevolution signal)
+# 4. Str -> Str update (node from MSA, edge from Pair)
+
+# Update MSA with biased self-attention. bias from Pair & Str
+class MSAPairStr2MSA(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, n_head=8, d_state=16,
+                 d_hidden=32, p_drop=0.15, use_global_attn=False):
+        super(MSAPairStr2MSA, self).__init__()
+        self.norm_pair = nn.LayerNorm(d_pair)
+        self.proj_pair = nn.Linear(d_pair+36, d_pair)
+        self.norm_state = nn.LayerNorm(d_state)
+        self.proj_state = nn.Linear(d_state, d_msa)
+        self.drop_row = Dropout(broadcast_dim=1, p_drop=p_drop)
+        self.row_attn = MSARowAttentionWithBias(d_msa=d_msa, d_pair=d_pair,
+                                                n_head=n_head, d_hidden=d_hidden) 
+        if use_global_attn:
+            self.col_attn = MSAColGlobalAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) 
+        else:
+            self.col_attn = MSAColAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) 
+        self.ff = FeedForwardLayer(d_msa, 4, p_drop=p_drop)
+        
+        # Do proper initialization
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # initialize weights to normal distrib
+        self.proj_pair = init_lecun_normal(self.proj_pair)
+        self.proj_state = init_lecun_normal(self.proj_state)
+
+        # initialize bias to zeros
+        nn.init.zeros_(self.proj_pair.bias)
+        nn.init.zeros_(self.proj_state.bias)
+
+    def forward(self, msa, pair, rbf_feat, state):
+        '''
+        Inputs:
+            - msa: MSA feature (B, N, L, d_msa)
+            - pair: Pair feature (B, L, L, d_pair)
+            - rbf_feat: Ca-Ca distance feature calculated from xyz coordinates (B, L, L, 36)
+            - xyz: xyz coordinates (B, L, n_atom, 3)
+            - state: updated node features after SE(3)-Transformer layer (B, L, d_state)
+        Output:
+            - msa: Updated MSA feature (B, N, L, d_msa)
+        '''
+        B, N, L = msa.shape[:3]
+
+        # prepare input bias feature by combining pair & coordinate info
+        pair = self.norm_pair(pair)
+        pair = torch.cat((pair, rbf_feat), dim=-1)
+        pair = self.proj_pair(pair) # (B, L, L, d_pair)
+        #
+        # update query sequence feature (first sequence in the MSA) with feedbacks (state) from SE3
+        state = self.norm_state(state)
+        state = self.proj_state(state).reshape(B, 1, L, -1)
+        msa = msa.index_add(1, torch.tensor([0,], device=state.device), state)
+        #
+        # Apply row/column attention to msa & transform 
+        msa = msa + self.drop_row(self.row_attn(msa, pair))
+        msa = msa + self.col_attn(msa)
+        msa = msa + self.ff(msa)
+
+        return msa
+
+class PairStr2Pair(nn.Module):
+    def __init__(self, d_pair=128, n_head=4, d_hidden=32, d_rbf=36, p_drop=0.15):
+        super(PairStr2Pair, self).__init__()
+        
+        self.emb_rbf = nn.Linear(d_rbf, d_hidden)
+        self.proj_rbf = nn.Linear(d_hidden, d_pair)
+
+        self.drop_row = Dropout(broadcast_dim=1, p_drop=p_drop)
+        self.drop_col = Dropout(broadcast_dim=2, p_drop=p_drop)
+
+        self.row_attn = BiasedAxialAttention(d_pair, d_pair, n_head, d_hidden, p_drop=p_drop, is_row=True)
+        self.col_attn = BiasedAxialAttention(d_pair, d_pair, n_head, d_hidden, p_drop=p_drop, is_row=False)
+
+        self.ff = FeedForwardLayer(d_pair, 2)
+        
+        self.reset_parameter()
+    
+    def reset_parameter(self):
+        nn.init.kaiming_normal_(self.emb_rbf.weight, nonlinearity='relu')
+        nn.init.zeros_(self.emb_rbf.bias)
+        
+        self.proj_rbf = init_lecun_normal(self.proj_rbf)
+        nn.init.zeros_(self.proj_rbf.bias)
+
+    def forward(self, pair, rbf_feat):
+        B, L = pair.shape[:2]
+
+        rbf_feat = self.proj_rbf(F.relu_(self.emb_rbf(rbf_feat)))
+
+        pair = pair + self.drop_row(self.row_attn(pair, rbf_feat))
+        pair = pair + self.drop_col(self.col_attn(pair, rbf_feat))
+        pair = pair + self.ff(pair)
+        return pair
+
+class MSA2Pair(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, d_hidden=32, p_drop=0.15):
+        super(MSA2Pair, self).__init__()
+        self.norm = nn.LayerNorm(d_msa)
+        self.proj_left = nn.Linear(d_msa, d_hidden)
+        self.proj_right = nn.Linear(d_msa, d_hidden)
+        self.proj_out = nn.Linear(d_hidden*d_hidden, d_pair)
+        
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # normal initialization
+        self.proj_left = init_lecun_normal(self.proj_left)
+        self.proj_right = init_lecun_normal(self.proj_right)
+        nn.init.zeros_(self.proj_left.bias)
+        nn.init.zeros_(self.proj_right.bias)
+
+        # zero initialize output
+        nn.init.zeros_(self.proj_out.weight)
+        nn.init.zeros_(self.proj_out.bias)
+
+    def forward(self, msa, pair):
+        B, N, L = msa.shape[:3]
+        msa = self.norm(msa)
+        left = self.proj_left(msa)
+        right = self.proj_right(msa)
+        right = right / float(N)
+        out = einsum('bsli,bsmj->blmij', left, right).reshape(B, L, L, -1)
+        out = self.proj_out(out)
+       
+        pair = pair + out
+        
+        return pair
+
+class SCPred(nn.Module):
+    def __init__(self, d_msa=256, d_state=32, d_hidden=128, p_drop=0.15):
+        super(SCPred, self).__init__()
+        self.norm_s0 = nn.LayerNorm(d_msa)
+        self.norm_si = nn.LayerNorm(d_state)
+        self.linear_s0 = nn.Linear(d_msa, d_hidden)
+        self.linear_si = nn.Linear(d_state, d_hidden)
+
+        # ResNet layers
+        self.linear_1 = nn.Linear(d_hidden, d_hidden)
+        self.linear_2 = nn.Linear(d_hidden, d_hidden)
+        self.linear_3 = nn.Linear(d_hidden, d_hidden)
+        self.linear_4 = nn.Linear(d_hidden, d_hidden)
+
+        # Final outputs
+        self.linear_out = nn.Linear(d_hidden, 20)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # normal initialization
+        self.linear_s0 = init_lecun_normal(self.linear_s0)
+        self.linear_si = init_lecun_normal(self.linear_si)
+        self.linear_out = init_lecun_normal(self.linear_out)
+        nn.init.zeros_(self.linear_s0.bias)
+        nn.init.zeros_(self.linear_si.bias)
+        nn.init.zeros_(self.linear_out.bias)
+        
+        # right before relu activation: He initializer (kaiming normal)
+        nn.init.kaiming_normal_(self.linear_1.weight, nonlinearity='relu')
+        nn.init.zeros_(self.linear_1.bias)
+        nn.init.kaiming_normal_(self.linear_3.weight, nonlinearity='relu')
+        nn.init.zeros_(self.linear_3.bias)
+
+        # right before residual connection: zero initialize
+        nn.init.zeros_(self.linear_2.weight)
+        nn.init.zeros_(self.linear_2.bias)
+        nn.init.zeros_(self.linear_4.weight)
+        nn.init.zeros_(self.linear_4.bias)
+    
+    def forward(self, seq, state):
+        '''
+        Predict side-chain torsion angles along with backbone torsions
+        Inputs:
+            - seq: hidden embeddings corresponding to query sequence (B, L, d_msa)
+            - state: state feature (output l0 feature) from previous SE3 layer (B, L, d_state)
+        Outputs:
+            - si: predicted torsion angles (phi, psi, omega, chi1~4 with cos/sin, Cb bend, Cb twist, CG) (B, L, 10, 2)
+        '''
+        B, L = seq.shape[:2]
+        seq = self.norm_s0(seq)
+        state = self.norm_si(state)
+        si = self.linear_s0(seq) + self.linear_si(state)
+
+        si = si + self.linear_2(F.relu_(self.linear_1(F.relu_(si))))
+        si = si + self.linear_4(F.relu_(self.linear_3(F.relu_(si))))
+
+        si = self.linear_out(F.relu_(si))
+        return si.view(B, L, 10, 2)
+
+
+class Str2Str(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, d_state=16, 
+            SE3_param={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, p_drop=0.1):
+        super(Str2Str, self).__init__()
+        
+        # initial node & pair feature process
+        self.norm_msa = nn.LayerNorm(d_msa)
+        self.norm_pair = nn.LayerNorm(d_pair)
+        self.norm_state = nn.LayerNorm(d_state)
+    
+        self.embed_x = nn.Linear(d_msa+d_state, SE3_param['l0_in_features'])
+        self.embed_e1 = nn.Linear(d_pair, SE3_param['num_edge_features'])
+        self.embed_e2 = nn.Linear(SE3_param['num_edge_features']+36+1, SE3_param['num_edge_features'])
+        
+        self.norm_node = nn.LayerNorm(SE3_param['l0_in_features'])
+        self.norm_edge1 = nn.LayerNorm(SE3_param['num_edge_features'])
+        self.norm_edge2 = nn.LayerNorm(SE3_param['num_edge_features'])
+        
+        self.se3 = SE3TransformerWrapper(**SE3_param)
+        self.sc_predictor = SCPred(d_msa=d_msa, d_state=SE3_param['l0_out_features'],
+                                   p_drop=p_drop)
+        
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # initialize weights to normal distribution
+        self.embed_x = init_lecun_normal(self.embed_x)
+        self.embed_e1 = init_lecun_normal(self.embed_e1)
+        self.embed_e2 = init_lecun_normal(self.embed_e2)
+
+        # initialize bias to zeros
+        nn.init.zeros_(self.embed_x.bias)
+        nn.init.zeros_(self.embed_e1.bias)
+        nn.init.zeros_(self.embed_e2.bias)
+    
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, top_k=64, eps=1e-5):
+        B, N, L = msa.shape[:3]
+
+        if motif_mask is None:
+            motif_mask = torch.zeros(L).bool()
+        
+        # process msa & pair features
+        node = self.norm_msa(msa[:,0])
+        pair = self.norm_pair(pair)
+        state = self.norm_state(state)
+       
+        node = torch.cat((node, state), dim=-1)
+        node = self.norm_node(self.embed_x(node))
+        pair = self.norm_edge1(self.embed_e1(pair))
+        
+        neighbor = get_seqsep(idx)
+        rbf_feat = rbf(torch.cdist(xyz[:,:,1], xyz[:,:,1]))
+        pair = torch.cat((pair, rbf_feat, neighbor), dim=-1)
+        pair = self.norm_edge2(self.embed_e2(pair))
+        
+        # define graph
+        if top_k != 0:
+            G, edge_feats = make_topk_graph(xyz[:,:,1,:], pair, idx, top_k=top_k)
+        else:
+            G, edge_feats = make_full_graph(xyz[:,:,1,:], pair, idx, top_k=top_k)
+        l1_feats = xyz - xyz[:,:,1,:].unsqueeze(2)
+        l1_feats = l1_feats.reshape(B*L, -1, 3)
+        
+        # apply SE(3) Transformer & update coordinates
+        shift = self.se3(G, node.reshape(B*L, -1, 1), l1_feats, edge_feats)
+
+        state = shift['0'].reshape(B, L, -1) # (B, L, C)
+        
+        offset = shift['1'].reshape(B, L, 2, 3)
+        offset[:,motif_mask,...] = 0            # NOTE: motif mask is all zeros if not freeezing the motif 
+
+        delTi = offset[:,:,0,:] / 10.0 # translation
+        R = offset[:,:,1,:] / 100.0 # rotation
+        
+        Qnorm = torch.sqrt( 1 + torch.sum(R*R, dim=-1) )
+        qA, qB, qC, qD = 1/Qnorm, R[:,:,0]/Qnorm, R[:,:,1]/Qnorm, R[:,:,2]/Qnorm
+
+        delRi = torch.zeros((B,L,3,3), device=xyz.device)
+        delRi[:,:,0,0] = qA*qA+qB*qB-qC*qC-qD*qD
+        delRi[:,:,0,1] = 2*qB*qC - 2*qA*qD
+        delRi[:,:,0,2] = 2*qB*qD + 2*qA*qC
+        delRi[:,:,1,0] = 2*qB*qC + 2*qA*qD
+        delRi[:,:,1,1] = qA*qA-qB*qB+qC*qC-qD*qD
+        delRi[:,:,1,2] = 2*qC*qD - 2*qA*qB
+        delRi[:,:,2,0] = 2*qB*qD - 2*qA*qC
+        delRi[:,:,2,1] = 2*qC*qD + 2*qA*qB
+        delRi[:,:,2,2] = qA*qA-qB*qB-qC*qC+qD*qD
+
+        Ri = einsum('bnij,bnjk->bnik', delRi, R_in)
+        Ti = delTi + T_in #einsum('bnij,bnj->bni', delRi, T_in) + delTi
+            
+        alpha = self.sc_predictor(msa[:,0], state)
+        return Ri, Ti, state, alpha
+
+class IterBlock(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128,
+                 n_head_msa=8, n_head_pair=4,
+                 use_global_attn=False,
+                 d_hidden=32, d_hidden_msa=None, p_drop=0.15,
+                 SE3_param={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}):
+        super(IterBlock, self).__init__()
+        if d_hidden_msa == None:
+            d_hidden_msa = d_hidden
+
+        self.msa2msa = MSAPairStr2MSA(d_msa=d_msa, d_pair=d_pair,
+                                      n_head=n_head_msa,
+                                      d_state=SE3_param['l0_out_features'],
+                                      use_global_attn=use_global_attn,
+                                      d_hidden=d_hidden_msa, p_drop=p_drop)
+        self.msa2pair = MSA2Pair(d_msa=d_msa, d_pair=d_pair,
+                                 d_hidden=d_hidden//2, p_drop=p_drop)
+                                 #d_hidden=d_hidden, p_drop=p_drop)
+        self.pair2pair = PairStr2Pair(d_pair=d_pair, n_head=n_head_pair, 
+                                      d_hidden=d_hidden, p_drop=p_drop)
+        self.str2str = Str2Str(d_msa=d_msa, d_pair=d_pair,
+                               d_state=SE3_param['l0_out_features'],
+                               SE3_param=SE3_param,
+                               p_drop=p_drop)
+
+    def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, use_checkpoint=False):
+        rbf_feat = rbf(torch.cdist(xyz[:,:,1,:], xyz[:,:,1,:]))
+        if use_checkpoint:
+            msa = checkpoint.checkpoint(create_custom_forward(self.msa2msa), msa, pair, rbf_feat, state)
+            pair = checkpoint.checkpoint(create_custom_forward(self.msa2pair), msa, pair)
+            pair = checkpoint.checkpoint(create_custom_forward(self.pair2pair), pair, rbf_feat)
+            R, T, state, alpha = checkpoint.checkpoint(create_custom_forward(self.str2str, top_k=0), msa, pair, R_in, T_in, xyz, state, idx, motif_mask)
+        else:
+            msa = self.msa2msa(msa, pair, rbf_feat, state)
+            pair = self.msa2pair(msa, pair)
+            pair = self.pair2pair(pair, rbf_feat)
+            R, T, state, alpha = self.str2str(msa, pair, R_in, T_in, xyz, state, idx, motif_mask=motif_mask, top_k=0) 
+        
+        return msa, pair, R, T, state, alpha
+
+class IterativeSimulator(nn.Module):
+    def __init__(self, n_extra_block=4, n_main_block=12, n_ref_block=4,
+                 d_msa=256, d_msa_full=64, d_pair=128, d_hidden=32,
+                 n_head_msa=8, n_head_pair=4,
+                 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},
+                 p_drop=0.15):
+        super(IterativeSimulator, self).__init__()
+        self.n_extra_block = n_extra_block
+        self.n_main_block = n_main_block
+        self.n_ref_block = n_ref_block
+        
+        self.proj_state = nn.Linear(SE3_param_topk['l0_out_features'], SE3_param_full['l0_out_features'])
+        # Update with extra sequences
+        if n_extra_block > 0:
+            self.extra_block = nn.ModuleList([IterBlock(d_msa=d_msa_full, d_pair=d_pair,
+                                                        n_head_msa=n_head_msa,
+                                                        n_head_pair=n_head_pair,
+                                                        d_hidden_msa=8,
+                                                        d_hidden=d_hidden,
+                                                        p_drop=p_drop,
+                                                        use_global_attn=True,
+                                                        SE3_param=SE3_param_full)
+                                                        for i in range(n_extra_block)])
+
+        # Update with seed sequences
+        if n_main_block > 0:
+            self.main_block = nn.ModuleList([IterBlock(d_msa=d_msa, d_pair=d_pair,
+                                                       n_head_msa=n_head_msa,
+                                                       n_head_pair=n_head_pair,
+                                                       d_hidden=d_hidden,
+                                                       p_drop=p_drop,
+                                                       use_global_attn=False,
+                                                       SE3_param=SE3_param_full)
+                                                       for i in range(n_main_block)])
+
+        self.proj_state2 = nn.Linear(SE3_param_full['l0_out_features'], SE3_param_topk['l0_out_features'])
+        # Final SE(3) refinement
+        if n_ref_block > 0:
+            self.str_refiner = Str2Str(d_msa=d_msa, d_pair=d_pair,
+                                       d_state=SE3_param_topk['l0_out_features'],
+                                       SE3_param=SE3_param_topk,
+                                       p_drop=p_drop)
+    
+        self.reset_parameter()
+    def reset_parameter(self):
+        self.proj_state = init_lecun_normal(self.proj_state)
+        nn.init.zeros_(self.proj_state.bias)
+        self.proj_state2 = init_lecun_normal(self.proj_state2)
+        nn.init.zeros_(self.proj_state2.bias)
+
+    def forward(self, seq, msa, msa_full, pair, xyz_in, state, idx, use_checkpoint=False, motif_mask=None):
+        """
+        input:
+           seq: query sequence (B, L)
+           msa: seed MSA embeddings (B, N, L, d_msa)
+           msa_full: extra MSA embeddings (B, N, L, d_msa_full)
+           pair: initial residue pair embeddings (B, L, L, d_pair)
+           xyz_in: initial BB coordinates (B, L, n_atom, 3)
+           state: initial state features containing mixture of query seq, sidechain, accuracy info (B, L, d_state)
+           idx: residue index
+           motif_mask: bool tensor, True if motif position that is frozen, else False(L,) 
+        """
+
+        B, L = pair.shape[:2]
+
+        if motif_mask is None:
+            motif_mask = torch.zeros(L).bool()
+
+        R_in = torch.eye(3, device=xyz_in.device).reshape(1,1,3,3).expand(B, L, -1, -1)
+        T_in = xyz_in[:,:,1].clone()
+        xyz_in = xyz_in - T_in.unsqueeze(-2)
+        
+        state = self.proj_state(state)
+
+        R_s = list()
+        T_s = list()
+        alpha_s = list()
+        for i_m in range(self.n_extra_block):
+            R_in = R_in.detach() # detach rotation (for stability)
+            T_in = T_in.detach()
+            # Get current BB structure
+            xyz = einsum('bnij,bnaj->bnai', R_in, xyz_in) + T_in.unsqueeze(-2)
+
+            msa_full, pair, R_in, T_in, state, alpha = self.extra_block[i_m](msa_full, 
+                                                                             pair,
+                                                                             R_in, 
+                                                                             T_in, 
+                                                                             xyz, 
+                                                                             state, 
+                                                                             idx,
+                                                                             motif_mask=motif_mask,
+                                                                             use_checkpoint=use_checkpoint)
+            R_s.append(R_in)
+            T_s.append(T_in)
+            alpha_s.append(alpha)
+
+        for i_m in range(self.n_main_block):
+            R_in = R_in.detach()
+            T_in = T_in.detach()
+            # Get current BB structure
+            xyz = einsum('bnij,bnaj->bnai', R_in, xyz_in) + T_in.unsqueeze(-2)
+            
+            msa, pair, R_in, T_in, state, alpha = self.main_block[i_m](msa, 
+                                                                       pair,
+                                                                       R_in, 
+                                                                       T_in, 
+                                                                       xyz, 
+                                                                       state, 
+                                                                       idx,
+                                                                       motif_mask=motif_mask,
+                                                                       use_checkpoint=use_checkpoint)
+            R_s.append(R_in)
+            T_s.append(T_in)
+            alpha_s.append(alpha)
+       
+        state = self.proj_state2(state)
+        for i_m in range(self.n_ref_block):
+            R_in = R_in.detach()
+            T_in = T_in.detach()
+            xyz = einsum('bnij,bnaj->bnai', R_in, xyz_in) + T_in.unsqueeze(-2)
+            R_in, T_in, state, alpha = self.str_refiner(msa, 
+                                                        pair, 
+                                                        R_in, 
+                                                        T_in, 
+                                                        xyz, 
+                                                        state, 
+                                                        idx, 
+                                                        top_k=64, 
+                                                        motif_mask=motif_mask)
+            R_s.append(R_in)
+            T_s.append(T_in)
+            alpha_s.append(alpha)
+
+        R_s = torch.stack(R_s, dim=0)
+        T_s = torch.stack(T_s, dim=0)
+        alpha_s = torch.stack(alpha_s, dim=0)
+
+        return msa, pair, R_s, T_s, alpha_s, state