From 3243539e3a76ac857180d006d8682f78a0bcc191 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Fri, 23 May 2025 15:53:52 +0100
Subject: [PATCH 01/10] Address typo from previous PR

---
 openfe/protocols/openmm_afe/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py
index df5bcc06c..d08ad3a71 100644
--- a/openfe/protocols/openmm_afe/base.py
+++ b/openfe/protocols/openmm_afe/base.py
@@ -168,7 +168,7 @@ def _pre_equilibrate(
         Parameters
         ----------
         system : openmm.System
-          An OpenMM System to equilibrate.
+          The OpenMM System to equilibrate.
         topology : openmm.app.Topology
           OpenMM Topology of the System.
         positions : openmm.unit.Quantity
@@ -502,7 +502,7 @@ def _get_omm_objects(
         topology : app.Topology
           OpenMM Topology object describing the parameterized system.
         system : openmm.System
-          An non-alchemical OpenMM System of the simulated system.
+          A non-alchemical OpenMM System of the simulated system.
         positions : openmm.unit.Quantity
           Positions of the system.
         comp_resids : dict[Component, npt.NDArray]

From 32196b2772248a59070a862c84d52079f0c1cee9 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Fri, 23 May 2025 16:05:46 +0100
Subject: [PATCH 02/10] move tests around

---
 openfe/tests/protocols/openmm/__init__.py                         | 0
 openfe/tests/protocols/openmm/absolute_hydration/__init__.py      | 0
 .../absolute_hydration/test_protocol.py}                          | 0
 .../absolute_hydration/test_simulation_slow.py}                   | 0
 .../absolute_hydration/test_tokenization.py}                      | 0
 openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py           | 0
 .../hybridtop_rfe/test_protocol.py}                               | 0
 .../hybridtop_rfe/test_simulation_slow.py}                        | 0
 .../hybridtop_rfe/test_tokenization.py}                           | 0
 openfe/tests/protocols/openmm/plain_md/__init__.py                | 0
 .../plain_md/test_protocol.py}                                    | 0
 openfe/tests/protocols/{ => openmm}/test_openmm_settings.py       | 0
 openfe/tests/protocols/{ => openmm}/test_openmmutils.py           | 0
 13 files changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 openfe/tests/protocols/openmm/__init__.py
 create mode 100644 openfe/tests/protocols/openmm/absolute_hydration/__init__.py
 rename openfe/tests/protocols/{test_openmm_afe_solvation_protocol.py => openmm/absolute_hydration/test_protocol.py} (100%)
 rename openfe/tests/protocols/{test_openmm_afe_slow.py => openmm/absolute_hydration/test_simulation_slow.py} (100%)
 rename openfe/tests/protocols/{test_solvation_afe_tokenization.py => openmm/absolute_hydration/test_tokenization.py} (100%)
 create mode 100644 openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py
 rename openfe/tests/protocols/{test_openmm_equil_rfe_protocols.py => openmm/hybridtop_rfe/test_protocol.py} (100%)
 rename openfe/tests/protocols/{test_openmm_rfe_slow.py => openmm/hybridtop_rfe/test_simulation_slow.py} (100%)
 rename openfe/tests/protocols/{test_rfe_tokenization.py => openmm/hybridtop_rfe/test_tokenization.py} (100%)
 create mode 100644 openfe/tests/protocols/openmm/plain_md/__init__.py
 rename openfe/tests/protocols/{test_openmm_plain_md_protocols.py => openmm/plain_md/test_protocol.py} (100%)
 rename openfe/tests/protocols/{ => openmm}/test_openmm_settings.py (100%)
 rename openfe/tests/protocols/{ => openmm}/test_openmmutils.py (100%)

diff --git a/openfe/tests/protocols/openmm/__init__.py b/openfe/tests/protocols/openmm/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/__init__.py b/openfe/tests/protocols/openmm/absolute_hydration/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/openfe/tests/protocols/test_openmm_afe_solvation_protocol.py b/openfe/tests/protocols/openmm/absolute_hydration/test_protocol.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_afe_solvation_protocol.py
rename to openfe/tests/protocols/openmm/absolute_hydration/test_protocol.py
diff --git a/openfe/tests/protocols/test_openmm_afe_slow.py b/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_afe_slow.py
rename to openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
diff --git a/openfe/tests/protocols/test_solvation_afe_tokenization.py b/openfe/tests/protocols/openmm/absolute_hydration/test_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/test_solvation_afe_tokenization.py
rename to openfe/tests/protocols/openmm/absolute_hydration/test_tokenization.py
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py b/openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/openmm/hybridtop_rfe/test_protocol.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
rename to openfe/tests/protocols/openmm/hybridtop_rfe/test_protocol.py
diff --git a/openfe/tests/protocols/test_openmm_rfe_slow.py b/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_rfe_slow.py
rename to openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
diff --git a/openfe/tests/protocols/test_rfe_tokenization.py b/openfe/tests/protocols/openmm/hybridtop_rfe/test_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/test_rfe_tokenization.py
rename to openfe/tests/protocols/openmm/hybridtop_rfe/test_tokenization.py
diff --git a/openfe/tests/protocols/openmm/plain_md/__init__.py b/openfe/tests/protocols/openmm/plain_md/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/openmm/plain_md/test_protocol.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_plain_md_protocols.py
rename to openfe/tests/protocols/openmm/plain_md/test_protocol.py
diff --git a/openfe/tests/protocols/test_openmm_settings.py b/openfe/tests/protocols/openmm/test_openmm_settings.py
similarity index 100%
rename from openfe/tests/protocols/test_openmm_settings.py
rename to openfe/tests/protocols/openmm/test_openmm_settings.py
diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/openmm/test_openmmutils.py
similarity index 100%
rename from openfe/tests/protocols/test_openmmutils.py
rename to openfe/tests/protocols/openmm/test_openmmutils.py

