Merge branch 'main' into KonnektorIntegration

OpenFreeEnergy · Dec 9, 2024 · 96de81b · 96de81b
2 parents 6d3e59a + 5e198a1
commit 96de81b
Show file tree

Hide file tree

Showing 26 changed files with 609 additions and 74 deletions.
diff --git a/.github/PULL_REQUEST_TEMPLATE/release_template.md b/.github/PULL_REQUEST_TEMPLATE/release_template.md
@@ -9,6 +9,7 @@ Make the PR:
 * [ ] Verify that`docs/CHANGELOG.md` looks correct.
 * [ ] Make the PR and verify that CI/CD passes. 
 * [ ] Merge the PR into `main`.
+* [ ] Make a PR into the [example notebooks repository](https://github.com/OpenFreeEnergy/ExampleNotebooks) to update the version used in `showcase/openfe_showcase.ipynb` and `.binder/environment.yml`
 
 After Merging the PR [follow this guide](https://github.com/OpenFreeEnergy/openfe/wiki/How-to-create-a-new-release)
 
diff --git a/.github/workflows/change.yaml b/.github/workflows/change.yaml
@@ -0,0 +1,41 @@
+name: Check for API breaks
+
+on:
+  pull_request_target:
+    branches:
+      - main
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+
+    - name: Check for API breaks
+      continue-on-error: true
+      id: check
+      run: |
+        pip install griffe
+        griffe check "openfe" --verbose
+
+    - name: Post Comment on Failure
+      if: steps.check.outcome == 'failure'
+      uses: actions/github-script@v7
+      with:
+        script: |
+          const prNumber = context.payload.pull_request.number;
+          github.rest.issues.createComment({
+            owner: context.repo.owner,
+            repo: context.repo.repo,
+            issue_number: prNumber,
+            body: '🚨 API breaking changes detected! 🚨'
+          });
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -30,7 +30,7 @@ defaults:
 jobs:
   tests:
     runs-on: ${{ matrix.os }}
-    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }} 🗃️${{ matrix.pydantic-version }} oechem: ${{ matrix.openeye }}"
+    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }} 🗃️${{ matrix.pydantic-version }} oechem-${{ matrix.openeye }} gufe-${{ matrix.gufe-version }}"
     strategy:
       fail-fast: false
       matrix:
@@ -41,22 +41,34 @@ jobs:
           - "3.11"
           - "3.12"
         openeye: ["no"]
+        gufe-version: ['a9b5982471eb3644ca8c4423c7001e6309a3f544']
         include:
-          - os: "macos-12"
+          - os: "macos-latest"
             python-version: "3.12"
             pydantic-version: ">1"
+            gufe-version: 'a9b5982471eb3644ca8c4423c7001e6309a3f544'
           - os: "ubuntu-latest"
             python-version: "3.11"
             pydantic-version: "<2"
+            gufe-version: 'a9b5982471eb3644ca8c4423c7001e6309a3f544'
           - os: "ubuntu-latest"
             python-version: "3.11"
             pydantic-version: ">1"
             openeye: "yes"
+            gufe-version: 'a9b5982471eb3644ca8c4423c7001e6309a3f544'
           - os: "macos-latest"
             python-version: "3.12"
             pydantic-version: ">1"
             omff-version: ">0.13"
             openeye: ["no"]
+            gufe-version: 'a9b5982471eb3644ca8c4423c7001e6309a3f544'
+
+          # temporary test against gufe main while we're pinned to a9b5
+          - os: "ubuntu-latest"
+            python-version: "3.11"
+            pydantic-version: ">1"
+            openeye: "yes"
+            gufe-version: 'main'
 
     env:
       OE_LICENSE: ${{ github.workspace }}/oe_license.txt
@@ -96,8 +108,9 @@ jobs:
           micromamba install -c openeye openeye-toolkits
           python -c "import openeye; assert openeye.oechem.OEChemIsLicensed(), 'oechem license check failed!'"
 
