| import os,sys
|
|
|
| from colabdesign.mpnn import mk_mpnn_model
|
| from colabdesign.af import mk_af_model
|
| from colabdesign.shared.protein import pdb_to_string
|
| from colabdesign.shared.parse_args import parse_args
|
|
|
| import pandas as pd
|
| import numpy as np
|
| from string import ascii_uppercase, ascii_lowercase
|
| alphabet_list = list(ascii_uppercase+ascii_lowercase)
|
|
|
| def get_info(contig):
|
| F = []
|
| free_chain = False
|
| fixed_chain = False
|
| sub_contigs = [x.split("-") for x in contig.split("/")]
|
| for n,(a,b) in enumerate(sub_contigs):
|
| if a[0].isalpha():
|
| L = int(b)-int(a[1:]) + 1
|
| F += [1] * L
|
| fixed_chain = True
|
| else:
|
| L = int(b)
|
| F += [0] * L
|
| free_chain = True
|
| return F,[fixed_chain,free_chain]
|
|
|
| def main(argv):
|
| ag = parse_args()
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| ag.txt("Designability Test")
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| ag.txt("REQUIRED")
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| ag.add(["pdb=" ], None, str, ["input pdb"])
|
| ag.add(["loc=" ], None, str, ["location to save results"])
|
| ag.add(["contigs=" ], None, str, ["contig definition"])
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| ag.txt("OPTIONAL")
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| ag.add(["copies=" ], 1, int, ["number of repeating copies"])
|
| ag.add(["num_seqs=" ], 8, int, ["number of mpnn designs to evaluate"])
|
| ag.add(["initial_guess" ], False, None, ["initialize previous coordinates"])
|
| ag.add(["use_multimer" ], False, None, ["use alphafold_multimer_v3"])
|
| ag.add(["use_soluble" ], False, None, ["use solubleMPNN"])
|
| ag.add(["num_recycles=" ], 3, int, ["number of recycles"])
|
| ag.add(["rm_aa="], "C", str, ["disable specific amino acids from being sampled"])
|
| ag.add(["num_designs=" ], 1, int, ["number of designs to evaluate"])
|
| ag.add(["mpnn_sampling_temp=" ], 0.1, float, ["sampling temperature used by proteinMPNN"])
|
| ag.txt("-------------------------------------------------------------------------------------")
|
| o = ag.parse(argv)
|
|
|
| if None in [o.pdb, o.loc, o.contigs]:
|
| ag.usage("Missing Required Arguments")
|
|
|
| if o.rm_aa == "":
|
| o.rm_aa = None
|
|
|
|
|
| contigs = []
|
| for contig_str in o.contigs.replace(" ",":").replace(",",":").split(":"):
|
| if len(contig_str) > 0:
|
| contig = []
|
| for x in contig_str.split("/"):
|
| if x != "0": contig.append(x)
|
| contigs.append("/".join(contig))
|
|
|
| chains = alphabet_list[:len(contigs)]
|
| info = [get_info(x) for x in contigs]
|
| fixed_pos = []
|
| fixed_chains = []
|
| free_chains = []
|
| both_chains = []
|
| for pos,(fixed_chain,free_chain) in info:
|
| fixed_pos += pos
|
| fixed_chains += [fixed_chain and not free_chain]
|
| free_chains += [free_chain and not fixed_chain]
|
| both_chains += [fixed_chain and free_chain]
|
|
|
| flags = {"initial_guess":o.initial_guess,
|
| "best_metric":"rmsd",
|
| "use_multimer":o.use_multimer,
|
| "model_names":["model_1_multimer_v3" if o.use_multimer else "model_1_ptm"]}
|
|
|
| if sum(both_chains) == 0 and sum(fixed_chains) > 0 and sum(free_chains) > 0:
|
| protocol = "binder"
|
| print("protocol=binder")
|
| target_chains = []
|
| binder_chains = []
|
| for n,x in enumerate(fixed_chains):
|
| if x: target_chains.append(chains[n])
|
| else: binder_chains.append(chains[n])
|
| af_model = mk_af_model(protocol="binder",**flags)
|
| prep_flags = {"target_chain":",".join(target_chains),
|
| "binder_chain":",".join(binder_chains),
|
| "rm_aa":o.rm_aa}
|
| opt_extra = {}
|
|
|
| elif sum(fixed_pos) > 0:
|
| protocol = "partial"
|
| print("protocol=partial")
|
| af_model = mk_af_model(protocol="fixbb",
|
| use_templates=True,
|
| **flags)
|
| rm_template = np.array(fixed_pos) == 0
|
| prep_flags = {"chain":",".join(chains),
|
| "rm_template":rm_template,
|
| "rm_template_seq":rm_template,
|
| "copies":o.copies,
|
| "homooligomer":o.copies>1,
|
| "rm_aa":o.rm_aa}
|
| else:
|
| protocol = "fixbb"
|
| print("protocol=fixbb")
|
| af_model = mk_af_model(protocol="fixbb",**flags)
|
| prep_flags = {"chain":",".join(chains),
|
| "copies":o.copies,
|
| "homooligomer":o.copies>1,
|
| "rm_aa":o.rm_aa}
|
|
|
| batch_size = 8
|
| if o.num_seqs < batch_size:
|
| batch_size = o.num_seqs
|
|
|
| print("running proteinMPNN...")
|
| sampling_temp = o.mpnn_sampling_temp
|
| mpnn_model = mk_mpnn_model(weights="soluble" if o.use_soluble else "original")
|
| outs = []
|
| pdbs = []
|
| for m in range(o.num_designs):
|
| if o.num_designs == 0:
|
| pdb_filename = o.pdb
|
| else:
|
| pdb_filename = o.pdb.replace("_0.pdb",f"_{m}.pdb")
|
| pdbs.append(pdb_filename)
|
| af_model.prep_inputs(pdb_filename, **prep_flags)
|
| if protocol == "partial":
|
| p = np.where(fixed_pos)[0]
|
| af_model.opt["fix_pos"] = p[p < af_model._len]
|
|
|
| mpnn_model.get_af_inputs(af_model)
|
| outs.append(mpnn_model.sample(num=o.num_seqs//batch_size, batch=batch_size, temperature=sampling_temp))
|
|
|
| if protocol == "binder":
|
| af_terms = ["plddt","i_ptm","i_pae","rmsd"]
|
| elif o.copies > 1:
|
| af_terms = ["plddt","ptm","i_ptm","pae","i_pae","rmsd"]
|
| else:
|
| af_terms = ["plddt","ptm","pae","rmsd"]
|
|
|
| labels = ["design","n","score"] + af_terms + ["seq"]
|
| data = []
|
| best = {"rmsd":np.inf,"design":0,"n":0}
|
| print("running AlphaFold...")
|
| os.system(f"mkdir -p {o.loc}/all_pdb")
|
| with open(f"{o.loc}/design.fasta","w") as fasta:
|
| for m,(out,pdb_filename) in enumerate(zip(outs,pdbs)):
|
| out["design"] = []
|
| out["n"] = []
|
| af_model.prep_inputs(pdb_filename, **prep_flags)
|
| for k in af_terms: out[k] = []
|
| for n in range(o.num_seqs):
|
| out["design"].append(m)
|
| out["n"].append(n)
|
| sub_seq = out["seq"][n].replace("/","")[-af_model._len:]
|
| af_model.predict(seq=sub_seq, num_recycles=o.num_recycles, verbose=False)
|
| for t in af_terms: out[t].append(af_model.aux["log"][t])
|
| if "i_pae" in out:
|
| out["i_pae"][-1] = out["i_pae"][-1] * 31
|
| if "pae" in out:
|
| out["pae"][-1] = out["pae"][-1] * 31
|
| rmsd = out["rmsd"][-1]
|
| if rmsd < best["rmsd"]:
|
| best = {"design":m,"n":n,"rmsd":rmsd}
|
| af_model.save_current_pdb(f"{o.loc}/all_pdb/design{m}_n{n}.pdb")
|
| af_model._save_results(save_best=True, verbose=False)
|
| af_model._k += 1
|
| score_line = [f'design:{m} n:{n}',f'mpnn:{out["score"][n]:.3f}']
|
| for t in af_terms:
|
| score_line.append(f'{t}:{out[t][n]:.3f}')
|
| print(" ".join(score_line)+" "+out["seq"][n])
|
| line = f'>{"|".join(score_line)}\n{out["seq"][n]}'
|
| fasta.write(line+"\n")
|
| data += [[out[k][n] for k in labels] for n in range(o.num_seqs)]
|
| af_model.save_pdb(f"{o.loc}/best_design{m}.pdb")
|
|
|
|
|
| with open(f"{o.loc}/best.pdb", "w") as handle:
|
| remark_text = f"design {best['design']} N {best['n']} RMSD {best['rmsd']:.3f}"
|
| handle.write(f"REMARK 001 {remark_text}\n")
|
| handle.write(open(f"{o.loc}/best_design{best['design']}.pdb", "r").read())
|
|
|
| labels[2] = "mpnn"
|
| df = pd.DataFrame(data, columns=labels)
|
| df.to_csv(f'{o.loc}/mpnn_results.csv')
|
|
|
| if __name__ == "__main__":
|
| main(sys.argv[1:])
|
|
|