From 1f3d447e94a5beb2d2befeaa0937c01267c4c667 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Fri, 23 May 2025 17:15:12 +0100
Subject: [PATCH 03/10] Add tokenization tests for MD protocol

---
 .../openmm/plain_md/test_tokenization.py      | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)
 create mode 100644 openfe/tests/protocols/openmm/plain_md/test_tokenization.py

diff --git a/openfe/tests/protocols/openmm/plain_md/test_tokenization.py b/openfe/tests/protocols/openmm/plain_md/test_tokenization.py
new file mode 100644
index 000000000..3f7a82b90
--- /dev/null
+++ b/openfe/tests/protocols/openmm/plain_md/test_tokenization.py
@@ -0,0 +1,83 @@
+# This ccode is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
+import json
+import pytest
+import gufe
+from gufe.tests.test_tokenization import GufeTokenizableTestsMixin
+import openfe
+from openfe.protocols import openmm_md
+
+
+@pytest.fixture
+def protocol():
+    return openmm_md.PlainMDProtocol(
+        openmm_md.PlainMDProtocol.default_settings()
+    )
+
+
+@pytest.fixture
+def protocol_unit(protocol, benzene_system):
+    pus = protocol.create(
+        stateA=benzene_system,
+        stateB=benzene_system,
+        mapping=None,
+    )
+    return list(pus.protocol_units)[0]
+
+
+@pytest.fixture
+def protocol_result(md_json):
+    d = json.loads(md_json, cls=gufe.tokenization.JSON_HANDLER.decoder)
+    pr = gufe.ProtocolResult.from_dict(d['protocol_result'])
+    return pr
+
+
+class TestPlainMDProtocol(GufeTokenizableTestsMixin):
+    cls = openmm_md.PlainMDProtocol
+    key = None
+    repr = "PlainMDProtocol-"
+
+    @pytest.fixture()
+    def instance(self, protocol):
+        return protocol
+    
+    def test_repr(self, instance):
+        """
+        Overwrites the base `test_repr` call to do a bit more.
+        """
+        assert isinstance(repr(instance), str)
+        assert self.repr in repr(instance)
+
+
+class TestPlainMDProtocolUnit(GufeTokenizableTestsMixin):
+    cls = openmm_md.PlainMDProtocolUnit
+    repr = "PlainMDProtocolUnit("
+    key = None
+
+    @pytest.fixture
+    def instance(self, protocol_unit):
+        return protocol_unit
+    
+    def test_repr(self, instance):
+        """
+        Overwrites the base `test_repr` call to do a bit more.
+        """
+        assert isinstance(repr(instance), str)
+        assert self.repr in repr(instance)
+
+
+class TestPlainMDProtocolResult(GufeTokenizableTestsMixin):
+    cls = openmm_md.PlainMDProtocolResult
+    key = None
+    repr = "PlainMDProtocolResult-"
+
+    @pytest.fixture()
+    def instance(self, protocol_result):
+        return protocol_result
+    
+    def test_repr(self, instance):
+        """
+        Overwrites the base `test_repr` call to do a bit more.
+        """
+        assert isinstance(repr(instance), str)
+        assert self.repr in repr(instance)
\ No newline at end of file

From d85cf97dc078edf8283cb6f67f2dfc577c18ef5b Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Fri, 23 May 2025 17:35:21 +0100
Subject: [PATCH 04/10] update import for HAS_INTERNET

---
 openfe/tests/protocols/openmm/test_openmmutils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openfe/tests/protocols/openmm/test_openmmutils.py b/openfe/tests/protocols/openmm/test_openmmutils.py
index 4a1a7dc4f..770b13e6b 100644
--- a/openfe/tests/protocols/openmm/test_openmmutils.py
+++ b/openfe/tests/protocols/openmm/test_openmmutils.py
@@ -32,7 +32,7 @@
 from openmmtools import multistate
 from pymbar.utils import ParameterError
 