-      - name: "Install GUFE from main@HEAD"
-        run: python -m pip install --no-deps git+https://github.com/OpenFreeEnergy/gufe@main
+      - name: "Install GUFE at {{ matrix.gufe-version }}"
+        if: ${{ matrix.gufe-version != 'main' }}
+        run: python -m pip install --no-deps git+https://github.com/OpenFreeEnergy/gufe@${{ matrix.gufe-version }}
 
       - name: "Install"
         run: python -m pip install --no-deps -e .

diff --git a/.github/workflows/conda_cron.yaml b/.github/workflows/conda_cron.yaml
@@ -20,15 +20,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ['ubuntu-latest', 'macos-latest', 'macos-12']
+        os: ['ubuntu-latest', 'macos-latest']
         python-version:
-          - "3.9"
           - "3.10"
           - "3.11"
           - "3.12"
-        exclude:
-          - os: 'macos-latest'
-            python-version: '3.9'
 
     steps:
       - name: Setup Micromamba

diff --git a/.github/workflows/installer.yaml b/.github/workflows/installer.yaml
@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [macos-14, macos-12, ubuntu-latest]
+        os: [macos-latest, ubuntu-latest]
 
     steps:
     - name: Checkout Code

diff --git a/news/986_input_validation.rst b/news/986_input_validation.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* ``openfe quickrun`` now fails up-front when the user-supplied output file (``-o``) already exists or has a non-existent parent directory.
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>
diff --git a/news/compute_selection_fixes.rst b/news/compute_selection_fixes.rst
@@ -0,0 +1,26 @@
+**Added:**
+
+* OpenMMEngineSettings now has a `gpu_device_index` attribute
+  allowing users to pass through a list of ints to select the
+  GPU devices to run their simulations on.
+
+**Changed:**
+
+* `openfe.protocols.openmm_rfe._rfe_utils.compute` has been moved
+  to `openfe.protocols.openmm_utils.omm_compute`.
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* OpenMM CPU vacuum calculations now enforce the use of a single CPU to avoid large performance losses.
+
+**Security:**
+
+* <news item>
diff --git a/news/no_more_macos12.rst b/news/no_more_macos12.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* openfe is no longer tested against macos-12. macos support is, for now, limited to osx-arm64 (macos-14+).
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>
diff --git a/openfe/analysis/plotting.py b/openfe/analysis/plotting.py
@@ -33,7 +33,7 @@ def plot_lambda_transition_matrix(matrix: npt.NDArray) -> Axes:
     Notes
     -----
     Borrowed from `alchemlyb <https://github.com/alchemistry/alchemlyb/blob/master/src/alchemlyb/visualisation/mbar_matrix.py>`_
-    which itself borrows from `alchemical-analysis <https://github.com/MobleyLab/alchemical-analysis>`_. 
+    which itself borrows from `alchemical-analysis <https://github.com/MobleyLab/alchemical-analysis>`_.
     """
     num_states = len(matrix)
 

diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py
@@ -58,12 +58,13 @@
     MultiStateSimulationSettings, OpenMMEngineSettings,
     IntegratorSettings, LambdaSettings, MultiStateOutputSettings,
     ThermoSettings, OpenFFPartialChargeSettings,
+    OpenMMSystemGeneratorFFSettings,
 )
-from openfe.protocols.openmm_rfe._rfe_utils import compute
 from openfe.protocols.openmm_md.plain_md_methods import PlainMDProtocolUnit
 from ..openmm_utils import (
     settings_validation, system_creation,
-    multistate_analysis, charge_generation
+    multistate_analysis, charge_generation,
+    omm_compute,
 )
 from openfe.utils import without_oechem_backend
 
@@ -175,6 +176,7 @@ def _pre_equilibrate(
         settings : dict[str, SettingsBaseModel]
           A dictionary of settings objects. Expects the
           following entries:
+          * `forcefield_settings`
           * `engine_settings`
           * `thermo_settings`
           * `integrator_settings`
@@ -189,8 +191,12 @@ def _pre_equilibrate(
           Equilibrated system positions
         """
         # Prep the simulation object
