Commit feb99932 authored by Vaclav Hanzl's avatar Vaclav Hanzl
Browse files

Make recently new cif_output arguments optional, do not break existing code

Colab notbook was broken by this change and potentially code of other users
as well. Keep backward compatibility by also allowing the old way.
parent b95099bb
...@@ -57,7 +57,7 @@ class AmberRelaxation(object): ...@@ -57,7 +57,7 @@ class AmberRelaxation(object):
self._use_gpu = use_gpu self._use_gpu = use_gpu
def process( def process(
self, *, prot: protein.Protein, cif_output: bool self, *, prot: protein.Protein, cif_output: bool = False
) -> Tuple[str, Dict[str, Any], np.ndarray]: ) -> Tuple[str, Dict[str, Any], np.ndarray]:
"""Runs Amber relax on a prediction, adds hydrogens, returns PDB string.""" """Runs Amber relax on a prediction, adds hydrogens, returns PDB string."""
out = amber_minimize.run_pipeline( out = amber_minimize.run_pipeline(
......
...@@ -228,7 +228,7 @@ def prep_output(out, batch, feature_dict, feature_processor, config_preset, mult ...@@ -228,7 +228,7 @@ def prep_output(out, batch, feature_dict, feature_processor, config_preset, mult
return unrelaxed_protein return unrelaxed_protein
def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name, cif_output): def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name, cif_output=False):
amber_relaxer = relax.AmberRelaxation( amber_relaxer = relax.AmberRelaxation(
use_gpu=(model_device != "cpu"), use_gpu=(model_device != "cpu"),
**config.relax, **config.relax,
...@@ -257,4 +257,4 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out ...@@ -257,4 +257,4 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out
with open(relaxed_output_path, 'w') as fp: with open(relaxed_output_path, 'w') as fp:
fp.write(struct_str) fp.write(struct_str)
logger.info(f"Relaxed output written to {relaxed_output_path}...") logger.info(f"Relaxed output written to {relaxed_output_path}...")
\ No newline at end of file
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