Update for openfe v1.9

.github/workflows/ci-dev-main.yaml

@@ -0,0 +1,68 @@
+name: "CI - main branch of deps"
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+  schedule:
+    # At 07:00 UTC on Monday and Thursday
+    - cron: "0 7 * * *"
+  workflow_dispatch:
+
+
+concurrency:
+  group: "${{ github.workflow }}-${{ github.ref }}"
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash -l {0}
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    name: "${{ matrix.os }} - ${{ matrix.python-version }}"
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version:
+          - "3.12"
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: "Setup Micromamba"
+        uses: mamba-org/setup-micromamba@v2
+        with:
+          environment-file: devtools/latest_environment.yml
+          environment-name: pontibus-dev
+          create-args: >-
+            python=${{ matrix.python-version }}
+          init-shell: bash
+
+      - name: "Install"
+        run: python -m pip install --no-deps -e .
+
+      - name: "Test imports"
+        run: |
+          # if we add more to this, consider changing to for + env vars
+          python -Ic "import pontibus; print(pontibus.__version__)"
+
+      - name: "Environment Information"
+        run: |
+          micromamba info
+          micromamba list
+          pip list
+
+      - name: "Run tests"
+        env:
+          PONTIBUS_SLOW_TESTS: TRUE
+          PONTIBUS_CPUVSLOW_TESTS: FALSE
+          PONTIBUS_GPU_TESTS: FALSE
+          DUECREDIT_ENABLE: 'yes'
+        run: |
+          pytest -n logical -v --durations=10

devtools/environment_dev.yml

@@ -0,0 +1,45 @@
+# env for testing against pre-release versions of openfe (and any other dependencies).
+name: pontibus-dev
+channels:
+  - conda-forge
+dependencies:
+  - pip
+  # keep these openff packages even though we pip install@main so that necessary deps get pulled in.
+  - openff-toolkit >0.16.0
+  - openff-interchange >=0.4.8
+  - openff-nagl-base >=0.3.3
+  # OpenFE stack deps
+  - duecredit<0.10
+  - kartograf>=1.0.0
+  - konnektor
+  - lomap2>=3.0.0
+  - numpy
+  - networkx
+  - rdkit
+  - packaging
+  - pip
+  - pydantic >=2.0
+  - pyyaml
+  - coverage
+  - cinnabar ~=0.5.0
+  - click
+  - typing-extensions
+  - openmm >=8.0.0,!=8.1.0,<8.3.0
+  - openmmtools >=0.25.0
+  - openmmforcefields
+  - openfe-analysis >=0.3.1
+  - plugcli
+  - tqdm
+  # Testing deps
+  - pytest
+  - pytest-cov
+  - pytest-xdist
+  # pip deps
+  - pip:
+    - git+https://github.com/OpenFreeEnergy/openfe@main
+    - git+https://github.com/OpenFreeEnergy/gufe@main
+    - git+https://github.com/openforcefield/openff-toolkit@main
+    - git+https://github.com/openforcefield/openff-interchange@main
+    - git+https://github.com/openforcefield/openff-nagl@main
+    - git+https://github.com/openforcefield/openff-units@main
+    - git+https://github.com/openforcefield/openff-utilities@main

devtools/gen_serialized_results.py

@@ -100,6 +100,7 @@ def generate_asfe_json_water(smc):
 def generate_asfe_json_octanol(smc):
     settings = generate_ahfe_settings()
     settings.solvation_settings.assign_solvent_charges = True
+    settings.solvation_settings.target_density = 0.715 * unit.grams / unit.mL
     protocol = ASFEProtocol(settings=settings)
     solvent = Molecule.from_smiles("CCCCCCCCO")
     solvent.assign_partial_charges(partial_charge_method="am1bcc")

docs/api/protocols/asfe_protocol.rst

@@ -18,8 +18,12 @@ Protocol API Specification
    ASFEProtocol
    ASFEProtocolResult
    ASFESettings
-   ASFESolventUnit
-   ASFEVacuumUnit
+   ASFESolventSetupUnit
+   ASFESolventSimUnit
+   ASFESolventAnalysisUnit
+   ASFEVacuumSetupUnit
+   ASFEVacuumSimUnit
+   ASFEVacuumAnalysisUnit
 
 
 Protocol Settings

docs/api/protocols/rbfe_protocol.rst

