Skip to content

GitLab

Unverified Commit 71735d5d authored by double-vin's avatar double-vin Committed by GitHub
Browse files

added relaxation argument (#146) 


Co-authored-by: default avatarzhuww <zhuww@sugon.com>
parent 4dc719f7
Hide whitespace changes
Inline Side-by-side
Showing with 31 additions and 26 deletions
inference.py
View file @ 71735d5d
...@@ -306,21 +306,22 @@ def inference_multimer_model(args): ...@@ -306,21 +306,22 @@ def inference_multimer_model(args):
with open(unrelaxed_output_path, 'w') as f: with open(unrelaxed_output_path, 'w') as f:
f.write(protein.to_pdb(unrelaxed_protein)) f.write(protein.to_pdb(unrelaxed_protein))
amber_relaxer = relax.AmberRelaxation( if(args.relaxation):
use_gpu=True, amber_relaxer = relax.AmberRelaxation(
**config.relax, use_gpu=True,
) **config.relax,
)
# Relax the prediction. # Relax the prediction.
t = time.perf_counter() t = time.perf_counter()
relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein) relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
print(f"Relaxation time: {time.perf_counter() - t}") print(f"Relaxation time: {time.perf_counter() - t}")
# Save the relaxed PDB. # Save the relaxed PDB.
relaxed_output_path = os.path.join(args.output_dir, relaxed_output_path = os.path.join(args.output_dir,
f'{output_prefix}_{args.model_name}_relaxed.pdb') f'{output_prefix}_{args.model_name}_relaxed.pdb')
with open(relaxed_output_path, 'w') as f: with open(relaxed_output_path, 'w') as f:
f.write(relaxed_pdb_str) f.write(relaxed_pdb_str)
def inference_monomer_model(args): def inference_monomer_model(args):
...@@ -445,21 +446,22 @@ def inference_monomer_model(args): ...@@ -445,21 +446,22 @@ def inference_monomer_model(args):
with open(unrelaxed_output_path, 'w') as f: with open(unrelaxed_output_path, 'w') as f:
f.write(protein.to_pdb(unrelaxed_protein)) f.write(protein.to_pdb(unrelaxed_protein))
amber_relaxer = relax.AmberRelaxation( if(args.relaxation):
use_gpu=True, amber_relaxer = relax.AmberRelaxation(
**config.relax, use_gpu=True,
) **config.relax,
)
# Relax the prediction. # Relax the prediction.
t = time.perf_counter() t = time.perf_counter()
relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein) relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
print(f"Relaxation time: {time.perf_counter() - t}") print(f"Relaxation time: {time.perf_counter() - t}")
# Save the relaxed PDB. # Save the relaxed PDB.
relaxed_output_path = os.path.join(args.output_dir, relaxed_output_path = os.path.join(args.output_dir,
f'{tag}_{args.model_name}_relaxed.pdb') f'{tag}_{args.model_name}_relaxed.pdb')
with open(relaxed_output_path, 'w') as f: with open(relaxed_output_path, 'w') as f:
f.write(relaxed_pdb_str) f.write(relaxed_pdb_str)
if __name__ == "__main__": if __name__ == "__main__":
...@@ -494,6 +496,9 @@ if __name__ == "__main__": ...@@ -494,6 +496,9 @@ if __name__ == "__main__":
help="""Path to model parameters. If None, parameters are selected help="""Path to model parameters. If None, parameters are selected
automatically according to the model name from automatically according to the model name from
./data/params""") ./data/params""")
parser.add_argument(
"--relaxation", action="store_false", default=False,
)
parser.add_argument("--cpus", parser.add_argument("--cpus",
type=int, type=int,
default=12, default=12,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment