Create setup, run, and analysis units for HybridTop Protocol

openfe/protocols/openmm_rfe/__init__.py

@@ -5,4 +5,8 @@
 from .equil_rfe_settings import RelativeHybridTopologyProtocolSettings
 from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
 from .hybridtop_protocols import RelativeHybridTopologyProtocol
-from .hybridtop_units import RelativeHybridTopologyProtocolUnit
+from .hybridtop_units import (
+    HybridTopologyMultiStateAnalysisUnit,
+    HybridTopologyMultiStateSimulationUnit,
+    HybridTopologySetupUnit,
+)

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -19,4 +19,8 @@
 from .equil_rfe_settings import RelativeHybridTopologyProtocolSettings
 from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
 from .hybridtop_protocols import RelativeHybridTopologyProtocol
-from .hybridtop_units import RelativeHybridTopologyProtocolUnit
+from .hybridtop_units import (
+    HybridTopologyMultiStateAnalysisUnit,
+    HybridTopologyMultiStateSimulationUnit,
+    HybridTopologySetupUnit,
+)

openfe/protocols/openmm_rfe/hybridtop_protocol_results.py

@@ -33,7 +33,7 @@ def __init__(self, **data):
     def compute_mean_estimate(dGs: list[Quantity]) -> Quantity:
         u = dGs[0].u
         # convert all values to units of the first value, then take average of magnitude
-        # this would avoid a screwy case where each value was in different units
+        # this would avoid an edge case where each value was in different units
         vals = np.asarray([dG.to(u).m for dG in dGs])
 
         return np.average(vals) * u
@@ -199,9 +199,9 @@ def is_file(filename: str):
         replica_states = []
 
         for pus in self.data.values():
-            nc = is_file(pus[0].outputs["nc"])
+            nc = is_file(pus[0].outputs["trajectory"])
             dir_path = nc.parents[0]
-            chk = is_file(dir_path / pus[0].outputs["last_checkpoint"]).name
+            chk = is_file(pus[0].outputs["checkpoint"]).name
             reporter = multistate.MultiStateReporter(
                 storage=nc, checkpoint_storage=chk, open_mode="r"
             )

openfe/protocols/openmm_rfe/hybridtop_protocols.py

@@ -49,7 +49,11 @@
     RelativeHybridTopologyProtocolSettings,
 )
 from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
-from .hybridtop_units import RelativeHybridTopologyProtocolUnit
+from .hybridtop_units import (
+    HybridTopologyMultiStateAnalysisUnit,
+    HybridTopologyMultiStateSimulationUnit,
+    HybridTopologySetupUnit,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -593,23 +597,47 @@ def _create(
         Anames = ",".join(c.name for c in alchem_comps["stateA"])
         Bnames = ",".join(c.name for c in alchem_comps["stateB"])
 
-        # our DAG has no dependencies, so just list units
-        n_repeats = self.settings.protocol_repeats
+        # DAG dependency is setup -> simulation -> analysis
+        #                     |--------------------->
+        setup_units = []
+        simulation_units = []
+        analysis_units = []
+
+        for i in range(self.settings.protocol_repeats):
+            repeat_id = int(uuid.uuid4())
 
-        units = [
-            RelativeHybridTopologyProtocolUnit(
+            setup = HybridTopologySetupUnit(
                 protocol=self,
                 stateA=stateA,
                 stateB=stateB,
                 ligandmapping=ligandmapping,
+                alchemical_components=alchem_comps,
+                generation=0,
+                repeat_id=repeat_id,
+                name=(f"HybridTopology Setup: {Anames} to {Bnames} repeat {i} generation 0"),
+            )
+
+            simulation = HybridTopologyMultiStateSimulationUnit(
+                protocol=self,
+                setup_results=setup,
+                generation=0,
+                repeat_id=repeat_id,
+                name=(f"HybridTopology Simulation: {Anames} to {Bnames} repeat {i} generation 0"),
+            )
+
+            analysis = HybridTopologyMultiStateAnalysisUnit(
+                protocol=self,
+                setup_results=setup,
+                simulation_results=simulation,
                 generation=0,
-                repeat_id=int(uuid.uuid4()),
-                name=f"{Anames} to {Bnames} repeat {i} generation 0",
+                repeat_id=repeat_id,
+                name=(f"HybridTopology Analysis: {Anames} to {Bnames} repeat {i} generation 0"),
             )
-            for i in range(n_repeats)
-        ]
+            setup_units.append(setup)
+            simulation_units.append(simulation)
+            analysis_units.append(analysis)
 
-        return units
+        return [*setup_units, *simulation_units, *analysis_units]
 
     def _gather(self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult]) -> dict[str, Any]:
         # result units will have a repeat_id and generations within this repeat_id
@@ -618,7 +646,8 @@ def _gather(self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult]) -> dic
         for d in protocol_dag_results:
             pu: gufe.ProtocolUnitResult
             for pu in d.protocol_unit_results:
-                if not pu.ok():
+                # We only need the analysis units that are ok
+                if ("Analysis" not in pu.name) or (not pu.ok()):
                     continue
 
                 unsorted_repeats[pu.outputs["repeat_id"]].append(pu)