Multi-leg HybridTop Protocols

src/openfe/protocols/openmm_afe/afe_protocol_results.py

@@ -31,16 +31,21 @@
 
 
 class AbsoluteProtocolResultMixin:
-    bound_state = "solvent"
-    unbound_state = "vacuum"
+    """
+    Subclasses must define the class attributes ``env_state`` and
+    ``ref_state``, naming the two legs of the thermodynamic cycle stored in
+    ``self.data`` (e.g. ``env_state = "complex"``, ``ref_state = "solvent"``).
+    """
+    env_state: str
+    ref_state: str
 
     def __init__(self, **data):
         super().__init__(**data)
         # TODO: Detect when we have extensions and stitch these together?
         if any(
             len(pur_list) > 2
             for pur_list in itertools.chain(
-                self.data[self.bound_state].values(), self.data[self.unbound_state].values()
+                self.data[self.env_state].values(), self.data[self.ref_state].values()
             )
         ):
             raise NotImplementedError("Can't stitch together results yet")
@@ -49,15 +54,16 @@ def get_forward_and_reverse_energy_analysis(
         self,
     ) -> dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]]:
         """
-        Get the reverse and forward analysis of the free energies.
+        Get the reverse and forward analysis of the free energies for both
+        legs of the thermodynamic cycle.
 
         Returns
         -------
         forward_reverse : dict[str, list[Optional[dict[str, Union[npt.NDArray, openff.units.Quantity]]]]]
             A dictionary, keyed for each leg of the thermodynamic cycle,
             either ``solvent`` and ``vacuum` for a solvation free energy or
           ``solvent`` and ``complex`` for a binding free energy,
-            with each containing a list of dictionaries containing the forward
+            with each containing a list of dictionaries with the forward
             and reverse analysis of each repeat of that simulation type.
 
             The forward and reverse analysis dictionaries contain:
@@ -87,7 +93,7 @@ def get_forward_and_reverse_energy_analysis(
 
         forward_reverse: dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]] = {}
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             forward_reverse[key] = [
                 pus[0].outputs["forward_and_reverse_energies"]
                 for pus in self.data[key].values()  # type: ignore[attr-defined]
@@ -128,7 +134,7 @@ def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]:
         # Loop through and get the repeats and get the matrices
         overlap_stats: dict[str, list[dict[str, npt.NDArray]]] = {}
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             overlap_stats[key] = [
                 pus[0].outputs["unit_mbar_overlap"]
                 for pus in self.data[key].values()  # type: ignore[attr-defined]
@@ -138,8 +144,8 @@ def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]:
 
     def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDArray]]]:
         """
-        Get the replica exchange transition statistics for all
-        legs of the simulation.
+        Get the replica exchange transition statistics for both
+        legs of the thermodynamic cycle.
 
         Note
         ----
@@ -163,7 +169,7 @@ def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDAr
         """
         repex_stats: dict[str, list[dict[str, npt.NDArray]]] = {}
         try:
-            for key in [self.bound_state, self.unbound_state]:
+            for key in [self.env_state, self.ref_state]:
                 repex_stats[key] = [
                     pus[0].outputs["replica_exchange_statistics"]
                     for pus in self.data[key].values()  # type: ignore[attr-defined]
@@ -182,14 +188,14 @@ def get_replica_states(self) -> dict[str, list[npt.NDArray]]:
         -------
         replica_states : dict[str, list[npt.NDArray]]
           Dictionary keyed for each leg of the thermodynamic cycle, either
-          `solvent` and `vacuum` for solvation free energies,
-          or `complex` and `solvent` for binding free energies,
+          ``solvent`` and ``vacuum`` for solvation free energies,
+          or ``complex`` and ``solvent`` for binding free energies,
           with lists of replica states timeseries for each repeat of that
           simulation type.
         """
         replica_states: dict[str, list[npt.NDArray]] = {
-            self.bound_state: [],
-            self.unbound_state: [],
+            self.env_state: [],
+            self.ref_state: [],
         }
 
         def is_file(filename: str):
