diff --git a/news/add_n_protocol_repeats.rst b/news/add_n_protocol_repeats.rst new file mode 100644 index 000000000..ffe3e9201 --- /dev/null +++ b/news/add_n_protocol_repeats.rst @@ -0,0 +1,23 @@ +**Added:** + +* Added ``--n-protocol-repeats`` CLI option to allow user-defined number of repeats per quickrun execution. This allows for parallelizing execution of repeats by setting ``--n-protocol-repeats=1`` and calling ``quickrun`` on the same input file multiple times. + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py b/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py index 526651488..cf79b2f3b 100644 --- a/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py +++ b/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py @@ -303,7 +303,7 @@ class RBFEAlchemicalNetworkPlanner(RelativeAlchemicalNetworkPlanner): """ def __init__( self, - name: str = "easy_rbfe", + name: str = "easy_rbfe", # TODO: change this default to "" in 2.0 mappers: Optional[Iterable[LigandAtomMapper]] = None, mapping_scorer: Callable[[LigandAtomMapping], float] = default_lomap_score, ligand_network_planner: Callable = generate_minimal_spanning_network, diff --git a/openfecli/commands/plan_rbfe_network.py b/openfecli/commands/plan_rbfe_network.py index 88beb4007..db1794793 100644 --- a/openfecli/commands/plan_rbfe_network.py +++ b/openfecli/commands/plan_rbfe_network.py @@ -5,7 +5,7 @@ from openfecli.utils import write, print_duration from openfecli import OFECommandPlugin from openfecli.parameters import ( - MOL_DIR, PROTEIN, OUTPUT_DIR, COFACTORS, YAML_OPTIONS, + MOL_DIR, PROTEIN, OUTPUT_DIR, COFACTORS, YAML_OPTIONS, N_PROTOCOL_REPEATS ) def plan_rbfe_network_main( @@ -16,6 +16,7 @@ def plan_rbfe_network_main( solvent, protein, cofactors, + n_protocol_repeats, ): """Utility method to plan a relative binding free energy network. @@ -34,7 +35,9 @@ def plan_rbfe_network_main( protein : ProteinComponent protein component for complex simulations, to which the ligands are bound cofactors : Iterable[SmallMoleculeComponent] - any cofactors alongisde the protein, can be empty list + any cofactors alongside the protein, can be empty list + n_protocol_repeats: int + number of completely independent repeats of the entire sampling process Returns ------- @@ -46,11 +49,17 @@ def plan_rbfe_network_main( from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import ( RBFEAlchemicalNetworkPlanner, ) + from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import RelativeHybridTopologyProtocol + + protocol_settings = RelativeHybridTopologyProtocol.default_settings() + protocol_settings.protocol_repeats = n_protocol_repeats + protocol = RelativeHybridTopologyProtocol(protocol_settings) network_planner = RBFEAlchemicalNetworkPlanner( mappers=mapper, mapping_scorer=mapping_scorer, ligand_network_planner=ligand_network_planner, + protocol=protocol ) alchemical_network = network_planner( ligands=small_molecules, solvent=solvent, protein=protein, @@ -83,11 +92,13 @@ def plan_rbfe_network_main( help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.", default="alchemicalNetwork", ) +@N_PROTOCOL_REPEATS.parameter(multiple=False, required=False, default=3, help=N_PROTOCOL_REPEATS.kwargs["help"]) @print_duration def plan_rbfe_network( molecules: list[str], protein: str, cofactors: tuple[str], yaml_settings: str, output_dir: str, + n_protocol_repeats: int, ): """ Plan a relative binding free energy network, saved as JSON files for @@ -169,6 +180,7 @@ def plan_rbfe_network( solvent=solvent, protein=protein, cofactors=cofactors, + n_protocol_repeats=n_protocol_repeats, ) write("\tDone") write("") diff --git a/openfecli/commands/plan_rhfe_network.py b/openfecli/commands/plan_rhfe_network.py index 90699d322..9620d3221 100644 --- a/openfecli/commands/plan_rhfe_network.py +++ b/openfecli/commands/plan_rhfe_network.py @@ -8,12 +8,12 @@ from openfecli.utils import write, print_duration from openfecli import OFECommandPlugin from openfecli.parameters import ( - MOL_DIR, MAPPER, OUTPUT_DIR, YAML_OPTIONS, + MOL_DIR, MAPPER, OUTPUT_DIR, YAML_OPTIONS, N_PROTOCOL_REPEATS ) def plan_rhfe_network_main( mapper, mapping_scorer, ligand_network_planner, small_molecules, - solvent, + solvent, n_protocol_repeats, ): """Utility method to plan a relative hydration free energy network. @@ -29,6 +29,8 @@ def plan_rhfe_network_main( molecules of the system solvent : SolventComponent Solvent component used for solvation + n_protocol_repeats: int + number of completely independent repeats of the entire sampling process Returns ------- @@ -39,11 +41,18 @@ def plan_rhfe_network_main( from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import ( RHFEAlchemicalNetworkPlanner ) + from openfe.setup.alchemical_network_planner.relative_alchemical_network_planner import RelativeHybridTopologyProtocol + + + protocol_settings = RelativeHybridTopologyProtocol.default_settings() + protocol_settings.protocol_repeats = n_protocol_repeats + protocol = RelativeHybridTopologyProtocol(protocol_settings) network_planner = RHFEAlchemicalNetworkPlanner( mappers=mapper, mapping_scorer=mapping_scorer, ligand_network_planner=ligand_network_planner, + protocol=protocol ) alchemical_network = network_planner( ligands=small_molecules, solvent=solvent @@ -70,8 +79,10 @@ def plan_rhfe_network_main( help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.", default="alchemicalNetwork", ) +@N_PROTOCOL_REPEATS.parameter(multiple=False, required=False, default=3, help=N_PROTOCOL_REPEATS.kwargs["help"]) + @print_duration -def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str): +def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str, n_protocol_repeats:int): """ Plan a relative hydration free energy network, saved as JSON files for the quickrun command. @@ -143,6 +154,7 @@ def plan_rhfe_network(molecules: List[str], yaml_settings: str, output_dir: str) ligand_network_planner=ligand_network_planner, small_molecules=small_molecules, solvent=solvent, + n_protocol_repeats=n_protocol_repeats, ) write("\tDone") write("") diff --git a/openfecli/parameters/__init__.py b/openfecli/parameters/__init__.py index f5a1cfa6f..171f0f84e 100644 --- a/openfecli/parameters/__init__.py +++ b/openfecli/parameters/__init__.py @@ -8,3 +8,4 @@ from .protein import PROTEIN from .molecules import MOL_DIR, COFACTORS from .plan_network_options import YAML_OPTIONS +from .misc import N_PROTOCOL_REPEATS diff --git a/openfecli/parameters/misc.py b/openfecli/parameters/misc.py new file mode 100644 index 000000000..41df1e0a4 --- /dev/null +++ b/openfecli/parameters/misc.py @@ -0,0 +1,11 @@ +import click +from plugcli.params import Option + +N_PROTOCOL_REPEATS = Option( + "--n-protocol-repeats", + type=click.INT, + help="Number of independent repeats(s) to be run per execution of a transformation using the `openfe quickrun` command.\n\n" + "For example:\n\n `--n-protocol-repeats=3` means `openfe quickrun` will execute 3 repeats in serial.\n\n" + " `--n-protocol-repeats=1` means `openfe quickrun` will execute only 1 repeat per call, " + "which allows for individual repeats to be submitted in parallel by calling `openfe quickrun` on the same input JSON file multiple times.", +) diff --git a/openfecli/tests/commands/test_plan_rbfe_network.py b/openfecli/tests/commands/test_plan_rbfe_network.py index 25b703811..3fe89128d 100644 --- a/openfecli/tests/commands/test_plan_rbfe_network.py +++ b/openfecli/tests/commands/test_plan_rbfe_network.py @@ -4,11 +4,13 @@ from importlib import resources import shutil from click.testing import CliRunner +from ..utils import assert_click_success from openfecli.commands.plan_rbfe_network import ( plan_rbfe_network, plan_rbfe_network_main, ) + from gufe import AlchemicalNetwork from gufe.tokenization import JSON_HANDLER import json @@ -63,7 +65,7 @@ def test_plan_rbfe_network_main(): for f in ['ligand_23.sdf', 'ligand_55.sdf'] ] with resources.files("openfe.tests.data") as d: - protein_compontent = ProteinComponent.from_pdb_file( + protein_component = ProteinComponent.from_pdb_file( str(d / "181l_only.pdb") ) @@ -74,8 +76,9 @@ def test_plan_rbfe_network_main(): ligand_network_planner=ligand_network_planning.generate_minimal_spanning_network, small_molecules=smallM_components, solvent=solvent_component, - protein=protein_compontent, + protein=protein_component, cofactors=[], + n_protocol_repeats=3, ) print(alchemical_network) @@ -123,6 +126,22 @@ def test_plan_rbfe_network(mol_dir_args, protein_args): for l1, l2 in zip(expected_output_1, expected_output_2): assert l1 in result.output or l2 in result.output +@pytest.mark.parametrize(['input_n_repeat', 'expected_n_repeat'], [([], 3), (["--n-protocol-repeats", "1"], 1)]) +def test_plan_rbfe_network_n_repeats(mol_dir_args, protein_args, input_n_repeat, expected_n_repeat): + runner = CliRunner() + + args = mol_dir_args + protein_args + input_n_repeat + + with runner.isolated_filesystem(): + result = runner.invoke(plan_rbfe_network, args) + assert_click_success(result) + + # make sure the number of repeats is correct + network = AlchemicalNetwork.from_dict( + json.load(open("alchemicalNetwork/alchemicalNetwork.json"), cls=JSON_HANDLER.decoder) + ) + for edge in network.edges: + assert edge.protocol.settings.protocol_repeats == expected_n_repeat @pytest.fixture def eg5_files(): diff --git a/openfecli/tests/commands/test_plan_rhfe_network.py b/openfecli/tests/commands/test_plan_rhfe_network.py index f5a72431f..97ee497a1 100644 --- a/openfecli/tests/commands/test_plan_rhfe_network.py +++ b/openfecli/tests/commands/test_plan_rhfe_network.py @@ -11,7 +11,6 @@ plan_rhfe_network_main, ) - @pytest.fixture(scope='session') def mol_dir_args(tmpdir_factory): ofe_dir_path = tmpdir_factory.mktemp('moldir') @@ -54,6 +53,7 @@ def test_plan_rhfe_network_main(): ligand_network_planner=ligand_network_planning.generate_minimal_spanning_network, small_molecules=smallM_components, solvent=solvent_component, + n_protocol_repeats=3, ) assert alchemical_network