diff --git a/openfecli/commands/plan_rbfe_network.py b/openfecli/commands/plan_rbfe_network.py index 85302376a..43a2f3eab 100644 --- a/openfecli/commands/plan_rbfe_network.py +++ b/openfecli/commands/plan_rbfe_network.py @@ -7,6 +7,7 @@ from openfecli import OFECommandPlugin from openfecli.parameters import ( MOL_DIR, PROTEIN, MAPPER, OUTPUT_DIR, COFACTORS, + NEW_STORAGE_OUTPUT, # separate line for easy delete later ) from openfecli.plan_alchemical_networks_utils import plan_alchemical_network_output @@ -82,9 +83,11 @@ def plan_rbfe_network_main( help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.", default="alchemicalNetwork", ) +@NEW_STORAGE_OUTPUT.parameter() @print_duration def plan_rbfe_network( - molecules: List[str], protein: str, cofactors: tuple[str], output_dir: str + molecules: List[str], protein: str, cofactors: tuple[str], + output_dir: str, new_storage: bool, ): """ Plan a relative binding free energy network, saved as JSON files for @@ -119,6 +122,8 @@ def plan_rbfe_network( from openfe.setup.atom_mapping.lomap_scorers import ( default_lomap_score, ) + from openfe.storage.resultclient import ResultClient + from gufe.storage.externalresource import FileStorage from openfe.setup import LomapAtomMapper from openfe.setup.ligand_network_planning import ( generate_minimal_spanning_network, @@ -175,12 +180,18 @@ def plan_rbfe_network( # OUTPUT write("Output:") - write("\tSaving to: " + str(output_dir)) - plan_alchemical_network_output( - alchemical_network=alchemical_network, - ligand_network=ligand_network, - folder_path=OUTPUT_DIR.get(output_dir), - ) + if new_storage: + write(f"Saving to: {output_dir}") + storage = ResultClient(FileStorage(output_dir)) + storage.store_network(alchemical_network) + # TODO: save out the ligand network as well + else: + write("\tSaving to: " + str(output_dir)) + plan_alchemical_network_output( + alchemical_network=alchemical_network, + ligand_network=ligand_network, + folder_path=OUTPUT_DIR.get(output_dir), + ) PLUGIN = OFECommandPlugin( diff --git a/openfecli/parameters/__init__.py b/openfecli/parameters/__init__.py index 8de7af8ab..336fb6da3 100644 --- a/openfecli/parameters/__init__.py +++ b/openfecli/parameters/__init__.py @@ -4,6 +4,6 @@ from .mol import MOL from .mapper import MAPPER from .output import OUTPUT_FILE_AND_EXT -from .output_dir import OUTPUT_DIR +from .output_dir import OUTPUT_DIR, NEW_STORAGE_OUTPUT from .protein import PROTEIN from .molecules import MOL_DIR, COFACTORS diff --git a/openfecli/parameters/output_dir.py b/openfecli/parameters/output_dir.py index 95c271aa6..a07915d4c 100644 --- a/openfecli/parameters/output_dir.py +++ b/openfecli/parameters/output_dir.py @@ -16,3 +16,10 @@ def get_dir(user_input, context): getter=get_dir, type=click.Path(file_okay=False, resolve_path=True), ) + +NEW_STORAGE_OUTPUT = Option( + "--new-storage", + help="use the new storage", + is_flag=True, + hidden=True, +)