Add app.py

app.py

@@ -0,0 +1,362 @@
+import sys
+sys.path.append('./diffab-repo')
+import os
+import shutil
+import pandas as pd
+import yaml
+import subprocess
+import streamlit as st
+import stmol
+import py3Dmol
+import tempfile
+import re
+import abnumber
+import gzip
+import tarfile
+import torch
+from tqdm.auto import tqdm
+from Bio import PDB
+from collections import OrderedDict
+
+from diffab.tools.renumber import renumber as renumber_antibody
+from diffab.tools.renumber.run import (
+    biopython_chain_to_sequence, 
+    assign_number_to_sequence,
+)
+
+DIFFAB_DIR = os.path.realpath('./diffab-repo')
+
+CDR_OPTIONS = OrderedDict()
+CDR_OPTIONS['H_CDR1'] = 'H1'
+CDR_OPTIONS['H_CDR2'] = 'H2'
+CDR_OPTIONS['H_CDR3'] = 'H3'
+CDR_OPTIONS['L_CDR1'] = 'L1'
+CDR_OPTIONS['L_CDR2'] = 'L2'
+CDR_OPTIONS['L_CDR3'] = 'L3'
+
+DESIGN_MODES = OrderedDict()
+DESIGN_MODES['denovo'] = 'De novo design'
+DESIGN_MODES['denovo_dock'] = 'De novo design (with HDOCK)'
+DESIGN_MODES['opt'] = 'Optimization'
+DESIGN_MODES['fixbb'] = 'Fix-backbone'
+
+MODE_CONFIG = {
+    'denovo': './configs/test/codesign_multicdrs.yml',
+    'denovo_dock': './configs/test/codesign_multicdrs.yml',
+    'opt': './configs/test/abopt_singlecdr.yml',
+    'fixbb': './configs/test/fixbb.yml',
+}
+
+GPU_AVAILABLE = torch.cuda.is_available()
+DEFAULT_NUM_SAMPLES = 5 if GPU_AVAILABLE else 1
+DEFAULT_NUM_DOCKS = 3
+
+
+def dict_to_func(d):
+    def f(x):
+        return d[x]
+    return f
+
+
+def get_config(save_dir, mode, cdrs, num_samples=5, optimization_step=4):
+    tmpl_path = MODE_CONFIG[mode]
+    with open(tmpl_path, 'r') as f:
+        cfg = yaml.safe_load(f)
+    cfg['sampling']['cdrs'] = cdrs
+    cfg['sampling']['num_samples'] = num_samples
+    cfg['sampling']['optimize_steps'] = [optimization_step, ]
+
+    save_path = os.path.join(save_dir, 'design.yml')
+    with open(save_path, 'w') as f:
+        yaml.dump(cfg, f)
+    return cfg, save_path
+
+
+def run_design(pdb_path, config_path, output_dir, docking, display_widget, num_docks=DEFAULT_NUM_DOCKS):
+    if docking:
+        cmd = f"python design_dock.py --antigen {pdb_path} --config {config_path} --num_docks {num_docks} "
+    else:
+        cmd = f"python design_pdb.py {pdb_path} --config {config_path} "
+    cmd += f"--batch_size 1 --out_root {output_dir} "
+
+    if GPU_AVAILABLE:
+        cmd += "--device cuda"
+    else:
+        cmd += "--device cpu"
+
+    result_dir = os.path.join(output_dir, 'design')
+    if os.path.exists(result_dir):
+        shutil.rmtree(result_dir)
+    
+    output_buffer = ''
+    proc = subprocess.Popen(
+        cmd,
+        shell=True,
+        env=os.environ.copy(),
+        bufsize=1,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        cwd=DIFFAB_DIR,
+    )
+    for line in iter(proc.stdout.readline, b''):
+        output_buffer += line.decode()
+
+        display_widget.code( 
+            '\n'.join(output_buffer.splitlines()[-10:]),
+        )
+    proc.stdout.close()
+    proc.wait()
+
+
+@st.cache
+def renumber_antibody_cached(in_pdb, out_pdb, file_id):
+    return renumber_antibody(
+        in_pdb, out_pdb, return_other_chains=True
+    )
+
+
+def gather_results(result_dir):
+    outputs = []
+    for root, dirs, files in os.walk(result_dir):
+        for fname in files:
+            if not re.match('^\d\d\d\d\.pdb$', fname):
+                continue
+            fpath = os.path.join(root, fname)
+            gname = os.path.basename(root)
+
+            outputs.append((gname, fname, fpath))
+    
+    parser = PDB.PDBParser(QUIET=True)
+    records = []
+    fpath_to_name = {}
+    for gname, fname, fpath in tqdm(outputs):
+        name = f"{gname}_{fname}"
+        structure = parser.get_structure(name, fpath)
+        model = structure[0]
+        record = {
+            'name': name,
+            'H1': None, 'H2': None, 'H3': None,
+            'L1': None, 'L2': None, 'L3': None,
+            'gname': gname, 'fname': fname, 'fpath': fpath,
+        }
+        for chain in model:
+            try:
+                seq, reslist = biopython_chain_to_sequence(chain)
+                numbers, abchain = assign_number_to_sequence(seq)
+                if abchain.chain_type == 'H':
+                    record['H1'] = abchain.cdr1_seq
+                    record['H2'] = abchain.cdr2_seq
+                    record['H3'] = abchain.cdr3_seq
+                elif abchain.chain_type in ('L', 'K'):
+                    record['L1'] = abchain.cdr1_seq
+                    record['L2'] = abchain.cdr2_seq
+                    record['L3'] = abchain.cdr3_seq
+            except abnumber.ChainParseError as e:
+                pass
+        records.append(record)
+        fpath_to_name[fpath] = name
+
+    with tarfile.open(os.path.join(result_dir, 'generated.tar.gz'), 'w:gz') as tar:
+        for record in records:
+            info = tar.gettarinfo(record['fpath'])
+            info.name = record['name']
+            tar.addfile(
+                tarinfo = info,
+                fileobj = open(record['fpath'], 'rb'),
+            )
+    
+    records = pd.DataFrame(records)
+
+    return records, fpath_to_name
+
+
+def main():
+    # Temporary workspace directory
+    if 'tempdir_path' not in st.session_state:
+        tempdir_path = tempfile.mkdtemp(prefix='streamlit')
+        st.session_state.tempdir_path = tempdir_path
+    else:
+        tempdir_path = st.session_state.tempdir_path
+    # Page layout
+    st.set_page_config(layout="wide")
+    st.markdown(
+        "# DiffAb \n\n"
+        "Antigen-Specific Antibody Design and Optimization with Diffusion-Based Generative Models for Protein Structures (NeurIPS 2022) \n\n"
+        "[[Paper](https://www.biorxiv.org/content/10.1101/2022.07.10.499510.abstract)] "
+        "[[Code](https://github.com/luost26/diffab)]"
+    )
+    left_col, right_col = st.columns(2)
+
+    # Step 1: Upload PDB or choose an example
+    uploaded_file = None
+    with left_col:
+        uploaded_file = st.file_uploader(
+            'Antigen structure or antibody-antigen complex',
+            # disabled=True
+        )
+
+        if uploaded_file is None:
+            with st.expander('Download examples', expanded=True):
+                with open('./data/examples/7DK2_AB_C.pdb', 'r') as f:
+                    st.download_button(
+                        'RBD + Antibody Complex', 
+                        data = f,
+                        file_name='RBD_AbAg.pdb',
+                    )
+                with open('./data/examples/Omicron_RBD.pdb', 'r') as f:
+                    st.download_button(
+                        'RBD Antigen Only', 
+                        data = f,
+                        file_name = 'RBD_AgOnly.pdb',
+                    )
+                st.text('Please upload the downloaded PDB file to run the demo.')
+
+    if 'submit' not in st.session_state:
+        st.session_state.submit = False
+    if 'done' not in st.session_state:
+        st.session_state.done = False
+
+    # Step 1.2: Retrieve uploaded PDB
+    if uploaded_file is not None:
+        pdb_path = os.path.join(tempdir_path, 'structure.pdb')
+        renum_path = os.path.join(tempdir_path, 'structure_renumber.pdb')
+        with open(pdb_path, 'w') as f:
+            f.write(uploaded_file.getvalue().decode())
+        H_chains, L_chains, Ag_chains = renumber_antibody_cached(
+            in_pdb = pdb_path,
+            out_pdb = renum_path,
+            file_id = uploaded_file.id
+        )
+        H_chain = H_chains[0] if H_chains else None
+        L_chain = L_chains[0] if L_chains else None
+        docking = H_chain is None and L_chain is None
+
+    # Step 2: Design options
+    if uploaded_file is not None:
+        with left_col:
+            st.dataframe(pd.DataFrame({
+                'Heavy': {'Chain': H_chain},
+                'Light': {'Chain': L_chain},
+                'Antigen': {'Chain': ','.join(Ag_chains)},
+            }), use_container_width=True)
+
+            form = st.form('design_form')
+            with form:
+                if H_chain is None and L_chain is None:
+                    # Antigen only
+                    cdr_options = ['H_CDR1', 'H_CDR2', 'H_CDR3', 'L_CDR1', 'L_CDR2', 'L_CDR3']
+                    cdr_default = ['H_CDR1', 'H_CDR2', 'H_CDR3']
+                    mode_options = ['denovo_dock']
+                elif H_chain is not None and L_chain is None:
+                    # Heavy chain + Antigen
+                    cdr_options = ['H_CDR1', 'H_CDR2', 'H_CDR3']
+                    cdr_default = ['H_CDR1', 'H_CDR2', 'H_CDR3']
+                    mode_options = ['denovo', 'opt', 'fixbb']
+                elif H_chain is None and L_chain is not None:
+                    # Light chain + Antigen
+                    cdr_options = ['L_CDR1', 'L_CDR2', 'L_CDR3']
+                    cdr_default = ['L_CDR1', 'L_CDR2', 'L_CDR3']
+                    mode_options = ['denovo', 'opt', 'fixbb']
+                else:
+                    # H + L + Ag
+                    cdr_options = ['H_CDR1', 'H_CDR2', 'H_CDR3', 'L_CDR1', 'L_CDR2', 'L_CDR3']
+                    cdr_default = ['H_CDR1', 'H_CDR2', 'H_CDR3']
+                    mode_options = ['denovo', 'opt', 'fixbb']
+                
+                design_mode = st.radio(
+                    'Mode',
+                    mode_options,
+                    format_func=dict_to_func(DESIGN_MODES),
+                    # disabled=True,
+                )
+                cdr_choices = st.multiselect(
+                    'CDRs',
+                    cdr_options,
+                    default = cdr_default,
+                    format_func=dict_to_func(CDR_OPTIONS),
+                    # disabled=True,
+                )
+
+                if docking:
+                    num_docks = st.slider(
+                        'Number of docking poses', 
+                        min_value=1, max_value=10, value=DEFAULT_NUM_DOCKS,
+                    )
+                else:
+                    num_docks = 0
+                num_designs = st.slider(
+                    'Number of samples',
+                    min_value=1, max_value=10, value=DEFAULT_NUM_SAMPLES,
+                )
+
+                submit = st.form_submit_button('Run')
+                st.session_state.submit = st.session_state.submit or submit
+                if submit:
+                    st.session_state.done = False
+
+    # Step 3: Prepare configuration and run design
+    if uploaded_file is not None and st.session_state.submit:
+        config, config_path = get_config(
+            save_dir = tempdir_path,
+            mode = design_mode,
+            cdrs = cdr_choices,
+            num_samples = num_designs,
+        )
+
+        with right_col:
+            result_molecule_display = st.empty()
+            result_select_widget = st.empty()
+            result_table_display = st.empty()
+            result_download_btn = st.empty()
+            output_display = st.empty()
+        if not st.session_state.done:
+            run_design(
+                pdb_path = renum_path,
+                config_path = config_path,
+                output_dir = tempdir_path,
+                docking = docking,
+                display_widget = output_display,
+                num_docks = num_docks,
+            )
+            st.session_state.done = True
+    
+            result_dir = os.path.join(tempdir_path, 'design')
+            df_cols = ['name'] + list(CDR_OPTIONS.values())
+            df_results, fpath_to_name = gather_results(result_dir)
+            st.session_state.results = (df_results, fpath_to_name)
+
+    # Step 5: Show results:
+    if st.session_state.submit and st.session_state.done:
+        result_dir = os.path.join(tempdir_path, 'design')
+        df_results, fpath_to_name = st.session_state.results
+
+        df_cols = ['name'] + list(CDR_OPTIONS.values())
+        result_table_display.dataframe(df_results[df_cols], use_container_width=True)
+
+        display_pdb_path = result_select_widget.selectbox(
+            label = "Visualize",
+            options = df_results['fpath'],
+            format_func = dict_to_func(fpath_to_name),
+        )
+
+        with open(os.path.join(result_dir, 'generated.tar.gz'), 'rb') as f:
+            result_download_btn.download_button(
+                label = "Download PDBs",
+                data = f,
+                file_name = "generated.tar.gz",
+            )
+
+        if not os.path.exists(display_pdb_path):
+            display_pdb_path = df_results['fpath'][0]
+        with open(display_pdb_path, 'r') as f:
+            pdb_str = f.read()
+        xyzview = py3Dmol.view(width=380, height=380)
+        xyzview.addModelsAsFrames(pdb_str)
+        xyzview.setStyle({'cartoon':{'color':'spectrum'}})
+        xyzview.zoomTo()
+        with result_molecule_display:
+            stmol.showmol(xyzview, width=380, height=380)
+    
+
+if __name__ == '__main__':
+    main()