-from ..conftest import HAS_INTERNET
+from openfe.tests.conftest import HAS_INTERNET
 
 
 @pytest.mark.parametrize('padding, number_solv, box_vectors, box_size', [

From 613a0e804cb9bc9d9ae7e6f2340e8938f36d0874 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Sat, 24 May 2025 21:45:56 +0100
Subject: [PATCH 05/10] Add gas phase slow test for MD Protocol, and fix PDB
 writing

---
 .../protocols/openmm_md/plain_md_methods.py   | 16 ++--
 openfe/tests/protocols/conftest.py            |  9 +++
 .../test_simulation_slow.py                   | 33 ++------
 .../hybridtop_rfe/test_simulation_slow.py     | 35 ++-------
 .../openmm/plain_md/test_simulation_slow.py   | 77 +++++++++++++++++++
 .../openmm/plain_md/test_tokenization.py      |  1 -
 6 files changed, 109 insertions(+), 62 deletions(-)
 create mode 100644 openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py

diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
index a5203263d..f14526757 100644
--- a/openfe/protocols/openmm_md/plain_md_methods.py
+++ b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -19,18 +19,17 @@
 from openff.units.openmm import from_openmm, to_openmm
 import openmm.unit as omm_unit
 from typing import Optional
-from openmm import app
 import pathlib
 from typing import Any, Iterable
 import uuid
 import time
-import numpy as np
 import mdtraj
 from mdtraj.reporters import XTCReporter
 from openfe.utils import without_oechem_backend, log_system_probe
 from gufe import (
-    settings, ChemicalSystem, SmallMoleculeComponent,
-    ProteinComponent, SolventComponent
+    settings,
+    ChemicalSystem,
+    SmallMoleculeComponent,
 )
 from openfe.protocols.openmm_utils.omm_settings import (
     BasePartialChargeSettings,
@@ -88,7 +87,7 @@ def get_traj_filename(self) -> list[pathlib.Path]:
         traj : list[pathlib.Path]
           list of paths (pathlib.Path) to the simulation trajectory
         """
-        traj = [pus[0].outputs['nc'] for pus in self.data.values()]
+        traj = [pus[0].outputs['trajectory'] for pus in self.data.values()]
 
         return traj
 
@@ -681,10 +680,13 @@ def run(self, *, dry=False, verbose=True,
             output = {
                 'system_pdb': shared_basepath / output_settings.preminimized_structure,
                 'minimized_pdb': shared_basepath / output_settings.minimized_structure,
-                'nc': shared_basepath / output_settings.production_trajectory_filename,
+                'trajectory': shared_basepath / output_settings.production_trajectory_filename,
                 'last_checkpoint': shared_basepath / output_settings.checkpoint_storage_filename,
             }
-            if output_settings.equil_nvt_structure:
+            if (
+                output_settings.equil_nvt_structure
+                and sim_settings.equilibration_length_nvt is not None
+            ):
                 output['nvt_equil_pdb'] = shared_basepath / output_settings.equil_nvt_structure
             if output_settings.equil_npt_structure:
                 output['npt_equil_pdb'] = shared_basepath / output_settings.equil_npt_structure
diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py
index 67cc21d1c..d705af40f 100644
--- a/openfe/tests/protocols/conftest.py
+++ b/openfe/tests/protocols/conftest.py
@@ -5,10 +5,19 @@
 from importlib import resources
 from rdkit import Chem
 from rdkit.Geometry import Point3D
+from openmm import Platform
 import openfe
 from openff.units import unit
 
 
+@pytest.fixture
+def available_platforms() -> set[str]:
+    return {
+        Platform.getPlatform(i).getName()
+        for i in range(Platform.getNumPlatforms())
+    }
+
+
 @pytest.fixture
 def benzene_vacuum_system(benzene_modifications):
     return openfe.ChemicalSystem(
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py b/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
index 37b5b83aa..bde8ec0af 100644
--- a/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
+++ b/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
@@ -4,42 +4,21 @@
 from gufe.protocols import execute_DAG
 import pytest
 from openff.units import unit
-from openmm import Platform
-import os
 import pathlib
 
 import openfe
 from openfe.protocols import openmm_afe
 
 
-@pytest.fixture
-def available_platforms() -> set[str]:
-    return {Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())}
-
-
-@pytest.fixture
-def set_openmm_threads_1():
-    # for vacuum sims, we want to limit threads to one
-    # this fixture sets OPENMM_CPU_THREADS='1' for a single test, then reverts to previously held value
-    previous: str | None = os.environ.get('OPENMM_CPU_THREADS')
-
-    try:
-        os.environ['OPENMM_CPU_THREADS'] = '1'
-        yield
-    finally:
-        if previous is None:
-            del os.environ['OPENMM_CPU_THREADS']
-        else:
-            os.environ['OPENMM_CPU_THREADS'] = previous
-
-
 @pytest.mark.integration  # takes too long to be a slow test ~ 4 mins locally
 @pytest.mark.flaky(reruns=3)  # pytest-rerunfailures; we can get bad minimisation
 @pytest.mark.parametrize('platform', ['CPU', 'CUDA'])
-def test_openmm_run_engine(platform,
-                           available_platforms,
-                           benzene_modifications,
-                           set_openmm_threads_1, tmpdir):
+def test_openmm_run_engine(
+    platform,
+    available_platforms,
+    benzene_modifications,
+    tmpdir
+):
     if platform not in available_platforms:
         pytest.skip(f"OpenMM Platform: {platform} not available")
 
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py b/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
index f0fe5a0a7..ffad3f501 100644
--- a/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
+++ b/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
@@ -5,41 +5,22 @@
 from gufe.protocols import execute_DAG
 import pytest
 from openff.units import unit
-from openmm import Platform
-import os
 import pathlib
 
 import openfe
 from openfe.protocols import openmm_rfe
 
 
-@pytest.fixture
-def available_platforms() -> set[str]:
-    return {Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())}
-
-
-@pytest.fixture
-def set_openmm_threads_1():
-    # for vacuum sims, we want to limit threads to one
-    # this fixture sets OPENMM_CPU_THREADS='1' for a single test, then reverts to previously held value
-    previous: str | None = os.environ.get('OPENMM_CPU_THREADS')
-
-    try:
-        os.environ['OPENMM_CPU_THREADS'] = '1'
-        yield
-    finally:
-        if previous is None:
-            del os.environ['OPENMM_CPU_THREADS']
-        else:
-            os.environ['OPENMM_CPU_THREADS'] = previous
-
-
 @pytest.mark.slow
-@pytest.mark.flaky(reruns=3)  # pytest-rerunfailures; we can get bad minimisation
+@pytest.mark.flaky(reruns=3)  # pytest-rerunfailures; we can get bad minimization
 @pytest.mark.parametrize('platform', ['CPU', 'CUDA'])
-def test_openmm_run_engine(benzene_vacuum_system, platform,
-                           available_platforms, benzene_modifications,
-                           set_openmm_threads_1, tmpdir):
+def test_openmm_run_engine(
+    benzene_vacuum_system,
+    platform,
+    available_platforms,
+    benzene_modifications,
+    tmpdir
+):
     if platform not in available_platforms:
         pytest.skip(f"OpenMM Platform: {platform} not available")
     # this test actually runs MD
diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
new file mode 100644
index 000000000..e09dc540e
--- /dev/null
+++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
@@ -0,0 +1,77 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
+import pathlib
+import pytest
+from openff.units import unit
+from gufe.protocols import execute_DAG
+from openfe.protocols import openmm_md
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize('platform', ['CPU', 'CUDA'])
+def test_vacuum_sim(
+    benzene_vacuum_system,
+    platform,
+    available_platforms,
+    tmpdir
+):
+    if platform not in available_platforms:
+        pytest.skip(f"OpenMM Platform: {platform} is not available")
+
+    # Run a vacuum MD simulation and check what files we get.
+    settings = openmm_md.PlainMDProtocol.default_settings()
+    settings.simulation_settings.equilibration_length_nvt = None
+    settings.simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.simulation_settings.production_length = 20 * unit.picosecond
+    settings.output_settings.checkpoint_interval = 10 * unit.picosecond
+    settings.forcefield_settings.nonbonded_method = "nocutoff"
+    settings.engine_settings.compute_platform = platform
+
+    prot = openmm_md.PlainMDProtocol(settings)
+
+    dag = prot.create(
+        stateA=benzene_vacuum_system,
+        stateB=benzene_vacuum_system,
+        mapping=None,
+    )
+
+    workdir = pathlib.Path(str(tmpdir))
+
+    r = execute_DAG(
+        dag,
+        shared_basedir=workdir,
+        scratch_basedir=workdir,
+        keep_shared=True
+    )
+
+    assert r.ok()
+
+    assert len(r.protocol_unit_results) == 1
+
+    pur = r.protocol_unit_results[0]
+    unit_shared = tmpdir / f"shared_{pur.source_key}_attempt_0"
+    assert unit_shared.exists()
+    assert pathlib.Path(unit_shared).is_dir()
+
+    # check the files
+    files = [
+        "checkpoint.chk",
+        "equil_npt.pdb",
+        "minimized.pdb",
+        "simulation.xtc",
+        "simulation.log",
+        "system.pdb"
+    ]
+    for file in files:
+        assert (unit_shared / file).exists()
+
+    # NVT PDB should not exist
+    assert not (unit_shared / "equil_nvt.pdb").exists()
+
+    # check that the output file paths are correct
+    assert pur.outputs['system_pdb'] == unit_shared / "system.pdb"
+    assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb"
+    assert pur.outputs['trajectory'] == unit_shared / "simulation.xtc"
+    assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
+    assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
+    assert "nvt_equil_pdb" not in pur.outputs.keys()
diff --git a/openfe/tests/protocols/openmm/plain_md/test_tokenization.py b/openfe/tests/protocols/openmm/plain_md/test_tokenization.py
index 3f7a82b90..1a7cf95ae 100644
--- a/openfe/tests/protocols/openmm/plain_md/test_tokenization.py
+++ b/openfe/tests/protocols/openmm/plain_md/test_tokenization.py
@@ -4,7 +4,6 @@
 import pytest
 import gufe
 from gufe.tests.test_tokenization import GufeTokenizableTestsMixin
-import openfe
 from openfe.protocols import openmm_md
 
 

From edf82800ee1faeb559bc9a8d47f2ef2713ed68f2 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Sat, 24 May 2025 22:17:15 +0100
Subject: [PATCH 06/10] Add complex sim long test

---
 .../openmm/plain_md/test_simulation_slow.py   | 67 +++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
index e09dc540e..c47839e12 100644
--- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
+++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
@@ -75,3 +75,70 @@ def test_vacuum_sim(
     assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
     assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
     assert "nvt_equil_pdb" not in pur.outputs.keys()
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize('platform', ['CUDA'])
+def test_complex_solvent_sim_gpu(
+    benzene_complex_system,
+    platform,
+    available_platforms,
+    tmpdir,
+):
+    if platform not in available_platforms:
+        pytest.skip(f"OpenMM Platform: {platform} is not available")
+
+    # Run an MD simulation and check what files we get.
+    settings = openmm_md.PlainMDProtocol.default_settings()
+    settings.simulation_settings.equilibration_length_nvt = 50 * unit.picosecond
+    settings.simulation_settings.equilibration_length = 50 * unit.picosecond
+    settings.simulation_settings.production_length = 100 * unit.picosecond
+    settings.output_settings.checkpoint_interval = 10 * unit.picosecond
+    settings.engine_settings.compute_platform = platform
+
+    prot = openmm_md.PlainMDProtocol(settings)
+
+    dag = prot.create(
+        stateA=benzene_complex_system,
+        stateB=benzene_complex_system,
+        mapping=None,
+    )
+
+    workdir = pathlib.Path(str(tmpdir))
+
+    r = execute_DAG(
+        dag,
+        shared_basedir=workdir,
+        scratch_basedir=workdir,
+        keep_shared=True
+    )
+
+    assert r.ok()
+
+    assert len(r.protocol_unit_results) == 1
+
+    pur = r.protocol_unit_results[0]
+    unit_shared = tmpdir / f"shared_{pur.source_key}_attempt_0"
+    assert unit_shared.exists()
+    assert pathlib.Path(unit_shared).is_dir()
+
+    # check the files
+    files = [
+        "checkpoint.chk",
+        "equil_nvt.pdb",
+        "equil_npt.pdb",
+        "minimized.pdb",
+        "simulation.xtc",
+        "simulation.log",
+        "system.pdb"
+    ]
+    for file in files:
+        assert (unit_shared / file).exists()
+
+    # check that the output file paths are correct
+    assert pur.outputs['system_pdb'] == unit_shared / "system.pdb"
+    assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb"
+    assert pur.outputs['trajectory'] == unit_shared / "simulation.xtc"
+    assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
+    assert pur.outputs['nvt_equil_pdb'] == unit_shared / "equil_nvt.pdb"
+    assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
\ No newline at end of file

From 18fc05630adcec334cede823bdc01cbaaf52e683 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Sat, 24 May 2025 22:33:19 +0100
Subject: [PATCH 07/10] Make file returns optional rather than not including
 fields

---
 openfe/protocols/openmm_md/plain_md_methods.py           | 9 +++++++++
 .../protocols/openmm/plain_md/test_simulation_slow.py    | 8 ++++----
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
index f14526757..39644de08 100644
--- a/openfe/protocols/openmm_md/plain_md_methods.py
+++ b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -683,11 +683,20 @@ def run(self, *, dry=False, verbose=True,
                 'trajectory': shared_basepath / output_settings.production_trajectory_filename,
                 'last_checkpoint': shared_basepath / output_settings.checkpoint_storage_filename,
             }
+            # The checkpoint file can not exist if frequency > sim length
+            if not output['last_checkpoint'].exists():
+                output['last_checkpoint'] = None
+
+            # The NVT PDB can be ommitted if we don't run the simulation
+            # Note: we could also just check the file exist
             if (
                 output_settings.equil_nvt_structure
                 and sim_settings.equilibration_length_nvt is not None
             ):
                 output['nvt_equil_pdb'] = shared_basepath / output_settings.equil_nvt_structure
+            else:
+                output['nvt_equil_pdb'] = None
+
             if output_settings.equil_npt_structure:
                 output['npt_equil_pdb'] = shared_basepath / output_settings.equil_npt_structure
 
diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
index c47839e12..f278016b7 100644
--- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
+++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
@@ -23,7 +23,7 @@ def test_vacuum_sim(
     settings.simulation_settings.equilibration_length_nvt = None
     settings.simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.simulation_settings.production_length = 20 * unit.picosecond
-    settings.output_settings.checkpoint_interval = 10 * unit.picosecond
+    settings.output_settings.checkpoint_interval = 40 * unit.picosecond
     settings.forcefield_settings.nonbonded_method = "nocutoff"
     settings.engine_settings.compute_platform = platform
 
@@ -55,7 +55,6 @@ def test_vacuum_sim(
 
     # check the files
     files = [
-        "checkpoint.chk",
         "equil_npt.pdb",
         "minimized.pdb",
         "simulation.xtc",
@@ -67,14 +66,15 @@ def test_vacuum_sim(
 
     # NVT PDB should not exist
     assert not (unit_shared / "equil_nvt.pdb").exists()
+    assert not (unit_shared / "checkpoint.chk").exists()
 
     # check that the output file paths are correct
     assert pur.outputs['system_pdb'] == unit_shared / "system.pdb"
     assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb"
     assert pur.outputs['trajectory'] == unit_shared / "simulation.xtc"
-    assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
+    assert pur.outputs['last_checkpoint'] is None
     assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
-    assert "nvt_equil_pdb" not in pur.outputs.keys()
+    assert pur.outputs['nvt_equil_pdb'] is None
 
 
 @pytest.mark.integration

From 7eaa6daaa9e29ef37f0490feaffed35b073ae307 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Sat, 24 May 2025 22:49:37 +0100
Subject: [PATCH 08/10] revert trajectory file key rename, we'll do this later

---
 openfe/protocols/openmm_md/plain_md_methods.py |  4 ++--
 .../protocols/openmm/plain_md/test_protocol.py | 18 ++++++++++++++----
 .../openmm/plain_md/test_simulation_slow.py    |  6 +++---
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
index 39644de08..51613b761 100644
--- a/openfe/protocols/openmm_md/plain_md_methods.py
+++ b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -87,7 +87,7 @@ def get_traj_filename(self) -> list[pathlib.Path]:
         traj : list[pathlib.Path]
           list of paths (pathlib.Path) to the simulation trajectory
         """
-        traj = [pus[0].outputs['trajectory'] for pus in self.data.values()]
+        traj = [pus[0].outputs['nc'] for pus in self.data.values()]
 
         return traj
 
@@ -680,7 +680,7 @@ def run(self, *, dry=False, verbose=True,
             output = {
                 'system_pdb': shared_basepath / output_settings.preminimized_structure,
                 'minimized_pdb': shared_basepath / output_settings.minimized_structure,
-                'trajectory': shared_basepath / output_settings.production_trajectory_filename,
+                'nc': shared_basepath / output_settings.production_trajectory_filename,
                 'last_checkpoint': shared_basepath / output_settings.checkpoint_storage_filename,
             }
             # The checkpoint file can not exist if frequency > sim length
diff --git a/openfe/tests/protocols/openmm/plain_md/test_protocol.py b/openfe/tests/protocols/openmm/plain_md/test_protocol.py
index 44a7018b1..edd731b82 100644
--- a/openfe/tests/protocols/openmm/plain_md/test_protocol.py
+++ b/openfe/tests/protocols/openmm/plain_md/test_protocol.py
@@ -434,8 +434,13 @@ def solvent_protocol_dag(benzene_system):
 def test_unit_tagging(solvent_protocol_dag, tmpdir):
     # test that executing the Units includes correct generation and repeat info
     dag_units = solvent_protocol_dag.protocol_units
-    with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run',
-                    return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}):
+    with mock.patch(
+        'openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run',
+        return_value={
+            'nc': 'simulation.xtc',
+            'last_checkpoint': 'checkpoint.chk'
+        }
+    ):
         results = []
         for u in dag_units:
             ret = u.execute(context=gufe.Context(tmpdir, tmpdir))
@@ -452,8 +457,13 @@ def test_unit_tagging(solvent_protocol_dag, tmpdir):
 
 def test_gather(solvent_protocol_dag, tmpdir):
     # check .gather behaves as expected
-    with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run',
-                    return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}):
+    with mock.patch(
+        'openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run',
+        return_value={
+            'nc': 'simulation.xtc',
+            'last_checkpoint': 'checkpoint.chk'
+        }
+    ):
         dagres = gufe.protocols.execute_DAG(solvent_protocol_dag,
                                             shared_basedir=tmpdir,
                                             scratch_basedir=tmpdir,
diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
index f278016b7..ecb7bc799 100644
--- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
+++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
@@ -71,7 +71,7 @@ def test_vacuum_sim(
     # check that the output file paths are correct
     assert pur.outputs['system_pdb'] == unit_shared / "system.pdb"
     assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb"
-    assert pur.outputs['trajectory'] == unit_shared / "simulation.xtc"
+    assert pur.outputs['nc'] == unit_shared / "simulation.xtc"
     assert pur.outputs['last_checkpoint'] is None
     assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
     assert pur.outputs['nvt_equil_pdb'] is None
@@ -138,7 +138,7 @@ def test_complex_solvent_sim_gpu(
     # check that the output file paths are correct
     assert pur.outputs['system_pdb'] == unit_shared / "system.pdb"
     assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb"
-    assert pur.outputs['trajectory'] == unit_shared / "simulation.xtc"
+    assert pur.outputs['nc'] == unit_shared / "simulation.xtc"
     assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
     assert pur.outputs['nvt_equil_pdb'] == unit_shared / "equil_nvt.pdb"
-    assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
\ No newline at end of file
+    assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"

From 67e2fc707cbce392c3c8272927a59562ed348a40 Mon Sep 17 00:00:00 2001
From: Alyssa Travitz <alyssa.travitz@omsf.io>
Date: Wed, 28 May 2025 15:25:45 -0700
Subject: [PATCH 09/10] reorganizing test directory structure

---
 openfe/tests/protocols/openmm/plain_md/__init__.py                | 0
 openfe/tests/protocols/{openmm => openmm_afe}/__init__.py         | 0
 .../test_protocol.py => openmm_afe/test_ahfe_protocol.py}         | 0
 .../test_simulation_slow.py => openmm_afe/test_ahfe_slow.py}      | 0
 .../test_tokenization.py => openmm_afe/test_ahfe_tokenization.py} | 0
 .../{openmm/absolute_hydration => openmm_md}/__init__.py          | 0
 .../test_protocol.py => openmm_md/test_plain_md_protocol.py}      | 0
 .../test_simulation_slow.py => openmm_md/test_plain_md_slow.py}   | 0
 .../test_plain_md_tokenization.py}                                | 0
 .../protocols/{openmm/hybridtop_rfe => openmm_rfe}/__init__.py    | 0
 .../test_protocol.py => openmm_rfe/test_hybrid_top_protocol.py}   | 0
 .../test_hybrid_top_slow.py}                                      | 0
 .../test_hybrid_top_tokenization.py}                              | 0
 openfe/tests/protocols/{openmm => }/test_openmm_settings.py       | 0
 openfe/tests/protocols/{openmm => }/test_openmmutils.py           | 0
 15 files changed, 0 insertions(+), 0 deletions(-)
 delete mode 100644 openfe/tests/protocols/openmm/plain_md/__init__.py
 rename openfe/tests/protocols/{openmm => openmm_afe}/__init__.py (100%)
 rename openfe/tests/protocols/{openmm/absolute_hydration/test_protocol.py => openmm_afe/test_ahfe_protocol.py} (100%)
 rename openfe/tests/protocols/{openmm/absolute_hydration/test_simulation_slow.py => openmm_afe/test_ahfe_slow.py} (100%)
 rename openfe/tests/protocols/{openmm/absolute_hydration/test_tokenization.py => openmm_afe/test_ahfe_tokenization.py} (100%)
 rename openfe/tests/protocols/{openmm/absolute_hydration => openmm_md}/__init__.py (100%)
 rename openfe/tests/protocols/{openmm/plain_md/test_protocol.py => openmm_md/test_plain_md_protocol.py} (100%)
 rename openfe/tests/protocols/{openmm/plain_md/test_simulation_slow.py => openmm_md/test_plain_md_slow.py} (100%)
 rename openfe/tests/protocols/{openmm/plain_md/test_tokenization.py => openmm_md/test_plain_md_tokenization.py} (100%)
 rename openfe/tests/protocols/{openmm/hybridtop_rfe => openmm_rfe}/__init__.py (100%)
 rename openfe/tests/protocols/{openmm/hybridtop_rfe/test_protocol.py => openmm_rfe/test_hybrid_top_protocol.py} (100%)
 rename openfe/tests/protocols/{openmm/hybridtop_rfe/test_simulation_slow.py => openmm_rfe/test_hybrid_top_slow.py} (100%)
 rename openfe/tests/protocols/{openmm/hybridtop_rfe/test_tokenization.py => openmm_rfe/test_hybrid_top_tokenization.py} (100%)
 rename openfe/tests/protocols/{openmm => }/test_openmm_settings.py (100%)
 rename openfe/tests/protocols/{openmm => }/test_openmmutils.py (100%)

