cpu requirements

app.py

@@ -69,10 +69,10 @@ def clear_mem():
         buf.delete()
 
 
-print("Is cuda available",torch.cuda.is_available())
-stream = os.popen('nvcc --version')
-output = stream.read()
-print(output)
+print("Is cuda available", torch.cuda.is_available())
+# stream = os.popen("nvcc --version")
+# output = stream.read()
+# print(output)
 
 
 def setup_af(seq, model_name="model_5_ptm"):
@@ -177,7 +177,7 @@ def align_structures(pdb1, pdb2, lenRes, index):
     io.set_structure(sample_structure)
     io.save(f"outputs/out_{index}_aligned.pdb")
     # Doing this to get around biopython CEALIGN bug
-    #subprocess.call("pymol -c -Q -r cealign.pml", shell=True)
+    # subprocess.call("pymol -c -Q -r cealign.pml", shell=True)
 
     return aligner.rms, "outputs/reference.pdb", f"outputs/out_{index}_aligned.pdb"
 
@@ -219,7 +219,7 @@ def run_alphafold(sequences, num_recycles):
         }
 
         positions = []
- 
+
         for r in range(recycles + 1):
             outs = RUNNER(x, OPT)
             outs = jax.tree_map(lambda x: np.asarray(x), outs)
@@ -228,10 +228,12 @@ def run_alphafold(sequences, num_recycles):
         plddts.append(outs["plddt"][:LEN])
         paes.append(outs["pae"])
         if os.path.exists("/home/duerr/phd/08_Code/ProteinMPNN"):
-            save_pdb(outs, f"/home/duerr/phd/08_Code/ProteinMPNN/outputs/out_{i}.pdb", LEN)
+            save_pdb(
+                outs, f"/home/duerr/phd/08_Code/ProteinMPNN/outputs/out_{i}.pdb", LEN
+            )
         else:
             save_pdb(outs, f"/home/user/app/outputs/out_{i}.pdb", LEN)
-    return plddts,paes, LEN
+    return plddts, paes, LEN
 
 
 if os.path.exists("/home/duerr/phd/08_Code/ProteinMPNN"):
@@ -308,7 +310,7 @@ def preprocess_mol(pdb_code="", filepath=""):
             return None
     else:
         mol = Molecule(pdb_code)
-    mol.write('original.pdb')
+    mol.write("original.pdb")
     # clean messy files and only include protein itself
     mol.filter("protein")
     # renumber using moleculekit 0...len(protein)
@@ -324,70 +326,73 @@ def preprocess_mol(pdb_code="", filepath=""):
     mol.write("cleaned.pdb")
     return "cleaned.pdb", df
 
+
 def assign_sasa(mol):
     from moleculekit.projections.metricsasa import MetricSasa
-    metr = MetricSasa(
-                mode="residue", filtersel="protein"
-            ) 
+
+    metr = MetricSasa(mode="residue", filtersel="protein")
     sasaR = metr.project(mol)[0]
     is_prot = mol.atomselect("protein")
-    resids=pd.DataFrame.from_dict({"resid":mol.resid, "is_prot":is_prot})
-    new_masses=[]
+    resids = pd.DataFrame.from_dict({"resid": mol.resid, "is_prot": is_prot})
+    new_masses = []
     i_without_non_prot = 0
-    for i, g in resids.groupby((resids['resid'].shift() != resids['resid']).cumsum()):
-        if g["is_prot"].unique()[0]==True:
-            g["sasa"]=sasaR[i_without_non_prot]
-            i_without_non_prot+=1
+    for i, g in resids.groupby((resids["resid"].shift() != resids["resid"]).cumsum()):
+        if g["is_prot"].unique()[0] == True:
+            g["sasa"] = sasaR[i_without_non_prot]
+            i_without_non_prot += 1
         else:
-            g["sasa"]=0
+            g["sasa"] = 0
         new_masses.extend(list(g.sasa))
     return np.array(new_masses)
 
+
 def process_atomsel(atomsel):
     """everything lowercase and replace some keywords not relevant for protein design"""
-    atomsel=re.sub('sasa', 'mass',atomsel, flags=re.I)
-    atomsel=re.sub('plddt', 'beta',atomsel, flags=re.I)
+    atomsel = re.sub("sasa", "mass", atomsel, flags=re.I)
+    atomsel = re.sub("plddt", "beta", atomsel, flags=re.I)
     return atomsel
 
 
 def make_fixed_positions_dict(atomsel, residue_index_df):
     # we use the uploaded file for the selection
-    mol = Molecule('original.pdb')
+    mol = Molecule("original.pdb")
     # use index for selection as resids will change
 
-  
     # set sasa to 0 for all non protein atoms (all non protein atoms are deleted later)
     mol.masses = assign_sasa(mol)
     print(mol.masses.shape)
     print(assign_sasa(mol).shape)
     atomsel = process_atomsel(atomsel)
-    selected_residues = mol.get("index",atomsel)
+    selected_residues = mol.get("index", atomsel)
 
     # clean up
     mol.filter("protein")
     mol.renumberResidues()
     # based on selected index now get resids
     selected_residues = [str(i) for i in selected_residues]
-    if len(selected_residues)==0:
+    if len(selected_residues) == 0:
         return None, []
     selected_residues_str = " ".join(selected_residues)
-    selected_residues=set(mol.get('resid', sel=f"index {selected_residues_str}"))
+    selected_residues = set(mol.get("resid", sel=f"index {selected_residues_str}"))
 
     # use the proteinMPNN index nomenclature to assemble fixed_positions_dict
-    fixed_positions_df = residue_index_df[residue_index_df['new_resid'].isin(selected_residues)]
+    fixed_positions_df = residue_index_df[
+        residue_index_df["new_resid"].isin(selected_residues)
+    ]
 
-    chains = set(mol.get('chain', sel="all"))
-    fixed_position_dict = {'cleaned':{}}
-    #store the selected residues in a list for the visualization later with cleaned.pdb
-    selected_residues = list(fixed_positions_df['new_resid'])
+    chains = set(mol.get("chain", sel="all"))
+    fixed_position_dict = {"cleaned": {}}
+    # store the selected residues in a list for the visualization later with cleaned.pdb
+    selected_residues = list(fixed_positions_df["new_resid"])
 
     for c in chains:
-        fixed_position_dict['cleaned'][c]=[]
+        fixed_position_dict["cleaned"][c] = []
 
     for i, row in fixed_positions_df.iterrows():
-        fixed_position_dict['cleaned'][row['chain']].append(row['proteinMPNN_index'])
+        fixed_position_dict["cleaned"][row["chain"]].append(row["proteinMPNN_index"])
     return fixed_position_dict, selected_residues
 