-        platform = compute.get_openmm_platform(
-            settings['engine_settings'].compute_platform
+        # Restrict CPU count if running vacuum simulation
+        restrict_cpu = settings['forcefield_settings'].nonbonded_method.lower() == 'nocutoff'
+        platform = omm_compute.get_openmm_platform(
+            platform_name=settings['engine_settings'].compute_platform,
+            gpu_device_index=settings['engine_settings'].gpu_device_index,
+            restrict_cpu_count=restrict_cpu
         )
 
         integrator = openmm.LangevinMiddleIntegrator(
@@ -710,14 +716,16 @@ def _get_reporter(
 
     def _get_ctx_caches(
         self,
+        forcefield_settings: OpenMMSystemGeneratorFFSettings,
         engine_settings: OpenMMEngineSettings
     ) -> tuple[openmmtools.cache.ContextCache, openmmtools.cache.ContextCache]:
         """
         Set the context caches based on the chosen platform
 
         Parameters
         ----------
-        engine_settings : OpenMMEngineSettings,
+        forcefield_settings: OpenMMSystemGeneratorFFSettings
+        engine_settings : OpenMMEngineSettings
 
         Returns
         -------
@@ -726,8 +734,13 @@ def _get_ctx_caches(
         sampler_context_cache : openmmtools.cache.ContextCache
           The sampler state context cache.
         """
-        platform = compute.get_openmm_platform(
-            engine_settings.compute_platform,
+        # Get the compute platform
+        # Set the number of CPUs to 1 if running a vacuum simulation
+        restrict_cpu = forcefield_settings.nonbonded_method.lower() == 'nocutoff'
+        platform = omm_compute.get_openmm_platform(
+            platform_name=engine_settings.compute_platform,
+            gpu_device_index=engine_settings.gpu_device_index,
+            restrict_cpu_count=restrict_cpu
         )
 
         energy_context_cache = openmmtools.cache.ContextCache(
@@ -1026,6 +1039,7 @@ def run(self, dry=False, verbose=True,
         try:
             # 12. Get context caches
             energy_ctx_cache, sampler_ctx_cache = self._get_ctx_caches(
+                settings['forcefield_settings'],
                 settings['engine_settings']
             )
 

diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -43,10 +43,9 @@
 )
 from openff.toolkit.topology import Molecule as OFFMolecule
 
-from openfe.protocols.openmm_rfe._rfe_utils import compute
 from openfe.protocols.openmm_utils import (
     system_validation, settings_validation, system_creation,
-    charge_generation,
+    charge_generation, omm_compute
 )
 
 logger = logging.getLogger(__name__)
@@ -623,8 +622,11 @@ def run(self, *, dry=False, verbose=True,
             )
 
         # 10. Get platform
-        platform = compute.get_openmm_platform(
-            protocol_settings.engine_settings.compute_platform
+        restrict_cpu = forcefield_settings.nonbonded_method.lower() == 'nocutoff'
+        platform = omm_compute.get_openmm_platform(
+            platform_name=protocol_settings.engine_settings.compute_platform,
+            gpu_device_index=protocol_settings.engine_settings.gpu_device_index,
+            restrict_cpu_count=restrict_cpu
         )
 
         # 11. Set the integrator

diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/__init__.py b/openfe/protocols/openmm_rfe/_rfe_utils/__init__.py
@@ -1,5 +1,4 @@
 from . import (
-    compute,
     lambdaprotocol,
     multistate,
     relative,

diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
@@ -61,7 +61,7 @@
 )
 from ..openmm_utils import (
     system_validation, settings_validation, system_creation,
-    multistate_analysis, charge_generation
+    multistate_analysis, charge_generation, omm_compute,
 )
 from . import _rfe_utils
 from ...utils import without_oechem_backend, log_system_probe
@@ -933,9 +933,13 @@ def run(self, *, dry=False, verbose=True,
                 bfactors=bfactors,
             )
 
-        # 10. Get platform
-        platform = _rfe_utils.compute.get_openmm_platform(
-            protocol_settings.engine_settings.compute_platform
+        # 10. Get compute platform
+        # restrict to a single CPU if running vacuum
+        restrict_cpu = forcefield_settings.nonbonded_method.lower() == 'nocutoff'
+        platform = omm_compute.get_openmm_platform(
+            platform_name=protocol_settings.engine_settings.compute_platform,
+            gpu_device_index=protocol_settings.engine_settings.gpu_device_index,
+            restrict_cpu_count=restrict_cpu
         )
 
         # 11. Set the integrator