OpenFreeEnergy · IAlibay · Jan 22, 2026 · Jan 1, 2026 · Jan 1, 2026 · Jan 1, 2026
diff --git a/docs/reference/api/openmm_binding_afe.rst b/docs/reference/api/openmm_binding_afe.rst
@@ -16,8 +16,12 @@ Protocol API specification
    :toctree: generated/
 
    AbsoluteBindingProtocol
-   AbsoluteBindingComplexUnit
-   AbsoluteBindingSolventUnit
+   ABFEComplexAnalysisUnit
+   ABFEComplexSetupUnit
+   ABFEComplexSimUnit
+   ABFESolventAnalysisUnit
+   ABFESolventSetupUnit
+   ABFESolventSimUnit
    AbsoluteBindingProtocolResult
 
 Protocol Settings

diff --git a/docs/reference/api/openmm_solvation_afe.rst b/docs/reference/api/openmm_solvation_afe.rst
@@ -16,8 +16,12 @@ Protocol API specification
    :toctree: generated/
 
    AbsoluteSolvationProtocol
-   AbsoluteSolvationVacuumUnit
-   AbsoluteSolvationSolventUnit
+   AHFESolventAnalysisUnit
+   AHFESolventSetupUnit
+   AHFESolventSimUnit
+   AHFEVacuumAnalysisUnit
+   AHFEVacuumSetupUnit
+   AHFEVacuumSimUnit
    AbsoluteSolvationProtocolResult
 
 Protocol Settings