+
 def update(
     inp,
     file,
@@ -398,7 +403,7 @@ def update(
     sampling_temp,
     model_name,
     backbone_noise,
-    atomsel
+    atomsel,
 ):
     from protein_mpnn_utils import (
         loss_nll,
@@ -471,8 +476,10 @@ def update(
     if atomsel == "":
         fixed_positions_dict, selected_residues = None, []
     else:
-        fixed_positions_dict, selected_residues=make_fixed_positions_dict(atomsel, mol_index)
-    
+        fixed_positions_dict, selected_residues = make_fixed_positions_dict(
+            atomsel, mol_index
+        )
+
     pssm_dict = None
     omit_AA_dict = None
     bias_AA_dict = None
@@ -489,7 +496,7 @@ def update(
         tied_positions_dict = make_tied_positions_for_homomers(pdb_dict_list)
     else:
         tied_positions_dict = None
-    
+
     chain_id_dict = {}
     chain_id_dict[pdb_dict_list[0]["name"]] = (designed_chain_list, fixed_chain_list)
     with torch.no_grad():
@@ -738,7 +745,7 @@ def update(
                         )
                         seq_list.append(seq + chain_s)
                         message += f"{line}\n"
-    if fixed_positions_dict!=None:
+    if fixed_positions_dict != None:
         message += f"\nfixed positions:* {fixed_positions_dict['cleaned']} \n\n*uses CHAIN:[1..len(chain)] residue numbering"
     # somehow sequences still contain X, remove again
     for i, x in enumerate(seq_list):
@@ -769,7 +776,7 @@ def update(
     )
 
     fig_tadjusted.update_xaxes(side="top")
-    seq_dict = {"seq_list":seq_list, "recovery":seq_recovery, "seq_score":seq_score}
+    seq_dict = {"seq_list": seq_list, "recovery": seq_recovery, "seq_score": seq_score}
     return (
         message,
         fig,
@@ -779,17 +786,17 @@ def update(
         pdb_path,
         gr.Dropdown.update(choices=seq_list),
         selected_residues,
-        seq_dict
+        seq_dict,
     )
 
 
-def update_AF(seq_dict, pdb, num_recycles,selectedResidues):
+def update_AF(seq_dict, pdb, num_recycles, selectedResidues):
 
     # # run alphafold using ray
     # plddts, pae, num_res = run_alphafold(
     #    startsequence, num_recycles
     # )
-    allSeqs = seq_dict['seq_list']
+    allSeqs = seq_dict["seq_list"]
     lenSeqs = len(allSeqs)
     if len(allSeqs[0]) > 700:
         return (
@@ -801,14 +808,22 @@ def update_AF(seq_dict, pdb, num_recycles,selectedResidues):
             plt.figure(),
             plt.figure(),
         )
-    
+
     plddts, paes, num_res = ray.get(run_alphafold.remote(allSeqs, num_recycles))
 
     sequences = {}
     for i in range(lenSeqs):
-        rms, input_pdb, aligned_pdb = align_structures(pdb, f"outputs/out_{i}.pdb", num_res,i)
-        sequences[i]={"Seq":i,"RMSD":f"{rms:.2f}","Score":seq_dict['seq_score'][i],"Recovery":seq_dict["recovery"][i],"Mean pLDDT":f"{np.mean(plddts[i]):.4f}"}
-    results=pd.DataFrame.from_dict(sequences, orient="index")
+        rms, input_pdb, aligned_pdb = align_structures(
+            pdb, f"outputs/out_{i}.pdb", num_res, i
+        )
+        sequences[i] = {
+            "Seq": i,
+            "RMSD": f"{rms:.2f}",
+            "Score": seq_dict["seq_score"][i],
+            "Recovery": seq_dict["recovery"][i],
+            "Mean pLDDT": f"{np.mean(plddts[i]):.4f}",
+        }
+    results = pd.DataFrame.from_dict(sequences, orient="index")
     print(results)
     plots = []
     for index, plddts_val in enumerate(plddts):
@@ -837,7 +852,7 @@ def update_AF(seq_dict, pdb, num_recycles,selectedResidues):
         legend=dict(yanchor="bottom", y=0.01, xanchor="left", x=0.99),
     )
     pae_plots = []
-    for i,pae in enumerate(paes):
+    for i, pae in enumerate(paes):
         plt.figure()
         plt.title(f"Predicted Aligned Error sequence {i}")
         Ln = pae.shape[0]
@@ -858,7 +873,20 @@ def update_AF(seq_dict, pdb, num_recycles,selectedResidues):
     # plotAF_pae.write_html("test.html")
     # plotAF_pae.update_layout(title="Predicted Aligned Error", template="simple_white")
 
-    return molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs, sequences), plotAF_plddt, pae_plots, results
+    return (
+        molecule(
+            input_pdb,
+            aligned_pdb,
+            lenSeqs,
+            num_res,
+            selectedResidues,
+            allSeqs,
+            sequences,
+        ),
+        plotAF_plddt,
+        pae_plots,
+        results,
+    )
 
 
 def read_mol(molpath):
@@ -870,26 +898,39 @@ def read_mol(molpath):
     return mol
 
 
-def molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs, sequences):
+def molecule(
+    input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs, sequences
+):
 
     mol = read_mol("outputs/reference.pdb")
-    options =""
+    options = ""
     pred_mol = "["
     seqdata = "{"
     selected = "selected"
     for i in range(lenSeqs):
-        seqdata+=str(i)+': { "score": '+sequences[i]["Score"]+', "rmsd": '+sequences[i]["RMSD"]+', "recovery": '+sequences[i]["Recovery"]+', "plddt": '+sequences[i]["Mean pLDDT"]+', "seq":"'+allSeqs[i]+'"}'
-        options+=f'<option {selected} value="{i}">sequence {i} </option>' #RMSD {sequences[i]["RMSD"]}, score {sequences[i]["Score"]}, recovery {sequences[i]["Recovery"]} pLDDT {sequences[i]["Mean pLDDT"]}
-        p=f"outputs/out_{i}_aligned.pdb"
-        pred_mol+=f"`{read_mol(p)}`"
+        seqdata += (
+            str(i)
+            + ': { "score": '
+            + sequences[i]["Score"]
+            + ', "rmsd": '
+            + sequences[i]["RMSD"]
+            + ', "recovery": '
+            + sequences[i]["Recovery"]
+            + ', "plddt": '
+            + sequences[i]["Mean pLDDT"]
+            + ', "seq":"'
+            + allSeqs[i]
+            + '"}'
+        )
+        options += f'<option {selected} value="{i}">sequence {i} </option>'  # RMSD {sequences[i]["RMSD"]}, score {sequences[i]["Score"]}, recovery {sequences[i]["Recovery"]} pLDDT {sequences[i]["Mean pLDDT"]}
+        p = f"outputs/out_{i}_aligned.pdb"
+        pred_mol += f"`{read_mol(p)}`"
         selected = ""
-        if i!=lenSeqs-1:
-            pred_mol+=","
-            seqdata+=","
-    pred_mol+="]"
-    seqdata+="}"
-
-
+        if i != lenSeqs - 1:
+            pred_mol += ","
+            seqdata += ","
+    pred_mol += "]"
+    seqdata += "}"
 
     x = (
         """<!DOCTYPE html>
@@ -934,7 +975,9 @@ def molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs
             a sequence</label>
         <select id="seq"
             class="bg-gray-50 border border-gray-300 text-gray-900 text-sm rounded-lg focus:ring-blue-500 focus:border-blue-500 block w-full p-2.5 dark:bg-gray-700 dark:border-gray-600 dark:placeholder-gray-400 dark:text-white dark:focus:ring-blue-500 dark:focus:border-blue-500">
-            """+options+"""
+            """
+        + options
+        + """
         </select>
     </div>
     <div class="font-mono bg-gray-100 py-3 px-2  font-sm rounded">
@@ -1033,16 +1076,18 @@ def molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs
         let element = null;
         let config = null;
         let currentIndex = 0;
-            let seqs = """+seqdata+"""
+            let seqs = """
+        + seqdata
+        + """
                let data = """
         + pred_mol
         + """  
                 let pdb = `"""
         + mol
         + """`  
-         var selectedResidues = """+
-                f"{selectedResidues}"
-                +"""
+         var selectedResidues = """
+        + f"{selectedResidues}"
+        + """
                 //AlphaFold code from https://gist.github.com/piroyon/30d1c1099ad488a7952c3b21a5bebc96
                 let colorAlpha = function (atom) {
                     if (atom.b < 50) {
@@ -1057,7 +1102,9 @@ def molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs
                 };
                
                 let colors = {}
-                for (let i=0; i<"""+str(num_res)+""";i++){
+                for (let i=0; i<"""
+        + str(num_res)
+        + """;i++){
                 if (selectedResidues.includes(i)){
                     colors[i]="hotpink"
                 }else{
@@ -1152,7 +1199,16 @@ def molecule(input_pdb, aligned_pdb, lenSeqs, num_res, selectedResidues, allSeqs
 
 
 def set_examples(example):
-    label, inp, designed_chain, fixed_chain, homomer, num_seqs, sampling_temp, atomsel = example
+    (
+        label,
+        inp,
+        designed_chain,
+        fixed_chain,
+        homomer,
+        num_seqs,
+        sampling_temp,
+        atomsel,
+    ) = example
     return [
         label,
         inp,
@@ -1161,7 +1217,7 @@ def set_examples(example):
         homomer,
         gr.Slider.update(value=num_seqs),
         gr.Radio.update(value=sampling_temp),
-        atomsel
+        atomsel,
     ]
 
 
@@ -1198,7 +1254,8 @@ with proteinMPNN:
                     value=0.1,
                     label="Sampling temperature",
                 )
-            gr.Markdown(""" Sampling temperature for amino acids, `T=0.0` means taking argmax, `T>>1.0` means sample randomly. Suggested values `0.1, 0.15, 0.2, 0.25, 0.3`. Higher values will lead to more diversity.
+            gr.Markdown(
+                """ Sampling temperature for amino acids, `T=0.0` means taking argmax, `T>>1.0` means sample randomly. Suggested values `0.1, 0.15, 0.2, 0.25, 0.3`. Higher values will lead to more diversity.
             """
             )
             with gr.Row():
@@ -1220,8 +1277,9 @@ with proteinMPNN:
                 gr.Markdown(
                     "for correct symmetric tying lenghts of homomer chains should be the same"
                 )
-            gr.Markdown('## Fixed positions')
-            gr.Markdown("""You can fix important positions in the protein. Resid should be specified with the same numbering as in the input pdb file. The fixed residues will be highlighted in the output. 
+            gr.Markdown("## Fixed positions")
+            gr.Markdown(
+                """You can fix important positions in the protein. Resid should be specified with the same numbering as in the input pdb file. The fixed residues will be highlighted in the output. 
                 The [VMD selection](http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.2/ug/node89.html) synthax is used. You can also select based on ligands or chains in the input structure to specify interfaces to be fixed.
 
                  - <code>within 5 of resid 94</code> All residues that have >1 atom closer than 5 Å to any atom of residue 94
@@ -1235,8 +1293,11 @@ with proteinMPNN:
                  - <code>pLDDT >70 </code> Redesign all residues with low pLDDT
                  
                  Note that <code>sasa</code> and <code>pLDDT</code> selectors modify default VMD behavior. SASA is calculated using moleculekit and written to the mass attribute. Selections based on mass do not work. 
-                 pLDDT is an alias for beta, it only works correctly with structures that contain the appropriate values in the beta column of the PDB file. """)
-            atomsel = gr.Textbox(placeholder="Specify atom selection ", label="Fixed positions")
+                 pLDDT is an alias for beta, it only works correctly with structures that contain the appropriate values in the beta column of the PDB file. """
+            )
+            atomsel = gr.Textbox(
+                placeholder="Specify atom selection ", label="Fixed positions"
+            )
 
         btn = gr.Button("Run")
     label = gr.Textbox(label="Label", visible=False)
@@ -1249,19 +1310,45 @@ with proteinMPNN:
             homomer,
             num_seqs,
             sampling_temp,
-            atomsel
+            atomsel,
         ],
         samples=[
-            ["Homomer design", "1O91", "A,B,C", "", True, 2, 0.1,""],
-            ["Monomer design", "6MRR", "A", "", False, 2, 0.1,""],
+            ["Homomer design", "1O91", "A,B,C", "", True, 2, 0.1, ""],
+            ["Monomer design", "6MRR", "A", "", False, 2, 0.1, ""],
             ["Redesign of Homomer to Heteromer", "3HTN", "A,B", "C", False, 2, 0.1, ""],
-            ["Redesign of MID1 scaffold keeping binding site fixed", "3V1C", "A,B", "", False, 2, 0.1, "within 5 of resname ZN"],
-            ["Redesign of DNA binding protein", "3JRD", "A,B", "", False, 2, 0.1, "within 8 of nucleic"],
-            ["Surface Redesign of miniprotein", "7JZM", "A,B", "", False, 2, 0.1, "chain B or (chain A and sasa < 20)"],
+            [
+                "Redesign of MID1 scaffold keeping binding site fixed",
+                "3V1C",
+                "A,B",
+                "",
+                False,
+                2,
+                0.1,
+                "within 5 of resname ZN",
+            ],
+            [
+                "Redesign of DNA binding protein",
+                "3JRD",
+                "A,B",
+                "",
+                False,
+                2,
+                0.1,
+                "within 8 of nucleic",
+            ],
+            [
+                "Surface Redesign of miniprotein",
+                "7JZM",
+                "A,B",
+                "",
+                False,
+                2,
+                0.1,
+                "chain B or (chain A and sasa < 20)",
+            ],
         ],
     )
 
-
     gr.Markdown("# Output")
 
     with gr.Tabs():
@@ -1296,8 +1383,9 @@ with proteinMPNN:
             with gr.Row():
                 with gr.Row():
                     chosen_seq = gr.Dropdown(
-                        choices=[], label="Select a sequence for validation",
-                        visible=False
+                        choices=[],
+                        label="Select a sequence for validation",
+                        visible=False,
                     )
                     num_recycles = gr.Dropdown(
                         choices=[0, 1, 3, 5], value=3, label="num Recycles"
@@ -1307,10 +1395,25 @@ with proteinMPNN:
                 mol = gr.HTML()
                 with gr.Column():
                     gr.Markdown("## Metrics")
-                    results = gr.Dataframe(nteractive=False, row_count=(0, 'dynamic'), headers=["Seq","RMSD","Score","Recovery","Mean pLDDT"])
+                    p = {
+                        0: {
+                            "Seq": "NA",
+                            "RMSD": "NA",
+                            "Score": "NA",
+                            "Recovery": "NA",
+                            "Mean pLDDT": "NA",
+                        }
+                    }
+                    placeholder = pd.DataFrame.from_dict(p, orient="index")
+                    results = gr.Dataframe(
+                        placeholder,
+                        interactive=False,
+                        row_count=(1, "dynamic"),
+                        headers=["Seq", "RMSD", "Score", "Recovery", "Mean pLDDT"],
+                    )
                     plotAF_plddt = gr.Plot(label="pLDDT")
                     # remove maxh80 class from css
-                    plotAF_pae = gr.Gallery(label="PAE plots") #gr.Plot(label="PAE")
+                    plotAF_pae = gr.Gallery(label="PAE plots")  # gr.Plot(label="PAE")
     tempFile = gr.Variable()
     selectedResidues = gr.Variable()
     seq_dict = gr.Variable()
@@ -1337,7 +1440,7 @@ with proteinMPNN:
             tempFile,
             chosen_seq,
             selectedResidues,
-            seq_dict
+            seq_dict,
         ],
     )
     btnAF.click(

requirements.txt

@@ -6,8 +6,8 @@ dm-haiku==0.0.5
 dm-tree==0.1.6
 docker==5.0.0
 immutabledict==2.0.0
-https://storage.googleapis.com/jax-releases/jax/jax-0.3.7.tar.gz
-https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp38-none-manylinux2014_x86_64.whl
+jax==0.3.7
+jaxlib==0.3.7
 ml-collections==0.1.0
 numpy
 pandas==1.3.4
@@ -18,7 +18,7 @@ plotly
 GPUtil
 ray
 tqdm
-gradio==3.0.11
+gradio==3.1.4
 protobuf<4
 mdtraj
 -f https://storage.googleapis.com/jax-releases/jax_releases.html