diff --git a/openfe/tests/protocols/openmm/plain_md/__init__.py b/openfe/tests/protocols/openmm/plain_md/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/openfe/tests/protocols/openmm/__init__.py b/openfe/tests/protocols/openmm_afe/__init__.py
similarity index 100%
rename from openfe/tests/protocols/openmm/__init__.py
rename to openfe/tests/protocols/openmm_afe/__init__.py
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/test_protocol.py b/openfe/tests/protocols/openmm_afe/test_ahfe_protocol.py
similarity index 100%
rename from openfe/tests/protocols/openmm/absolute_hydration/test_protocol.py
rename to openfe/tests/protocols/openmm_afe/test_ahfe_protocol.py
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py b/openfe/tests/protocols/openmm_afe/test_ahfe_slow.py
similarity index 100%
rename from openfe/tests/protocols/openmm/absolute_hydration/test_simulation_slow.py
rename to openfe/tests/protocols/openmm_afe/test_ahfe_slow.py
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/test_tokenization.py b/openfe/tests/protocols/openmm_afe/test_ahfe_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/openmm/absolute_hydration/test_tokenization.py
rename to openfe/tests/protocols/openmm_afe/test_ahfe_tokenization.py
diff --git a/openfe/tests/protocols/openmm/absolute_hydration/__init__.py b/openfe/tests/protocols/openmm_md/__init__.py
similarity index 100%
rename from openfe/tests/protocols/openmm/absolute_hydration/__init__.py
rename to openfe/tests/protocols/openmm_md/__init__.py
diff --git a/openfe/tests/protocols/openmm/plain_md/test_protocol.py b/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py
similarity index 100%
rename from openfe/tests/protocols/openmm/plain_md/test_protocol.py
rename to openfe/tests/protocols/openmm_md/test_plain_md_protocol.py
diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm_md/test_plain_md_slow.py
similarity index 100%
rename from openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py
rename to openfe/tests/protocols/openmm_md/test_plain_md_slow.py
diff --git a/openfe/tests/protocols/openmm/plain_md/test_tokenization.py b/openfe/tests/protocols/openmm_md/test_plain_md_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/openmm/plain_md/test_tokenization.py
rename to openfe/tests/protocols/openmm_md/test_plain_md_tokenization.py
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py b/openfe/tests/protocols/openmm_rfe/__init__.py
similarity index 100%
rename from openfe/tests/protocols/openmm/hybridtop_rfe/__init__.py
rename to openfe/tests/protocols/openmm_rfe/__init__.py
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/test_protocol.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
similarity index 100%
rename from openfe/tests/protocols/openmm/hybridtop_rfe/test_protocol.py
rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_slow.py
similarity index 100%
rename from openfe/tests/protocols/openmm/hybridtop_rfe/test_simulation_slow.py
rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_slow.py
diff --git a/openfe/tests/protocols/openmm/hybridtop_rfe/test_tokenization.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/openmm/hybridtop_rfe/test_tokenization.py
rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_tokenization.py
diff --git a/openfe/tests/protocols/openmm/test_openmm_settings.py b/openfe/tests/protocols/test_openmm_settings.py
similarity index 100%
rename from openfe/tests/protocols/openmm/test_openmm_settings.py
rename to openfe/tests/protocols/test_openmm_settings.py
diff --git a/openfe/tests/protocols/openmm/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py
similarity index 100%
rename from openfe/tests/protocols/openmm/test_openmmutils.py
rename to openfe/tests/protocols/test_openmmutils.py

