add support for one repeat per json

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:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

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,

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("")

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("")

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

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.",
+)

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():

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