diff --git a/news/multi-unit-afe.rst b/news/multi-unit-afe.rst
@@ -0,0 +1,26 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* The absolute free energy protocols have been broken into multiple
+  protocol units, allowing for setup, run, and analysis to happen
+  separately in the future when relevant changes to protocol execution are
+  made (PR #1776).
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>
diff --git a/openfe/protocols/openmm_afe/__init__.py b/openfe/protocols/openmm_afe/__init__.py
@@ -6,16 +6,24 @@
 """
 
 from .abfe_units import (
-    AbsoluteBindingComplexUnit,
-    AbsoluteBindingSolventUnit,
+    ABFEComplexAnalysisUnit,
+    ABFEComplexSetupUnit,
+    ABFEComplexSimUnit,
+    ABFESolventAnalysisUnit,
+    ABFESolventSetupUnit,
+    ABFESolventSimUnit,
 )
 from .afe_protocol_results import (
     AbsoluteBindingProtocolResult,
     AbsoluteSolvationProtocolResult,
 )
 from .ahfe_units import (
-    AbsoluteSolvationSolventUnit,
-    AbsoluteSolvationVacuumUnit,
+    AHFESolventAnalysisUnit,
+    AHFESolventSetupUnit,
+    AHFESolventSimUnit,
+    AHFEVacuumAnalysisUnit,
+    AHFEVacuumSetupUnit,
+    AHFEVacuumSimUnit,
 )
 from .equil_binding_afe_method import (
     AbsoluteBindingProtocol,
@@ -30,11 +38,19 @@
     "AbsoluteSolvationProtocol",
     "AbsoluteSolvationSettings",
     "AbsoluteSolvationProtocolResult",
-    "AbsoluteVacuumUnit",
-    "AbsoluteSolventUnit",
+    "AHFESolventSetupUnit",
+    "AHFESolventSimUnit",
+    "AHFESolventAnalysisUnit",
+    "AHFEVacuumSetupUnit",
+    "AHFEVacuumSimUnit",
+    "AHFEVacuumAnalysisUnit",
     "AbsoluteBindingProtocol",
     "AbsoluteBindingSettings",
     "AbsoluteBindingProtocolResult",
-    "AbsoluteBindingComplexUnit",
-    "AbsoluteBindingSolventUnit",
+    "ABFEComplexSetupUnit",
+    "ABFEComplexSimUnit",
+    "ABFEComplexAnalysisUnit",
+    "ABFESolventSetupUnit",
+    "ABFESolventSimUnit",
+    "ABFESolventAnalysisUnit",
 ]
diff --git a/openfe/protocols/openmm_afe/abfe_units.py b/openfe/protocols/openmm_afe/abfe_units.py
@@ -1,7 +1,5 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
-# This code is part of OpenFE and is licensed under the MIT license.
-# For details, see https://github.com/OpenFreeEnergy/openfe
 """ABFE Protocol Units --- :mod:`openfe.protocols.openmm_afe.abfe_units`
 ========================================================================
 This module defines the ProtocolUnits for the
@@ -23,7 +21,7 @@
 from openmm import System
 from openmm import unit as ommunit
 from openmm.app import Topology as omm_topology
-from openmmtools.states import GlobalParameterState, ThermodynamicState
+from openmmtools.states import ThermodynamicState
 from rdkit import Chem
 
 from openfe.protocols.openmm_afe.equil_afe_settings import (
@@ -36,18 +34,16 @@
 from openfe.protocols.restraint_utils.openmm import omm_restraints
 from openfe.protocols.restraint_utils.openmm.omm_restraints import BoreschRestraint
 
-from .base_afe_units import BaseAbsoluteUnit
+from .base_afe_units import (
+    BaseAbsoluteMultiStateAnalysisUnit,
+    BaseAbsoluteMultiStateSimulationUnit,
+    BaseAbsoluteSetupUnit,
+)
 
 logger = logging.getLogger(__name__)
 
 
-class AbsoluteBindingComplexUnit(BaseAbsoluteUnit):
-    """
-    Protocol Unit for the complex phase of an absolute binding free energy
-    """
-
-    simtype = "complex"
-
+class ComplexComponentsMixin:
     def _get_components(self):
         """
         Get the relevant components for a complex transformation.
@@ -75,7 +71,9 @@ def _get_components(self):
         # Similarly we don't need to check prot_comp
         return alchem_comps, solv_comp, prot_comp, off_comps
 
-    def _handle_settings(self) -> dict[str, SettingsBaseModel]:
+
+class ComplexSettingsMixin:
+    def _get_settings(self) -> dict[str, SettingsBaseModel]:
         """
         Extract the relevant settings for a complex transformation.
 
@@ -97,7 +95,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * output_settings: MultiStateOutputSettings
             * restraint_settings: BaseRestraintSettings
         """
-        prot_settings = self._inputs["protocol"].settings
+        prot_settings = self._inputs["protocol"].settings  # type: ignore[attr-defined]
 
         settings = {}
         settings["forcefield_settings"] = prot_settings.forcefield_settings
@@ -116,6 +114,15 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
 
         return settings
 
+
+class ABFEComplexSetupUnit(ComplexComponentsMixin, ComplexSettingsMixin, BaseAbsoluteSetupUnit):
+    """
+    Setup unit for the complex phase of absolute binding free energy
+    transformations.
+    """
+
+    simtype = "complex"
+
     @staticmethod
     def _get_mda_universe(
         topology: omm_topology,
@@ -261,7 +268,6 @@ def _add_restraints(
         comp_resids: dict[Component, npt.NDArray],
         settings: dict[str, SettingsBaseModel],
     ) -> tuple[
-        GlobalParameterState,
         Quantity,
         System,
         geometry.HostGuestRestraintGeometry,
@@ -295,9 +301,6 @@ def _add_restraints(
 
         Returns
         -------
-        restraint_parameter_state : RestraintParameterState
-          A RestraintParameterState object that defines the control
-          parameter for the restraint.
         correction : openff.units.Quantity
           The standard state correction for the restraint.
         system : openmm.System
@@ -380,10 +383,7 @@ def _add_restraints(
             rest_geom,
         )
 
-        # Get the GlobalParameterState for the restraint
-        restraint_parameter_state = omm_restraints.RestraintParameterState(lambda_restraints=1.0)
         return (
-            restraint_parameter_state,
             correction,
             # Remove the thermostat, otherwise you'll get an
             # Andersen thermostat by default!
@@ -392,13 +392,28 @@ def _add_restraints(
         )
 
 
-class AbsoluteBindingSolventUnit(BaseAbsoluteUnit):
+class ABFEComplexSimUnit(
+    ComplexComponentsMixin, ComplexSettingsMixin, BaseAbsoluteMultiStateSimulationUnit
+):
     """
-    Protocol Unit for the solvent phase of an absolute binding free energy
+    Multi-state simulation (e.g. multi replica methods like Hamiltonian
+    replica exchange) unit for the complex phase of absolute binding
+    free energy transformations.
     """
 
-    simtype = "solvent"
+    simtype = "complex"
 
+
+class ABFEComplexAnalysisUnit(ComplexSettingsMixin, BaseAbsoluteMultiStateAnalysisUnit):
+    """
+    Analysis unit for multi-state simulations with the complex phase
+    of absolute binding free energy transformations.
+    """
+
+    simtype = "complex"
+
+
+class SolventComponentsMixin:
     def _get_components(self):
         """
         Get the relevant components for a solvent transformation.
@@ -426,7 +441,9 @@ def _get_components(self):
         # Similarly we don't need to check prot_comp just return None
         return alchem_comps, solv_comp, None, off_comps
 
-    def _handle_settings(self) -> dict[str, SettingsBaseModel]:
+
+class SolventSettingsMixin:
+    def _get_settings(self) -> dict[str, SettingsBaseModel]:
         """
         Extract the relevant settings for a solvent transformation.
 
@@ -447,7 +464,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
             * simulation_settings : MultiStateSimulationSettings
             * output_settings: MultiStateOutputSettings
         """
-        prot_settings = self._inputs["protocol"].settings
+        prot_settings = self._inputs["protocol"].settings  # type: ignore[attr-defined]
 
         settings = {}
         settings["forcefield_settings"] = prot_settings.forcefield_settings
@@ -464,3 +481,33 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
         settings["output_settings"] = prot_settings.solvent_output_settings
 
         return settings
+
+
+class ABFESolventSetupUnit(SolventComponentsMixin, SolventSettingsMixin, BaseAbsoluteSetupUnit):
+    """
+    Setup unit for the solvent phase of absolute binding free energy
+    transformations.
+    """
+
+    simtype = "solvent"
+
+
+class ABFESolventSimUnit(
+    SolventComponentsMixin, SolventSettingsMixin, BaseAbsoluteMultiStateSimulationUnit
+):
+    """
+    Multi-state simulation (e.g. multi replica methods like Hamiltonian
+    replica exchange) unit for the solvent phase of absolute binding
+    free energy transformations.
+    """
+
+    simtype = "solvent"
+
+
+class ABFESolventAnalysisUnit(SolventSettingsMixin, BaseAbsoluteMultiStateAnalysisUnit):
+    """
+    Analysis unit for multi-state simulations with the solvent phase
+    of absolute binding free energy transformations.
+    """
+
+    simtype = "solvent"
diff --git a/openfe/protocols/openmm_afe/afe_protocol_results.py b/openfe/protocols/openmm_afe/afe_protocol_results.py
@@ -214,8 +214,8 @@ def get_replica_state(nc, chk):
         for key in [self.bound_state, self.unbound_state]:
             for pus in self.data[key].values():  # type: ignore[attr-defined]
                 states = get_replica_state(
-                    pus[0].outputs["nc"],
-                    pus[0].outputs["last_checkpoint"],
+                    pus[0].outputs["trajectory"],
+                    pus[0].outputs["checkpoint"],
                 )
                 replica_states[key].append(states)
 
@@ -295,7 +295,9 @@ def selection_indices(self) -> dict[str, list[Optional[npt.NDArray]]]:
 
 
 class AbsoluteSolvationProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResultMixin):
-    """Dict-like container for the output of a AbsoluteSolvationProtocol"""
+    """
+    Protocol results with the output of a AbsoluteSolvationProtocol
+    """
 
     bound_state = "solvent"
     unbound_state = "vacuum"
@@ -375,7 +377,9 @@ def _get_stdev(estimates):
 
 
 class AbsoluteBindingProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResultMixin):
-    """Dict-like container for the output of a AbsoluteBindingProtocol"""
+    """
+    Protocol results with the output of a AbsoluteBindingProtocol.
+    """
 
     bound_state = "complex"
     unbound_state = "solvent"