@@ -215,7 +221,7 @@ def get_replica_state(nc, chk):
 
             return retval
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             for pus in self.data[key].values():  # type: ignore[attr-defined]
                 states = get_replica_state(
                     pus[0].outputs["trajectory"],
@@ -227,20 +233,21 @@ def get_replica_state(nc, chk):
 
     def equilibration_iterations(self) -> dict[str, list[float]]:
         """
-        Get the number of equilibration iterations for each simulation.
+        Get the number of equilibration iterations for each repeat of
+        both legs of the calculation.
 
         Returns
         -------
         equilibration_lengths : dict[str, list[float]]
           Dictionary keyed for each leg of the thermodynamic cycle, either
-         `solvent` and `vacuum` for solvation free energies,
-          or `complex` and `solvent` for binding free energies,
+          ``solvent`` and ``vacuum`` for solvation free energies,
+          or ``complex`` and ``solvent`` for binding free energies,
           with lists containing the number of equilibration iterations for
           each repeat of that simulation type.
         """
         equilibration_lengths: dict[str, list[float]] = {}
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             equilibration_lengths[key] = [
                 pus[0].outputs["equilibration_iterations"]
                 for pus in self.data[key].values()  # type: ignore[attr-defined]
@@ -250,22 +257,21 @@ def equilibration_iterations(self) -> dict[str, list[float]]:
 
     def production_iterations(self) -> dict[str, list[float]]:
         """
-        Get the number of production iterations for each simulation.
         Returns the number of uncorrelated production samples for each
-        repeat of the calculation.
+        repeat of both legs of the calculation.
 
         Returns
         -------
         production_lengths : dict[str, list[float]]
           Dictionary keyed for each leg of the thermodynamic cycle, either
-         `solvent` and `vacuum` for solvation free energies,
-          or `complex` and `solvent` for binding free energies,
+         ``solvent`` and ``vacuum`` for solvation free energies,
+          or ``complex`` and ``solvent`` for binding free energies,
           with lists containing the number of equilibration iterations for
           each repeat of that simulation type.
         """
         production_lengths: dict[str, list[float]] = {}
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             production_lengths[key] = [
                 pus[0].outputs["production_iterations"]
                 for pus in self.data[key].values()  # type: ignore[attr-defined]
@@ -282,15 +288,15 @@ def selection_indices(self) -> dict[str, list[Optional[npt.NDArray]]]:
         -------
         indices : dict[str, list[npt.NDArray]]
           A dictionary keyed for each state, either
-         `solvent` and `vacuum` for solvation free energies,
-          or `complex` and `solvent` for binding free energies,
+         ``solvent`` and ``vacuum`` for solvation free energies,
+          or ``complex`` and ``solvent`` for binding free energies,
           each containing a list of NDArrays containing the corresponding
-          full system atom indices for each atom written in the production
-          trajectory files for each replica.
+          full system atom indices for each atom written in the PDB or
+          production trajectory files for each replica.
         """
         indices: dict[str, list[Optional[npt.NDArray]]] = {}
 
-        for key in [self.bound_state, self.unbound_state]:
+        for key in [self.env_state, self.ref_state]:
             indices[key] = []
             for pus in self.data[key].values():  # type: ignore[attr-defined]
                 indices[key].append(pus[0].outputs["selection_indices"])
@@ -303,8 +309,8 @@ class AbsoluteSolvationProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResul
     Protocol results with the output of a AbsoluteSolvationProtocol
     """
 
-    bound_state = "solvent"
-    unbound_state = "vacuum"
+    env_state = "solvent"
+    ref_state = "vacuum"
 
     def get_individual_estimates(self) -> dict[str, list[tuple[Quantity, Quantity]]]:
         """
@@ -320,7 +326,7 @@ def get_individual_estimates(self) -> dict[str, list[tuple[Quantity, Quantity]]]
         """
         dGs = {}
 
-        for state in [self.bound_state, self.unbound_state]:
+        for state in [self.env_state, self.ref_state]:
             state_dGs = [
                 (pus[0].outputs["unit_estimate"], pus[0].outputs["unit_estimate_error"])
                 for pus in self.data[state].values()
@@ -389,8 +395,8 @@ class AbsoluteBindingProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResultM
     Protocol results with the output of a AbsoluteBindingProtocol.
     """
 
-    bound_state = "complex"
-    unbound_state = "solvent"
+    env_state = "complex"
+    ref_state = "solvent"
 
     def get_individual_estimates(
         self,

src/openfe/protocols/openmm_rfe/__init__.py

@@ -2,11 +2,35 @@
 # For details, see https://github.com/OpenFreeEnergy/openfe
 
 from . import _rfe_utils
-from .equil_rfe_settings import RelativeHybridTopologyProtocolSettings
-from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
-from .hybridtop_protocols import RelativeHybridTopologyProtocol
+from .equil_rfe_settings import (
+    RBFEHTopProtocolSettings,
+    RelativeHybridTopologyProtocolSettings,
+    RHFEHTopProtocolSettings,
+)
+from .hybridtop_protocol_results import (
+    RBFEHTopProtocolResult,
+    RelativeHybridTopologyProtocolResult,
+    RHFEHTopProtocolResult,
+)
+from .hybridtop_protocols import (
+    RBFEHTopProtocol,
+    RelativeHybridTopologyProtocol,
+    RHFEHTopProtocol,
+)
 from .hybridtop_units import (
     HybridTopologyMultiStateAnalysisUnit,
     HybridTopologyMultiStateSimulationUnit,
     HybridTopologySetupUnit,
+    RBFEHTopComplexAnalysisUnit,
+    RBFEHTopComplexSetupUnit,
+    RBFEHTopComplexSimulationUnit,
+    RBFEHTopSolventAnalysisUnit,
+    RBFEHTopSolventSetupUnit,
+    RBFEHTopSolventSimulationUnit,
+    RHFEHTopSolventAnalysisUnit,
+    RHFEHTopSolventSetupUnit,
+    RHFEHTopSolventSimulationUnit,
+    RHFEHTopVacuumAnalysisUnit,
+    RHFEHTopVacuumSetupUnit,
+    RHFEHTopVacuumSimulationUnit,
 )