This commit is contained in:
Simon Duerr 2024-04-09 23:01:16 +02:00 committed by GitHub
parent 78cae70931
commit f6c038563d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 322 additions and 3 deletions

View file

@ -80,7 +80,7 @@ tar xfz pdb100_2021Mar03.tar.gz
We use a library called Hydra to compose config files for predictions. The actual script that runs the model is in `rf2aa/run_inference.py` and default parameters that were used to train the model are in `rf2aa/config/inference/base.yaml`. We highly suggest using the default parameters since those are closest to the training task for RFAA but we have found that increasing loader_params.MAXCYCLE=10 (default set to 4) gives better results for hard cases (as noted in the paper).
The general way to run the model is as follows:
The general way to run the model using the command line is as follows:
```
python -m rf2aa.run_inference --config-name {your inference config}
```
@ -91,6 +91,10 @@ The main inputs into the model are split into:
- covalent bonds between protein chains and small molecule chains
- modified or unnatural amino acids (COMING SOON)
You can generate inputs for the model using this [webapp](https://hf.co/spaces/simonduerr/RFAA).
Predictions can also be run interactively by launching `python ui.py` and opening the displayed link in a web browser.
In the following sections, we will describe how to set up configs for different prediction tasks that we described in the paper.
<a id="protein-pred"></a>
@ -253,7 +257,7 @@ becomes this so it can be parsed correctly:
"[((\"A\", \"74\", \"ND2\"), (\"B\", \"1\"), (\"CW\", \"null\"))]"
```
We know this syntax is hard to work with and we are happy to review PRs if anyone in the community can figure out how to specify all the necessary requirements in a more user friendly way!
To generate this input you can use the [webapp](https://hf.co/spaces/simonduerr/RFAA) to generate this input automatically and delete any undesired atoms/leaving groups.
<a id="outputs"></a>
### Understanding model outputs

View file

@ -324,4 +324,9 @@ dependencies:
- shortuuid==1.0.12
- smmap==5.0.1
- subprocess32==3.5.4
- wandb==0.12.0
- wandb==0.12.0
- gradio
- gradio_molecule3d
- gradio_cofoldinginput
- biopython
- plotly

310
ui.py Normal file
View file

@ -0,0 +1,310 @@
"""
Input UI for RoseTTAfold All Atom
using two custom gradio components: gradio_molecule3d and gradio_cofoldinginput
"""
import json
import yaml
import os
import zipfile
import torch
import gradio as gr
import plotly.express as px
from openbabel import openbabel
from gradio_cofoldinginput import CofoldingInput
from gradio_molecule3d import Molecule3D
baseconfig = """job_name: "structure_prediction"
output_path: ""
checkpoint_path: RFAA_paper_weights.pt
database_params:
sequencedb: ""
hhdb: "pdb100_2021Mar03/pdb100_2021Mar03"
command: make_msa.sh
num_cpus: 4
mem: 64
protein_inputs: null
na_inputs: null
sm_inputs: null
covale_inputs: null
residue_replacement: null
chem_params:
use_phospate_frames_for_NA: True
use_cif_ordering_for_trp: True
loader_params:
n_templ: 4
MAXLAT: 128
MAXSEQ: 1024
MAXCYCLE: 4
BLACK_HOLE_INIT: False
seqid: 150.0
legacy_model_param:
n_extra_block: 4
n_main_block: 32
n_ref_block: 4
n_finetune_block: 0
d_msa: 256
d_msa_full: 64
d_pair: 192
d_templ: 64
n_head_msa: 8
n_head_pair: 6
n_head_templ: 4
d_hidden_templ: 64
p_drop: 0.0
use_chiral_l1: True
use_lj_l1: True
use_atom_frames: True
recycling_type: "all"
use_same_chain: True
lj_lin: 0.75
SE3_param:
num_layers: 1
num_channels: 32
num_degrees: 2
l0_in_features: 64
l0_out_features: 64
l1_in_features: 3
l1_out_features: 2
num_edge_features: 64
n_heads: 4
div: 4
SE3_ref_param:
num_layers: 2
num_channels: 32
num_degrees: 2
l0_in_features: 64
l0_out_features: 64
l1_in_features: 3
l1_out_features: 2
num_edge_features: 64
n_heads: 4
div: 4
"""
def convert_format(input_file, jobname, chain, deleteIndexes, attachmentIndex):
conv = openbabel.OBConversion()
conv.SetInAndOutFormats('cdjson', 'sdf')
# Add options
conv.AddOption("c", openbabel.OBConversion.OUTOPTIONS, "1")
with open(f"{jobname}_sm_{chain}.json", "w+") as fp:
fp.write(input_file)
mol = openbabel.OBMol()
conv.ReadFile(mol, f"{jobname}_sm_{chain}.json")
deleted_count = 0
# delete atoms in delete indexes
for index in sorted(deleteIndexes, reverse=True):
if index < attachmentIndex:
deleted_count += 1
atom = mol.GetAtom(index)
mol.DeleteAtom(atom)
attachmentIndex -= deleted_count
conv.WriteFile(mol, f"{jobname}_sm_{chain}.sdf")
return attachmentIndex
def prepare_input(input, jobname, baseconfig, hard_case):
input_categories = {"protein":"protein_inputs", "DNA":"na_inputs","RNA":"na_inputs", "ligand":"sm_inputs"}
# convert input to yaml format
yaml_dict = {"defaults":["base"], "job_name":jobname, "output_path": jobname}
list_of_input_files = []
if len(input["chains"]) == 0:
raise gr.Error("At least one chain must be provided")
for chain in input["chains"]:
if input_categories[chain["class"]] not in yaml_dict.keys():
yaml_dict[input_categories[chain["class"]]] = {}
if input_categories[chain["class"]] in ["protein_inputs", "na_inputs"]:
#write fasta
with open(f"{jobname}_{chain['chain']}.fasta", "w+") as fp:
fp.write(f">chain A\n{chain['sequence']}")
if input_categories[chain["class"]] == "na_inputs":
entry = {"input_type":chain["class"].lower(), "fasta":f"{jobname}/{jobname}_{chain['chain']}.fasta"}
else:
entry = {"fasta_file": f"{jobname}/{jobname}_{chain['chain']}.fasta"}
list_of_input_files.append(f"{jobname}_{chain['chain']}.fasta")
yaml_dict[input_categories[chain["class"]]][chain['chain']] = entry
if input_categories[chain['class']] == "sm_inputs":
if "smiles" in chain.keys():
entry = {"input_type": "smiles", "input": chain["smiles"]}
elif "sdf" in chain.keys():
# write to file
with open(f"{jobname}_sm_{chain['chain']}.sdf", "w+") as fp:
fp.write(chain["sdf"])
list_of_input_files.append(f"{jobname}_sm_{chain['chain']}.sdf")
entry = {"input_type": "sdf", "input": f"{jobname}/{jobname}_sm_{chain['chain']}.sdf"}
elif "name" in chain.keys():
list_of_input_files.append(f"metal_sdf/{chain['name']}_ideal.sdf")
entry = {"input_type": "sdf", "input": f"{jobname}/{chain['name']}_ideal.sdf"}
yaml_dict["sm_inputs"][chain['chain']] = entry
covale_inputs = []
if len(input["covMods"])>0:
yaml_dict["covale_inputs"]=""
for covMod in input["covMods"]:
if len(covMod["deleteIndexes"])>0:
new_attachment_index = convert_format(covMod["mol"],jobname, covMod["ligand"], covMod["deleteIndexes"], covMod["attachmentIndex"])
chirality_ligand = "null"
chirality_protein = "null"
if covMod["protein_symmetry"] in ["CW", "CCW"]:
chirality_protein = covMod["protein_symmetry"]
if covMod["ligand_symmetry"] in ["CW", "CCW"]:
chirality_ligand = covMod["ligand_symmetry"]
covale_inputs.append(((covMod[ "protein"], covMod["residue"], covMod["atom"]), (covMod["ligand"], new_attachment_index), (chirality_protein, chirality_ligand)))
if len(input["covMods"])>0:
yaml_dict["covale_inputs"] = json.dumps(json.dumps(covale_inputs))[1:-1].replace("'", "\"")
if hard_case:
yaml_dict["loader_params"]= {}
yaml_dict["loader_params"]["MAXCYCLE"] = 10
# write yaml to tmp
with open(f"/tmp/{jobname}.yaml", "w+") as fp:
# need to convert single quotes to double quotes
fp.write(yaml.dump(yaml_dict).replace("'", "\""))
# write baseconfig
with open(f"/tmp/base.yaml", "w+") as fp:
fp.write(baseconfig)
list_of_input_files.append(f"/tmp/{jobname}.yaml")
list_of_input_files.append(f"/tmp/base.yaml")
# convert dictionary to YAML
with zipfile.ZipFile(os.path.join("/tmp/", f"{jobname}.zip"), 'w') as zip_archive:
for file in set(list_of_input_files):
zip_archive.write(file, arcname= os.path.join(jobname,os.path.basename(file)),compress_type=zipfile.ZIP_DEFLATED)
return yaml.dump(yaml_dict).replace("'", "\""),os.path.join("/tmp/", f"{jobname}.zip")
def convert_bfactors(pdb_path):
with open(pdb_path, 'r') as f:
lines = f.readlines()
for i,line in enumerate(lines):
# multiple each bfactor by 100
if line[0:6] == 'ATOM ' or line[0:6] == 'HETATM':
bfactor = float(line[60:66])
bfactor *= 100
line = line[:60] + f'{bfactor:6.2f}' + line[66:]
lines[i] = line
with open(pdb_path.replace(".pdb", "_processed.pdb"), 'w') as f:
f.write(''.join(lines))
def run_rf2aa(jobname, zip_archive):
current_dir = os.getcwd()
try:
with zipfile.ZipFile(zip_archive, 'r') as zip_ref:
zip_ref.extractall(os.path.join(current_dir))
os.system(f"python -m rf2aa.run_inference --config-name {jobname}.yaml --config-path {current_dir}/{jobname}")
# scale pLDDT to 0-100 range in pdb output file
convert_bfactors(f"{current_dir}/{jobname}/{jobname}.pdb")
aux = torch.load(f"{current_dir}/{jobname}/{jobname}_aux.pt")
fig_pde = px.imshow(aux["pde"][0], template="simple_white", labels={"x": "Scored residue", "y": "Aligned residue", "color":"PDE",})
fig_pde.update_layout(coloraxis_colorbar=dict(title="PDE (Å)"))
fig_pae = px.imshow(aux["pae"][0], template="simple_white", labels={"x": "Scored residue", "y": "Aligned residue", "color":"PAE",})
fig_pae.update_layout(coloraxis_colorbar=dict(title="PAE (Å)"))
fig_plddt = px.line(y=aux["plddts"].flatten().numpy()*100, template="simple_white", labels={"y": "pLDDT", "x":"residue"})
except Exception as e:
raise gr.Error(f"Error running RFAA: {e}")
return f"{current_dir}/{jobname}/{jobname}_processed.pdb", fig_plddt, fig_pae, fig_pde
def predict(input, jobname, dry_run, baseconfig, hard_case):
yaml_input, zip_archive = prepare_input(input, jobname, baseconfig, hard_case)
reps = []
for chain in input["chains"]:
if chain["class"] in ["protein", "RNA", "DNA"]:
reps.append({
"model": 0,
"chain": chain["chain"],
"resname": "",
"style": "cartoon",
"color": "alphafold",
"residue_range": "",
"around": 0,
"byres": False
})
elif chain["class"] == "ligand" and "name" not in chain.keys():
reps.append({
"model": 0,
"chain": chain["chain"],
"resname": "LG1",
"style": "stick",
"color": "whiteCarbon",
"residue_range": "",
"around": 0,
"byres": False
})
else:
reps.append({
"model": 0,
"chain": chain["chain"],
"resname": "LG1",
"style": "sphere",
"color": "whiteCarbon",
"residue_range": "",
"around": 0,
"byres": False
})
if dry_run:
return gr.Code(yaml_input, visible=True), gr.File(zip_archive, visible=True), gr.Markdown(f"""You can run your RFAA job using the following command: <pre>python -m rf2aa.run_inference --config-name {jobname}.yaml --config-path absolute/path/to/unzipped/{jobname}</pre>""", visible=True), Molecule3D(visible=False), gr.Plot(visible=False), gr.Plot(visible=False), gr.Plot(visible=False)
else:
pdb_file, pldtt_plot, pae_plot, pde_plot = run_rf2aa(jobname, zip_archive)
return gr.Code(yaml_input, visible=True), gr.File(zip_archive, visible=True),gr.Markdown(visible=False), Molecule3D(pdb_file,reps=reps,visible=True), gr.Plot(pldtt_plot, visible=True), gr.Plot(pae_plot, visible=True), gr.Plot(pde_plot, visible=True)
with gr.Blocks() as demo:
gr.Markdown("# RoseTTAFold All Atom UI")
gr.Markdown("""This UI allows you to generate input files for RoseTTAFold All Atom (RFAA) using the CofoldingInput widget. The input files can be used to run RFAA on your local machine. <br />
If you launch the UI directly on your local machine you can also directly run the RFAA prediction. <br />
More information in the official GitHub repository: [baker-laboratory/RoseTTAFold-All-Atom](https://github.com/baker-laboratory/RoseTTAFold-All-Atom)
""")
jobname = gr.Textbox("job1", label="Job Name")
with gr.Tab("Input"):
inp=CofoldingInput(label="Input")
hard_case = gr.Checkbox(False, label="Hard case (increase MAXCYCLE to 10)")
# only allow running the predictions if local
if os.environ.get("SPACE_HOST")!=None:
dry_run = gr.Checkbox(True, label="Only generate input files (dry run)", interactive=False)
else:
dry_run = gr.Checkbox(True, label="Only generate input files (dry run)")
with gr.Tab("Base config"):
base_config = gr.Code(baseconfig, label="Base config")
btn = gr.Button("Run")
config_file = gr.Code(label="YAML Hydra config for RFAA", visible=True)
runfiles = gr.File(label="files to run RFAA", visible=False)
instructions = gr.Markdown(visible=False)
out = Molecule3D(visible=False, label="Predicted Structure")
with gr.Row():
plddt = gr.Plot(visible=False, label="pLDDT")
pae = gr.Plot(visible=False, label="Predicted aligned error")
pde = gr.Plot(visible=False, label="Predicted distance error")
btn.click(predict, inputs=[inp, jobname, dry_run, base_config, hard_case], outputs=[config_file, runfiles, instructions, out, plddt, pae, pde])
if __name__ == "__main__":
demo.launch(share=True)