@@ -20,7 +20,7 @@ Protocol API Specification
    HybridTopProtocol
    HybridTopProtocolResult
    HybridTopProtocolSettings
-   HybridTopProtocolUnit
+   HybridTopProtocolSetupUnit
 
 
 Protocol Settings

environment.yml

@@ -8,7 +8,7 @@ dependencies:
   - openff-interchange >=0.5
   - openff-nagl-base >=0.3.3
   # OpenFE stack deps
-  - openfe ~= 1.8.0
+  - openfe >= 1.9.0
   - duecredit<0.10
   - kartograf>=1.0.0
   - konnektor

news/ofev1.9.rst

@@ -0,0 +1,25 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* The ASFE and HybridTop Protocols have been updated to work
+  with openfe v1.9. Notably this means that the Protocols have
+  now been split into three ProtocolUnits (PR #190).
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/pontibus/components/extended_solvent_component.py

@@ -33,16 +33,16 @@ def __init__(
           positive and negative monoatomic ions, defaults "Na+", "Cl-"
         neutralize : bool, optional
           if the net charge on the chemical state is neutralized by the ions in
-          this solvent component.  Default `True`
+          this solvent component.  Default `False`
         ion_concentration : openff.units.Quantity, optional
-          ionic concentration required, default 0.15 * unit.molar
+          ionic concentration required, default 0.0 * unit.molar
           this must be supplied with units, e.g. "1.5 * unit.molar"
 
         Examples
         --------
         To create a sodium chloride solution at 0.2 molar concentration::
 
-          >>> s = SolventComponent(position_ion='Na', negative_ion='Cl',
+          >>> s = SolventComponent(positive_ion='Na', negative_ion='Cl',
           ...                      ion_concentration=0.2 * unit.molar)
 
         To create a methane solvent::

src/pontibus/protocols/relative/__init__.py

@@ -4,16 +4,14 @@
 Run relative free energy calculations using OpenMM and OpenMMTools.
 """
 
-from .hybridtop_protocol import (
-    HybridTopProtocol,
-    HybridTopProtocolResult,
-)
-from .hybridtop_units import HybridTopProtocolUnit
+from .hybridtop_protocol import HybridTopProtocol
+from .hybridtop_protocol_results import HybridTopProtocolResult
+from .hybridtop_units import HybridTopProtocolSetupUnit
 from .settings import HybridTopProtocolSettings
 
 __all__ = [
     "HybridTopProtocol",
     "HybridTopProtocolSettings",
     "HybridTopProtocolResult",
-    "HybridTopProtocolUnit",
+    "HybridTopProtocolSetupUnit",
 ]

src/pontibus/protocols/relative/hybridtop_protocol.py

@@ -22,14 +22,16 @@
 from gufe.settings import ThermoSettings
 from openfe.protocols.openmm_rfe.equil_rfe_methods import (
     RelativeHybridTopologyProtocol,
-    RelativeHybridTopologyProtocolResult,
-    _validate_alchemical_components,
 )
 from openfe.protocols.openmm_rfe.equil_rfe_settings import (
     AlchemicalSettings,
     LambdaSettings,
 )
-from openfe.protocols.openmm_utils import system_validation
+from openfe.protocols.openmm_rfe.hybridtop_units import (
+    HybridTopologyMultiStateAnalysisUnit,
+    HybridTopologyMultiStateSimulationUnit,
+)
+from openfe.protocols.openmm_utils import settings_validation, system_validation
 from openfe.protocols.openmm_utils.omm_settings import (
     IntegratorSettings,
     MultiStateOutputSettings,
@@ -39,22 +41,15 @@
 )
 from openff.units import unit
 
-from pontibus.protocols.relative.hybridtop_units import HybridTopProtocolUnit
+from pontibus.protocols.relative.hybridtop_protocol_results import HybridTopProtocolResult
+from pontibus.protocols.relative.hybridtop_units import HybridTopProtocolSetupUnit
 from pontibus.protocols.relative.settings import HybridTopProtocolSettings
 from pontibus.utils.settings import (
     InterchangeFFSettings,
     InterchangeOpenMMSolvationSettings,
 )
 
 
-class HybridTopProtocolResult(RelativeHybridTopologyProtocolResult):
-    """
-    Results class for the HybridTopologyProtocol class.
-    Inherits from
-    :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocolResult`.
-    """
-
-
 class HybridTopProtocol(RelativeHybridTopologyProtocol):
     """
     Relative Free Energy calculations using OpenMM and OpenMMTools.
@@ -67,7 +62,7 @@ class HybridTopProtocol(RelativeHybridTopologyProtocol):
     :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocol`
     :class:`pontibus.protocols.relative.HybridTopSettings`
     :class:`pontibus.protocols.relative.HybridTopResult`
-    :class:`pontibus.protocols.relative.HybridTopProtocolUnit`
+    :class:`pontibus.protocols.relative.HybridTopProtocolSetupUnit`
     """
 
     result_cls = HybridTopProtocolResult
@@ -107,20 +102,26 @@ def _default_settings(cls):
             output_settings=MultiStateOutputSettings(),
         )
 
-    def _create(
+    def _validate(
         self,
         stateA: ChemicalSystem,
         stateB: ChemicalSystem,
         mapping: ComponentMapping | list[ComponentMapping] | None,
-        extends: ProtocolDAGResult | None,
-    ) -> list[ProtocolUnit]:
-        if extends:  # pragma: no cover
-            raise NotImplementedError("Can't extend simulations yet")
+        extends: ProtocolDAGResult | None = None,
+    ) -> None:
+        # Check we're not trying to extend
+        if extends:
+            raise ValueError("Can't extend simulations yet")
+
+        # Validate the end states
+        self._validate_endstates(stateA, stateB)
 
-        # Get alchemical components & validate them + mapping
+        # Validate the mapping
         alchem_comps = system_validation.get_alchemical_components(stateA, stateB)
-        _validate_alchemical_components(alchem_comps, mapping)
-        ligandmapping = mapping[0] if isinstance(mapping, list) else mapping
+        self._validate_mapping(mapping, alchem_comps)
+
+        # Validate the small molecule components
+        self._validate_smcs(stateA, stateB)
 
         # Validate solvent component
         nonbond = self.settings.forcefield_settings.nonbonded_method
@@ -129,22 +130,67 @@ def _create(
         # Validate protein component
         system_validation.validate_protein(stateA)
 
-        # actually create and return Units
+        # Validate integrator things
+        settings_validation.validate_timestep(
+            self.settings.forcefield_settings.hydrogen_mass,
+            self.settings.integrator_settings.timestep,
+        )
+
+    def _create(
+        self,
+        stateA: ChemicalSystem,
+        stateB: ChemicalSystem,
+        mapping: ComponentMapping | list[ComponentMapping] | None,
+        extends: ProtocolDAGResult | None,
+    ) -> list[ProtocolUnit]:
+        # Run validation
+        self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends)
+
+        # Get alchemical components & mapping
+        alchem_comps = system_validation.get_alchemical_components(stateA, stateB)
+        ligandmapping = mapping[0] if isinstance(mapping, list) else mapping
+
         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
-        units = [
-            HybridTopProtocolUnit(
+
+        setup_units = []
+        simulation_units = []
+        analysis_units = []
+
+        for i in range(n_repeats):
+            repeat_id = int(uuid.uuid4())
+
+            setup = HybridTopProtocolSetupUnit(
                 protocol=self,
                 stateA=stateA,
                 stateB=stateB,
                 ligandmapping=ligandmapping,
+                alchemical_components=alchem_comps,
                 generation=0,
-                repeat_id=int(uuid.uuid4()),
-                name=f"{Anames} to {Bnames} repeat {i} generation 0",
+                repeat_id=repeat_id,
+                name=f"HybridTop Setup: {Anames} to {Bnames} repeat {i} generation 0",
             )
-            for i in range(n_repeats)
-        ]
 
-        return units
+            simulation = HybridTopologyMultiStateSimulationUnit(
+                protocol=self,
+                setup_results=setup,
+                generation=0,
+                repeat_id=repeat_id,
+                name=f"HybridTop Simulation: {Anames} to {Bnames} repeat {i} generation 0",
+            )
+
+            analysis = HybridTopologyMultiStateAnalysisUnit(
+                protocol=self,
+                setup_results=setup,
+                simulation_results=simulation,
+                generation=0,
+                repeat_id=repeat_id,
+                name=f"HybridTop Analysis: {Anames} to {Bnames} repeat {i} generation 0",
+            )
+
+            setup_units.append(setup)
+            simulation_units.append(simulation)
+            analysis_units.append(analysis)
+
+        return [*setup_units, *simulation_units, *analysis_units]