From 8963fd21dc741d38d99a8b4845b0385a4d0217e5 Mon Sep 17 00:00:00 2001
From: IAlibay <IAlibay@users.noreply.github.com>
Date: Thu, 29 May 2025 00:03:44 +0100
Subject: [PATCH 10/10] AFE to AHFE

---
 openfe/tests/protocols/{openmm_afe => openmm_ahfe}/__init__.py    | 0
 .../protocols/{openmm_afe => openmm_ahfe}/test_ahfe_protocol.py   | 0
 .../tests/protocols/{openmm_afe => openmm_ahfe}/test_ahfe_slow.py | 0
 .../{openmm_afe => openmm_ahfe}/test_ahfe_tokenization.py         | 0
 4 files changed, 0 insertions(+), 0 deletions(-)
 rename openfe/tests/protocols/{openmm_afe => openmm_ahfe}/__init__.py (100%)
 rename openfe/tests/protocols/{openmm_afe => openmm_ahfe}/test_ahfe_protocol.py (100%)
 rename openfe/tests/protocols/{openmm_afe => openmm_ahfe}/test_ahfe_slow.py (100%)
 rename openfe/tests/protocols/{openmm_afe => openmm_ahfe}/test_ahfe_tokenization.py (100%)

diff --git a/openfe/tests/protocols/openmm_afe/__init__.py b/openfe/tests/protocols/openmm_ahfe/__init__.py
similarity index 100%
rename from openfe/tests/protocols/openmm_afe/__init__.py
rename to openfe/tests/protocols/openmm_ahfe/__init__.py
diff --git a/openfe/tests/protocols/openmm_afe/test_ahfe_protocol.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py
similarity index 100%
rename from openfe/tests/protocols/openmm_afe/test_ahfe_protocol.py
rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py
diff --git a/openfe/tests/protocols/openmm_afe/test_ahfe_slow.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_slow.py
similarity index 100%
rename from openfe/tests/protocols/openmm_afe/test_ahfe_slow.py
rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_slow.py
diff --git a/openfe/tests/protocols/openmm_afe/test_ahfe_tokenization.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_tokenization.py
similarity index 100%
rename from openfe/tests/protocols/openmm_afe/test_ahfe_tokenization.py
rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_tokenization.py