new b

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 RosettaCommons
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,17 @@
+---
+title: PROTEIN GENERATOR
+emoji: 🧪
+thumbnail: http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landscape.png
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 3.24.1
+app_file: app.py
+pinned: false
+---
+
+![fig](./utils/figs/thumbnail_github.png)
+
+## Code Accessibility
+
+To download code and for more details please visit the [github](https://github.com/RosettaCommons/protein_generator)!

app.py

@@ -0,0 +1,404 @@
+import os,sys
+
+# install environment goods
+#os.system("pip -q install dgl -f https://data.dgl.ai/wheels/cu113/repo.html")
+os.system('pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html')
+#os.system('pip install gradio')
+os.environ["DGLBACKEND"] = "pytorch"
+#os.system(f'pip install -r ./PROTEIN_GENERATOR/requirements.txt')
+print('Modules installed')
+
+os.system('pip install --force gradio==3.28.3')
+
+os.environ["DGLBACKEND"] = "pytorch"
+
+if not os.path.exists('./SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'):
+    print('Downloading model weights 1')
+    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
+    print('Successfully Downloaded')
+
+if not os.path.exists('./SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'):
+    print('Downloading model weights 2')
+    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
+    print('Successfully Downloaded')
+
+import numpy as np
+import gradio as gr
+import py3Dmol
+from io import StringIO
+import json
+import secrets
+import copy
+import matplotlib.pyplot as plt
+from utils.sampler import HuggingFace_sampler
+
+plt.rcParams.update({'font.size': 13})
+
+with open('./tmp/args.json','r') as f:
+    args = json.load(f)
+
+# manually set checkpoint to load
+args['checkpoint'] = None
+args['dump_trb'] = False
+args['dump_args'] = True
+args['save_best_plddt'] = True
+args['T'] = 25
+args['strand_bias'] = 0.0
+args['loop_bias'] = 0.0
+args['helix_bias'] = 0.0
+
+
+
+def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias, 
+                    secondary_structure, aa_bias, aa_bias_potential, 
+                    #target_charge, target_ph, charge_potential, 
+                    num_steps, noise, hydrophobic_target_score, hydrophobic_potential):
+    
+    dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
+    og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'
+
+    model_args = copy.deepcopy(args)
+
+    # make sampler
+    S = HuggingFace_sampler(args=model_args)
+
+    # get random prefix 
+    S.out_prefix = './tmp/'+secrets.token_hex(nbytes=10).upper()
+
+    # set args
+    S.args['checkpoint'] = None
+    S.args['dump_trb'] = False
+    S.args['dump_args'] = True
+    S.args['save_best_plddt'] = True
+    S.args['T'] = 20
+    S.args['strand_bias'] = 0.0
+    S.args['loop_bias'] = 0.0
+    S.args['helix_bias'] = 0.0
+    S.args['potentials'] = None
+    S.args['potential_scale'] = None
+    S.args['aa_composition'] = None
+
+
+    # get sequence if entered and make sure all chars are valid
+    alt_aa_dict = {'B':['D','N'],'J':['I','L'],'U':['C'],'Z':['E','Q'],'O':['K']}
+    if sequence not in ['',None]:
+        L = len(sequence)
+        aa_seq = []
+        for aa in sequence.upper():
+            if aa in alt_aa_dict.keys():
+                aa_seq.append(np.random.choice(alt_aa_dict[aa]))
+            else:
+                aa_seq.append(aa)
+
+        S.args['sequence'] = aa_seq
+    else:
+        S.args['contigs'] = [f'{seq_len}']
+        L = int(seq_len)
+    
+    if secondary_structure in ['',None]:
+        secondary_structure = None
+    else:
+        secondary_structure = ''.join(['E' if x == 'S' else x for x in secondary_structure])
+        if L < len(secondary_structure):
+            secondary_structure = secondary_structure[:len(sequence)]
+        elif L == len(secondary_structure):
+            pass
+        else:
+            dseq = L - len(secondary_structure)
+            secondary_structure += secondary_structure[-1]*dseq
+    
+
+    # potentials
+    potential_list = []
+    potential_bias_list = []
+
+    if aa_bias not in ['',None]:
+        potential_list.append('aa_bias')
+        S.args['aa_composition'] = aa_bias
+        if aa_bias_potential in ['',None]:
+            aa_bias_potential = 3
+        potential_bias_list.append(str(aa_bias_potential))
+    '''
+    if target_charge not in ['',None]:
+        potential_list.append('charge')
+        if charge_potential in ['',None]:
+            charge_potential = 1
+        potential_bias_list.append(str(charge_potential))
+        S.args['target_charge'] = float(target_charge)
+        if target_ph in ['',None]:
+            target_ph = 7.4
+        S.args['target_pH'] = float(target_ph)
+    '''
+    
+    if hydrophobic_target_score not in ['',None]:
+        potential_list.append('hydrophobic')
+        S.args['hydrophobic_score'] = float(hydrophobic_target_score)
+        if hydrophobic_potential in ['',None]:
+            hydrophobic_potential = 3
+        potential_bias_list.append(str(hydrophobic_potential))
+    
+
+    if len(potential_list) > 0:
+        S.args['potentials'] = ','.join(potential_list)
+        S.args['potential_scale'] = ','.join(potential_bias_list)
+
+
+    # normalise secondary_structure bias from range 0-0.3
+    S.args['secondary_structure'] = secondary_structure
+    S.args['helix_bias'] = helix_bias
+    S.args['strand_bias'] = strand_bias
+    S.args['loop_bias'] = loop_bias
+    
+    # set T
+    if num_steps in ['',None]:
+        S.args['T'] = 20
+    else:
+        S.args['T'] = int(num_steps)
+
+    # noise
+    if 'normal' in noise:
+        S.args['sample_distribution'] = noise
+        S.args['sample_distribution_gmm_means'] = [0]
+        S.args['sample_distribution_gmm_variances'] = [1]
+    elif 'gmm2' in noise:
+        S.args['sample_distribution'] = noise
+        S.args['sample_distribution_gmm_means'] = [-1,1]
+        S.args['sample_distribution_gmm_variances'] = [1,1]
+    elif 'gmm3' in noise:
+        S.args['sample_distribution'] = noise
+        S.args['sample_distribution_gmm_means'] = [-1,0,1]
+        S.args['sample_distribution_gmm_variances'] = [1,1,1]
+
+
+
+    if secondary_structure not in ['',None] or helix_bias+strand_bias+loop_bias > 0:
+        S.args['checkpoint'] = dssp_checkpoint
+        S.args['d_t1d'] = 29
+        print('using dssp checkpoint')
+    else:
+        S.args['checkpoint'] = og_checkpoint
+        S.args['d_t1d'] = 24
+        print('using og checkpoint')
+    
+
+    for k,v in S.args.items():
+        print(f"{k} --> {v}")
+    
+    # init S
+    S.model_init()
+    S.diffuser_init()
+    S.setup()
+
+    # sampling loop
+    plddt_data = []
+    for j in range(S.max_t):
+        output_seq, output_pdb, plddt = S.take_step_get_outputs(j)
+        plddt_data.append(plddt)
+        yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
+    
+    output_seq, output_pdb, plddt = S.get_outputs()
+
+    yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
+
+def get_plddt_plot(plddt_data, max_t):
+    x = [i+1 for i in range(len(plddt_data))]
+    fig, ax = plt.subplots(figsize=(15,6))
+    ax.plot(x,plddt_data,color='#661dbf', linewidth=3,marker='o')
+    ax.set_xticks([i+1 for i in range(max_t)])
+    ax.set_yticks([(i+1)/10 for i in range(10)])
+    ax.set_ylim([0,1])
+    ax.set_ylabel('model confidence (plddt)')
+    ax.set_xlabel('diffusion steps (t)')
+    return fig
+
+def display_pdb(path_to_pdb):
+    '''
+        #function to display pdb in py3dmol
+    '''
+    pdb = open(path_to_pdb, "r").read()
+    
+    view = py3Dmol.view(width=500, height=500)
+    view.addModel(pdb, "pdb")
+    view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})#'linear', 'min': 0, 'max': 1, 'colors': ["#ff9ef0","#a903fc",]}}}) 
+    view.zoomTo()
+    output = view._make_html().replace("'", '"')
+    print(view._make_html())
+    x = f"""<!DOCTYPE html><html></center> {output} </center></html>"""  # do not use ' in this input
+    
+    return f"""<iframe height="500px" width="100%"  name="result" allow="midi; geolocation; microphone; camera;
+                            display-capture; encrypted-media;" sandbox="allow-modals allow-forms
+                            allow-scripts allow-same-origin allow-popups
+                            allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
+                            allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+
+'''
+
+    return f"""<iframe  style="width: 100%; height:700px" name="result" allow="midi; geolocation; microphone; camera; 
+                            display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
+                            allow-scripts allow-same-origin allow-popups 
+                            allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+                            allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+'''
+
+def toggle_seq_input(choice):
+    if choice == "protein length":
+        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
+    elif choice == "custom sequence":
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
+
+def toggle_secondary_structure(choice):
+    if choice == "sliders":
+        return gr.update(visible=True, value=None),gr.update(visible=True, value=None),gr.update(visible=True, value=None),gr.update(visible=False, value=None)
+    elif choice == "explicit":
+        return gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=True, value=None)
+
+# Define the Gradio interface
+with gr.Blocks(theme='ParityError/Interstellar') as demo:
+    
+    gr.Markdown(f"""# Protein Generation via Diffusion in Sequence Space""")
+
+    with gr.Row():
+        with gr.Column(min_width=500):
+            gr.Markdown(f"""
+                    ## How does it work?\n
+                    --- [PREPRINT](https://biorxiv.org/content/10.1101/2023.05.08.539766v1) ---
+
+                    Protein sequence and structure co-generation is a long outstanding problem in the field of protein design. By implementing [ddpm](https://arxiv.org/abs/2006.11239) style diffusion over protein seqeuence space we generate protein sequence and structure pairs. Starting with [RoseTTAFold](https://www.science.org/doi/10.1126/science.abj8754), a protein structure prediction network, we finetuned it to predict sequence and structure given a partially noised sequence. By applying losses to both the predicted sequence and structure the model is forced to generate meaningful pairs. Diffusing in sequence space makes it easy to implement potentials to guide the diffusive process toward particular amino acid composition, net charge, and more! Furthermore, you can sample proteins from a family of sequences or even train a small sequence to function classifier to guide generation toward desired sequences.
+                    ![fig1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landscape.png)
+                    
+                    ## How to use it?\n
+                    A user can either design a custom input sequence to diffuse from or specify a length below. To scaffold a sequence use the following format where X represent residues to diffuse: XXXXXXXXSCIENCESCIENCEXXXXXXXXXXXXXXXXXXX. You can even design a protein with your name XXXXXXXXXXXXNAMEHEREXXXXXXXXXXXXX!
+                    
+                    ### Acknowledgements\n
+                    Thank you to Simon Dürr and the Hugging Face team for setting us up with a community GPU grant!
+                    """)
+            
+        gr.Markdown("""
+        ## Model in Action
+        ![gif1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/seqdiff_anim_720p.gif)
+        """)
+
+    with gr.Row().style(equal_height=False):
+        with gr.Column():
+            gr.Markdown("""## INPUTS""")
+            gr.Markdown("""#### Start Sequence
+                        Specify the protein length for complete unconditional generation, or scaffold a motif (or your name) using the custom sequence input""")
+            seq_opt = gr.Radio(["protein length","custom sequence"], label="How would you like to specify the starting sequence?", value='protein length')
+
+            sequence = gr.Textbox(label="custom sequence", lines=1, placeholder='AMINO ACIDS: A,C,D,E,F,G,H,I,K,L,M,N,P,Q,R,S,T,V,W,Y\n  MASK TOKEN: X', visible=False)
+            seq_len = gr.Slider(minimum=5.0, maximum=250.0, label="protein length", value=100, visible=True)
+            
+            seq_opt.change(fn=toggle_seq_input,
+                                    inputs=[seq_opt],
+                                    outputs=[seq_len, sequence],
+                                    queue=False)
+
+            gr.Markdown("""### Optional Parameters""")
+            with gr.Accordion(label='Secondary Structure',open=True):
+                gr.Markdown("""Try changing the sliders or inputing explicit secondary structure conditioning for each residue""")
+                sec_str_opt = gr.Radio(["sliders","explicit"], label="How would you like to specify secondary structure?", value='sliders')
+
+                secondary_structure = gr.Textbox(label="secondary structure", lines=1, placeholder='HELIX = H  STRAND = S  LOOP = L  MASK = X(must be the same length as input sequence)', visible=False)
+                
+                with gr.Column():
+                    helix_bias = gr.Slider(minimum=0.0, maximum=0.05, label="helix bias", visible=True)
+                    strand_bias = gr.Slider(minimum=0.0, maximum=0.05, label="strand bias", visible=True)
+                    loop_bias = gr.Slider(minimum=0.0, maximum=0.20, label="loop bias", visible=True)
+            
+                sec_str_opt.change(fn=toggle_secondary_structure,
+                                        inputs=[sec_str_opt],
+                                        outputs=[helix_bias,strand_bias,loop_bias,secondary_structure],
+                                        queue=False)
+                
+            with gr.Accordion(label='Amino Acid Compositional Bias',open=False):
+                gr.Markdown("""Bias sequence composition for particular amino acids by specifying the one letter code followed by the fraction to bias. This can be input as a list for example: W0.2,E0.1""")
+                with gr.Row():
+                    aa_bias = gr.Textbox(label="aa bias", lines=1, placeholder='specify one letter AA and fraction to bias, for example W0.1 or M0.1,K0.1' )
+                    aa_bias_potential = gr.Textbox(label="aa bias scale", lines=1, placeholder='AA Bias potential scale (recomended range 1.0-5.0)')
+            
+            '''
+            with gr.Accordion(label='Charge Bias',open=False):
+                gr.Markdown("""Bias for a specified net charge at a particular pH using the boxes below""")
+                with gr.Row():
+                    target_charge = gr.Textbox(label="net charge", lines=1, placeholder='net charge to target')
+                    target_ph = gr.Textbox(label="pH", lines=1, placeholder='pH at which net charge is desired')
+                    charge_potential = gr.Textbox(label="charge potential scale", lines=1, placeholder='charge potential scale (recomended range 1.0-5.0)')
+            '''
+
+            with gr.Accordion(label='Hydrophobic Bias',open=False):
+                gr.Markdown("""Bias for or against hydrophobic composition, to get more soluble proteins, bias away with a negative target score (ex. -5)""")
+                with gr.Row():
+                    hydrophobic_target_score = gr.Textbox(label="hydrophobic score", lines=1, placeholder='hydrophobic score to target (negative score is good for solublility)')
+                    hydrophobic_potential = gr.Textbox(label="hydrophobic potential scale", lines=1, placeholder='hydrophobic potential scale (recomended range 1.0-2.0)')
+            
+            with gr.Accordion(label='Diffusion Params',open=False):
+                gr.Markdown("""Increasing T to more steps can be helpful for harder design challenges, sampling from different distributions can change the sequence and structural composition""")
+                with gr.Row():
+                    num_steps = gr.Textbox(label="T", lines=1, placeholder='number of diffusion steps (25 or less will speed things up)')
+                    noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal') 
+
+            btn = gr.Button("GENERATE")
+
+        #with gr.Row():
+        with gr.Column():
+            gr.Markdown("""## OUTPUTS""")
+            gr.Markdown("""#### Confidence score for generated structure at each timestep""")
+            plddt_plot = gr.Plot(label='plddt at step t')
+            gr.Markdown("""#### Output protein sequnece""")
+            output_seq = gr.Textbox(label="sequence")
+            gr.Markdown("""#### Download PDB file""")
+            output_pdb = gr.File(label="PDB file")
+            gr.Markdown("""#### Structure viewer""")
+            output_viewer = gr.HTML()
+
+    gr.Markdown("""### Don't know where to get started? Click on an example below to try it out!""")
+    gr.Examples(
+        [["","125",0.0,0.0,0.2,"","","","20","normal",'',''],
+         ["","100",0.0,0.0,0.0,"","W0.2","2","20","normal",'',''],
+         ["","100",0.0,0.0,0.0,"XXHHHHHHHHHXXXXXXXHHHHHHHHHXXXXXXXHHHHHHHHXXXXSSSSSSSSSSSXXXXXXXXSSSSSSSSSSSSXXXXXXXSSSSSSSSSXXXXXXX","","","25","normal",'',''],
+         ["XXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXXXXXPEPSEQXXXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXX","",0.0,0.0,0.0,"","","","25","normal",'','']],
+        inputs=[sequence, 
+                seq_len, 
+                helix_bias, 
+                strand_bias, 
+                loop_bias, 
+                secondary_structure, 
+                aa_bias, 
+                aa_bias_potential,
+                #target_charge, 
+                #target_ph, 
+                #charge_potential, 
+                num_steps, 
+                noise, 
+                hydrophobic_target_score, 
+                hydrophobic_potential],
+        outputs=[output_seq,
+                 output_pdb,
+                 output_viewer,
+                 plddt_plot],
+        fn=protein_diffusion_model,
+        ) 
+    btn.click(protein_diffusion_model, 
+                [sequence, 
+                 seq_len, 
+                 helix_bias, 
+                 strand_bias, 
+                 loop_bias, 
+                 secondary_structure, 
+                 aa_bias, 
+                 aa_bias_potential, 
+                 #target_charge, 
+                 #target_ph, 
+                 #charge_potential, 
+                 num_steps, 
+                 noise, 
+                 hydrophobic_target_score, 
+                 hydrophobic_potential], 
+                [output_seq,
+                 output_pdb,
+                 output_viewer,
+                 plddt_plot])
+
+demo.queue()
+demo.launch(debug=True)
+

examples/aa_weights.json

@@ -0,0 +1,22 @@
+{
+    "A": 0,
+    "R": 0,
+    "N": 0,
+    "D": 0,
+    "C": 0,
+    "Q": 0,
+    "E": 0,
+    "G": 0,
+    "H": 0,
+    "I": 0,
+    "L": 0,
+    "K": 0,
+    "M": 0,
+    "F": 0,
+    "P": 0,
+    "S": 0,
+    "T": 0,
+    "W": 0,
+    "Y": 0,
+    "V": 0 
+}

examples/binder_design.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/binder_design \
+    --pdb pdbs/cd86.pdb \
+    --T 25 --save_best_plddt \
+    --contigs B1-110,0 25-75 \
+    --hotspots B40,B32,B87,B96,B30

examples/loop_design.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --pdb pdbs/G12D_manual_mut.pdb \
+    --out out/ab_loop \
+    --contigs A2-176,0 C7-16,0 H2-95,12-15,H111-116,0 L1-45,10-12,L56-107 \
+    --T 25 --save_best_plddt --loop_design

examples/motif_scaffolding.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --pdb pdbs/rsv5_5tpn.pdb \
+    --contigs 0-25,A163-181,25-30 --T 25 --save_best_plddt

examples/partial_diffusion.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --pdb out/design_000.pdb \
+    --trb out/design_000.trb \
+    --out out/partial_diffusion_design \
+    --contigs 0 --sampling_temp 0.3 --T 50 --save_best_plddt

examples/secondary_structure.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --contigs 100 \
+    --T 25 --save_best_plddt \
+    --secondary_structure XXXXXHHHHXXXLLLXXXXXXXXXXHHHHXXXLLLXXXXXXXXXXHHHHXXXLLLXXXXXXXXXXHHHHXXXLLLXXXXXXXXXXHHHHXXXLLLXXXXX
+
+# FOR SECONDARY STRUCTURE:
+#   X - mask
+#   H - helix
+#   E - strand
+#   L - loop

examples/secondary_structure_bias.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --contigs 100 \
+    --T 25 --save_best_plddt \
+    --helix_bias 0.01 --strand_bias 0.01 --loop_bias 0.0 

examples/secondary_structure_from_pdb.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --contigs 110 \
+    --T 25 --save_best_plddt \
+    --dssp_pdb ./pdbs/cd86.pdb
+
+# FOR SECONDARY STRUCTURE:
+#   X - mask
+#   H - helix
+#   E - strand
+#   L - loop

examples/symmetric_design.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/symmetric_design \
+    --contigs 25,0 25,0 25,0 \
+    --T 50 \
+    --save_best_plddt \
+    --symmetry 3

examples/weighted_sequence.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --contigs 100 \
+    --T 25 --save_best_plddt \
+    --aa_composition W0.2 --potential_scale 1.75 

examples/weighted_sequence_json.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH -J seq_diff
+#SBATCH -p gpu
+#SBATCH --mem=8g
+#SBATCH --gres=gpu:a6000:1
+#SBATCH -o ./out/slurm/slurm_%j.out
+
+source activate /software/conda/envs/SE3nv
+
+srun python ../inference.py \
+    --num_designs 10 \
+    --out out/design \
+    --contigs 75 \
+    --aa_weights_json aa_weights.json \
+    --add_weight_every_n 5 --add_weight_every_n \ 
+    --T 25 --save_best_plddt

model/.ipynb_checkpoints/RoseTTAFoldModel-checkpoint.py

@@ -0,0 +1,140 @@
+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

model/Attention_module.py

@@ -0,0 +1,411 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from opt_einsum import contract as einsum
+from util_module import init_lecun_normal
+from icecream import ic
+
+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, p_drop=0.1):
+        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, same_chain = None):
+        # 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)
+
+        if same_chain is not None:
+            ic(same_chain)
+            ic(attn)
+            ic(attn[~same_chain])
+            attn[~same_chain] *= 1.1
+        
+        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
+

model/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, p_drop=0.1):
+        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)
+
+
+

model/Embeddings.py

@@ -0,0 +1,307 @@
+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 util import get_tips
+from util_module import Dropout, create_custom_forward, rbf, init_lecun_normal
+from Attention_module import Attention, FeedForwardLayer, AttentionWithBias
+from Track_module import PairStr2Pair
+from icecream import ic
+
+# 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):
+        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.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, seq1hot=None):
+        # 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
+        seq = seq.long()
+        tmp = self.emb_q(seq).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 
+        if seq1hot is not None:
+            left = (seq1hot @ self.emb_left.weight)[:,None] # (B, 1, L, d_pair)
+            right = (seq1hot @ self.emb_right.weight)[:,:,None] # (B, L, 1, d_pair)
+        else:
+            left = self.emb_left(seq)[:,None] # (B, 1, L, d_pair)
+            right = self.emb_right(seq)[:,:,None] # (B, L, 1, d_pair)
+        #ic(torch.norm(self.emb_left.weight, dim=1))
+        #ic(torch.norm(self.emb_right.weight, dim=1))
+        pair = left + right # (B, L, L, d_pair)
+        pair = self.pos(pair, idx) # add relative position
+
+        # state embedding
+        state = self.drop(self.emb_state(seq))
+
+        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):
+        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.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, seq1hot=None):
+        # 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
+        if seq1hot is not None:
+            seq = (seq1hot @ self.emb_q.weight).unsqueeze(1) # (B, 1, L, d_model) -- query embedding
+        else:
+            seq = self.emb_q(seq).unsqueeze(1) # (B, 1, L, d_model) -- query embedding
+        #ic(torch.norm(self.emb_q.weight, dim=1))
+        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)
+    #   - seq confidence (1)
+    #   - global time step (1)
+    #   - struc confidence (1)
+    #   
+    def __init__(self, d_t1d=21+1+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, p_drop=p_drop)
+        
+        # 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, p_drop=p_drop)
+
+        self.reset_parameter()
+    
+    def reset_parameter(self):
+        self.emb = init_lecun_normal(self.emb)
+        #nn.init.zeros_(self.emb.weight) #init weights to zero
+        nn.init.zeros_(self.emb.bias)
+
+        nn.init.kaiming_normal_(self.emb_t1d.weight, nonlinearity='relu')
+        #nn.init.zeros_(self.emb_t1d.weight)
+        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) 24 SL
+        #   - 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, 88)
+        
+        #ic(templ.shape)
+        #ic(templ.dtype)
+        #ic(self.emb.weight.dtype)
+        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, 22+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
+

model/RoseTTAFoldModel.py

@@ -0,0 +1,140 @@
+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

model/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 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)

model/Track_module.py

@@ -0,0 +1,476 @@
+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 util import cross_product_matrix
+from util_module import *
+from Attention_module import *
+from SE3_network import SE3TransformerWrapper
+from icecream import ic
+
+# 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.type(torch.float32))
+        #
+        # 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, top_k=64, eps=1e-5):
+        B, N, L = msa.shape[:3]
+        
+        state = state.type(torch.float32)
+        mas = msa.type(torch.float32)
+        pair = pair.type(torch.float32)
+        R_in = R_in.type(torch.float32)
+        T_in = T_in.type(torch.float32)
+        xyz = xyz.type(torch.float32)
+        
+        #ic(msa.dtype)
+        #ic(pair.dtype)
+        #ic(R_in.dtype)
+        #ic(T_in.dtype)
+        #ic(xyz.dtype)
+        #ic(state.dtype)
+        #ic(idx.dtype)
+        
+
+        # 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)
+        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
+        #
+        ## convert vector to rotation matrix
+        #R_angle = torch.norm(R, dim=-1, keepdim=True) # (B, L, 1)
+        #cos_angle = torch.cos(R_angle).unsqueeze(2) # (B, L, 1, 1)
+        #sin_angle = torch.sin(R_angle).unsqueeze(2) # (B, L, 1, 1)
+        #R_vector = R / (R_angle+eps) # (B, L, 3)
+
+        #delRi = cos_angle*torch.eye(3, device=R.device).reshape(1,1,3,3) \
+        #      + sin_angle*cross_product_matrix(R_vector) \
+        #      + (1.0-cos_angle)*einsum('bni,bnj->bnij', R_vector, R_vector)
+
+        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, 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)
+        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, 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):
+        # 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
+
+        B, L = pair.shape[:2]
+
+        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,
+                                                                             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,
+                                                                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)
+            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

model/apply_masks.py

@@ -0,0 +1,196 @@
+import sys, os
+import torch
+from icecream import ic
+import random
+import numpy as np
+from kinematics import get_init_xyz
+sys.path.append('../')
+from utils.calc_dssp import annotate_sse
+
+ic.configureOutput(includeContext=True)
+
+def mask_inputs(seq, 
+        msa_masked, 
+        msa_full, 
+        xyz_t, 
+        t1d, 
+        mask_msa, 
+        input_seq_mask=None, 
+        input_str_mask=None, 
+        input_floating_mask=None, 
+        input_t1dconf_mask=None, 
+        loss_seq_mask=None, 
+        loss_str_mask=None, 
+        loss_str_mask_2d=None, 
+        dssp=False,
+        hotspots=False,
+        diffuser=None, 
+        t=None, 
+        freeze_seq_emb=False, 
+        mutate_seq=False, 
+        no_clamp_seq=False, 
+        norm_input=False,
+        contacts=None,
+        frac_provide_dssp=0.5,
+        dssp_mask_percentage=[0,100],
+        frac_provide_contacts=0.5,
+        struc_cond=False):
+    """
+    Parameters:
+        seq (torch.tensor, required): (I,L) integer sequence 
+
+        msa_masked (torch.tensor, required): (I,N_short,L,48)
+
+        msa_full  (torch,.tensor, required): (I,N_long,L,25)
+        
+        xyz_t (torch,tensor): (T,L,27,3) template crds BEFORE they go into get_init_xyz 
+        
+        t1d (torch.tensor, required): (I,L,22) this is the t1d before tacking on the chi angles 
+        
+        str_mask_1D (torch.tensor, required): Shape (L) rank 1 tensor where structure is masked at False positions 
+
+        seq_mask_1D (torch.tensor, required): Shape (L) rank 1 tensor where seq is masked at False positions
+        t1d_24: is there an extra dimension to input structure confidence?
+
+        diffuser: diffuser class
+        
+        t: time step
+
+    NOTE: in the MSA, the order is 20aa, 1x unknown, 1x mask token. We set the masked region to 22 (masked).
+        For the t1d, this has 20aa, 1x unkown, and 1x template conf. Here, we set the masked region to 21 (unknown).
+        This, we think, makes sense, as the template in normal RF training does not perfectly correspond to the MSA.
+    """
+
+    
+    
+    #ic(input_seq_mask.shape)
+    #ic(seq.shape)
+    #ic(msa_masked.shape)
+    #ic(msa_full.shape)
+    #ic(t1d.shape)
+    #ic(xyz_t.shape)
+    #ic(input_str_mask.shape)
+    #ic(mask_msa.shape)
+
+    ###########
+    seq_mask = input_seq_mask
+
+
+    ######################
+    ###sequence diffusion###
+    ######################
+    
+    str_mask     = input_str_mask
+    
+    x_0          = torch.nn.functional.one_hot(seq[0,...],num_classes=22).float()*2-1
+    seq_diffused = diffuser.q_sample(x_0,t,mask=seq_mask)
+    
+    seq_tmp=torch.argmax(seq_diffused,axis=-1).to(device=seq.device)
+    seq=seq_tmp.repeat(seq.shape[0], 1)
+
+    ###################
+    ###msa diffusion###
+    ###################
+
+    ### msa_masked ###
+    #ic(msa_masked.shape)
+    B,N,L,_=msa_masked.shape
+    msa_masked[:,0,:,:22] = seq_diffused
+
+    x_0_msa = msa_masked[0,1:,:,:22].float()*2-1
+    msa_seq_mask = seq_mask.unsqueeze(0).repeat(N-1, 1)
+    msa_diffused = diffuser.q_sample(x_0_msa,torch.tensor([t]),mask=msa_seq_mask)
+    
+    msa_masked[:,1:,:,:22] = torch.clone(msa_diffused)
+
+    # index 44/45 is insertion/deletion
+    # index 43 is the masked token NOTE check this
+    # index 42 is the unknown token 
+    msa_masked[:,0,:,22:44] = seq_diffused
+    msa_masked[:,1:,:,22:44] = msa_diffused
+
+    # insertion/deletion stuff 
+    msa_masked[:,0,~seq_mask,44:46] = 0
+
+    ### msa_full ### 
+    ################
+    #make msa_full same size as msa_masked
+    #ic(msa_full.shape)
+    msa_full = msa_full[:,:msa_masked.shape[1],:,:]
+    msa_full[:,0,:,:22] = seq_diffused
+    msa_full[:,1:,:,:22] = msa_diffused
+
+    ### t1d ###
+    ########### 
+    # NOTE: adjusting t1d last dim (confidence) from sequence mask
+    t1d = torch.cat((t1d, torch.zeros((t1d.shape[0],t1d.shape[1],1)).float()), -1).to(seq.device)
+    t1d[:,:,:21] = seq_diffused[...,:21]
+
+    #t1d[:,:,21] *= input_t1dconf_mask
+    #set diffused conf to 0 and everything else to 1
+    t1d[:,~seq_mask,21] = 0.0
+    t1d[:,seq_mask,21] = 1.0
+
+    t1d[:1,:,22] = 1-t/diffuser.num_timesteps
+    
+    #to do add structure confidence metric; need to expand dimensions of chkpt b4
+    #if t1d_24: JG - changed to be default
+    t1d = torch.cat((t1d, torch.zeros((t1d.shape[0],t1d.shape[1],1)).float()), -1).to(seq.device)
+    t1d[:,~str_mask,23] = 0.0
+    t1d[:,str_mask,23] = 1.0
+
+    if dssp:
+        print(f'adding dssp {frac_provide_dssp} of time')
+        t1d = torch.cat((t1d, torch.zeros((t1d.shape[0],t1d.shape[1],4)).float()), -1).to(seq.device)
+        #dssp info
+        #mask some percentage of dssp info in range dssp_mask_percentage[0],dssp_mask_percentage[1]
+        percentage_mask=random.randint(dssp_mask_percentage[0], dssp_mask_percentage[1])
+        dssp=annotate_sse(np.array(xyz_t[0,:,1,:].squeeze()), percentage_mask=percentage_mask)
+        #dssp_unmasked = annotate_sse(np.array(xyz_t[0,:,1,:].squeeze()), percentage_mask=0)
+        if np.random.rand()>frac_provide_dssp:
+            print('masking dssp')
+            dssp[...]=0 #replace with mask token
+            dssp[:,-1]=1
+        t1d[...,24:]=dssp
+    
+    if hotspots:
+        print(f"adding hotspots {frac_provide_contacts} of time")
+        t1d = torch.cat((t1d, torch.zeros((t1d.shape[0],t1d.shape[1],1)).float()), -1).to(seq.device)
+        #mask all contacts some fraction of the time
+        if np.random.rand()>frac_provide_contacts:
+            print('masking contacts')
+            contacts = torch.zeros(L) 
+        t1d[...,-1] = contacts
+
+    ### xyz_t ###
+    #############
+    xyz_t = get_init_xyz(xyz_t[None])
+    xyz_t = xyz_t[0]
+    #Sequence masking
+    xyz_t[:,:,3:,:] = float('nan')
+    # Structure masking
+    if struc_cond:
+        print("non-autoregressive structure conditioning")
+        r = diffuser.alphas_cumprod[t]
+        xyz_mask = (torch.rand(xyz_t.shape[1]) > r).to(torch.bool).to(seq.device)
+        xyz_mask = torch.logical_and(xyz_mask,~str_mask)
+        xyz_t[:,xyz_mask,:,:] = float('nan')
+    else:
+        xyz_t[:,~str_mask,:,:] = float('nan')
+    
+    ### mask_msa ###
+    ################
+    # NOTE: this is for loss scoring
+    mask_msa[:,:,~loss_seq_mask] = False
+    
+    out=dict(
+            seq= seq,
+            msa_masked= msa_masked,
+            msa_full= msa_full,
+            xyz_t= xyz_t,
+            t1d= t1d,
+            mask_msa= mask_msa,
+            seq_diffused= seq_diffused
+            )
+    
+    return out