Make CUDA the default compute platform

news/cuda.rst

@@ -0,0 +1,24 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* The default platform for OpenMM-based Protocols is now CUDA and will fail
+  by default on a non-Nvidia GPU enabled system.
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

openfe/protocols/openmm_utils/omm_settings.py

@@ -314,10 +314,13 @@ class OpenMMEngineSettings(SettingsBaseModel):
     * In the future make precision and deterministic forces user defined too.
     """
 
-    compute_platform: Optional[str] = None
+    compute_platform: Optional[str] = 'cuda'
     """
     OpenMM compute platform to perform MD integration with. If ``None``, will
-    choose fastest available platform. Default ``None``.
+    choose fastest available platform.
+    Allowed platforms are; ``cuda``, ``opencl``, ``cpu``.
+    Default ``cuda``.
+
     """
     gpu_device_index: Optional[list[int]] = None
     """
@@ -332,6 +335,15 @@ class OpenMMEngineSettings(SettingsBaseModel):
     Default ``None``.
     """
 
+    @field_validator('compute_platform')
+    def supported_sampler(cls, v):
+        supported = ['cpu', 'opencl', 'cuda']
+        if v is not None and v.lower() not in supported:
+            errmsg = ("Only the following OpenMM compute backends are "
+                      f"supported: {supported}")
+            raise ValueError(errmsg)
+        return v
+
 
 class IntegratorSettings(SettingsBaseModel):
     """Settings for the `LangevinDynamicsMove integrator <https://openmmtools.readthedocs.io/en/latest/api/generated/openmmtools.mcmc.LangevinDynamicsMove.html>`_.

openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py

@@ -29,6 +29,15 @@
 )
 
 
+@pytest.fixture()
+def protocol_dry_settings():
+    settings = AbsoluteSolvationProtocol.default_settings()
+    settings.vacuum_engine_settings.compute_platform = None
+    settings.solvent_engine_settings.compute_platform = None
+    settings.protocol_repeats = 1
+    return settings
+
+
 @pytest.fixture()
 def default_settings():
     return AbsoluteSolvationProtocol.default_settings()
@@ -55,14 +64,16 @@ def test_serialize_protocol(default_settings):
 @pytest.mark.parametrize('method', [
     'repex', 'sams', 'independent', 'InDePeNdENT'
 ])
-def test_dry_run_vac_benzene(benzene_modifications,
-                             method, tmpdir):
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.vacuum_simulation_settings.sampler_method = method
+def test_dry_run_vac_benzene(
+    benzene_modifications,
+    method,
+    protocol_dry_settings,
+    tmpdir
+):
+    protocol_dry_settings.vacuum_simulation_settings.sampler_method = method
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -98,13 +109,10 @@ def test_dry_run_vac_benzene(benzene_modifications,
         assert not vac_sampler.is_periodic
 
 
-def test_confgen_fail_AFE(benzene_modifications,  tmpdir):
+def test_confgen_fail_AFE(benzene_modifications, protocol_dry_settings, tmpdir):
     # check system parametrisation works even if confgen fails
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-        settings=s,
+        settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -134,13 +142,13 @@ def test_confgen_fail_AFE(benzene_modifications,  tmpdir):
             assert vac_sampler
 
 
-def test_dry_run_solv_benzene(benzene_modifications, tmpdir):
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.solvent_output_settings.output_indices = "resname UNK"
+def test_dry_run_solv_benzene(
+    benzene_modifications, protocol_dry_settings, tmpdir
+):
+    protocol_dry_settings.solvent_output_settings.output_indices = "resname UNK"
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -179,24 +187,24 @@ def test_dry_run_solv_benzene(benzene_modifications, tmpdir):
         assert pdb.n_atoms == 12
 
 
-def test_dry_run_solv_benzene_tip4p(benzene_modifications, tmpdir):
-    s = AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.vacuum_forcefield_settings.forcefields = [
+def test_dry_run_solv_benzene_tip4p(
+    benzene_modifications, protocol_dry_settings, tmpdir
+):
+    protocol_dry_settings.vacuum_forcefield_settings.forcefields = [
         "amber/ff14SB.xml",    # ff14SB protein force field
         "amber/tip4pew_standard.xml",  # FF we are testsing with the fun VS
         "amber/phosaa10.xml",  # Handles THE TPO
     ]
-    s.solvent_forcefield_settings.forcefields = [
+    protocol_dry_settings.solvent_forcefield_settings.forcefields = [
         "amber/ff14SB.xml",    # ff14SB protein force field
         "amber/tip4pew_standard.xml",  # FF we are testsing with the fun VS
         "amber/phosaa10.xml",  # Handles THE TPO
     ]
-    s.solvation_settings.solvent_model = 'tip4pew'
-    s.integrator_settings.reassign_velocities = True
+    protocol_dry_settings.solvation_settings.solvent_model = 'tip4pew'
+    protocol_dry_settings.integrator_settings.reassign_velocities = True
 
     protocol = AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -226,13 +234,12 @@ def test_dry_run_solv_benzene_tip4p(benzene_modifications, tmpdir):
 
 
 def test_dry_run_solv_benzene_noncubic(
-    benzene_modifications, tmpdir
+    benzene_modifications, protocol_dry_settings, tmpdir
 ):
-    s = AbsoluteSolvationProtocol.default_settings()
-    s.solvation_settings.solvent_padding = 1.5 * offunit.nanometer
-    s.solvation_settings.box_shape = 'dodecahedron'
+    protocol_dry_settings.solvation_settings.solvent_padding = 1.5 * offunit.nanometer
+    protocol_dry_settings.solvation_settings.box_shape = 'dodecahedron'
 
-    protocol = AbsoluteSolvationProtocol(settings=s)
+    protocol = AbsoluteSolvationProtocol(settings=protocol_dry_settings)
 
     stateA = ChemicalSystem({
         'benzene': benzene_modifications['benzene'],
@@ -276,17 +283,16 @@ def test_dry_run_solv_benzene_noncubic(
         )
 
 
-def test_dry_run_solv_user_charges_benzene(benzene_modifications, tmpdir):
+def test_dry_run_solv_user_charges_benzene(
+    benzene_modifications, protocol_dry_settings, tmpdir
+):
     """
     Create a test system with fictitious user supplied charges and
     ensure that they are properly passed through to the constructed
     alchemical system.
     """
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     def assign_fictitious_charges(offmol):
@@ -385,19 +391,18 @@ def assign_fictitious_charges(offmol):
     ),
 ])
 def test_dry_run_charge_backends(
-    CN_molecule, tmpdir, method, backend, ref_key, am1bcc_ref_charges
+    CN_molecule, tmpdir, method, backend, ref_key,
+    protocol_dry_settings, am1bcc_ref_charges
 ):
     """
     Check that partial charge generation with different backends
     works as expected.
     """
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.partial_charge_settings.partial_charge_method = method
-    s.partial_charge_settings.off_toolkit_backend = backend
-    s.partial_charge_settings.nagl_model = 'openff-gnn-am1bcc-0.1.0-rc.1.pt'
+    protocol_dry_settings.partial_charge_settings.partial_charge_method = method
+    protocol_dry_settings.partial_charge_settings.off_toolkit_backend = backend
+    protocol_dry_settings.partial_charge_settings.nagl_model = 'openff-gnn-am1bcc-0.1.0-rc.1.pt'
 
-    protocol = openmm_afe.AbsoluteSolvationProtocol(settings=s)
+    protocol = openmm_afe.AbsoluteSolvationProtocol(settings=protocol_dry_settings)
 
     # Create ChemicalSystems
     stateA = ChemicalSystem({
@@ -441,14 +446,12 @@ def test_dry_run_charge_backends(
     )
 
 
-def test_high_timestep(benzene_modifications, tmpdir):
-    s = AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.solvent_forcefield_settings.hydrogen_mass = 1.0
-    s.vacuum_forcefield_settings.hydrogen_mass = 1.0
+def test_high_timestep(benzene_modifications, protocol_dry_settings, tmpdir):
+    protocol_dry_settings.solvent_forcefield_settings.hydrogen_mass = 1.0
+    protocol_dry_settings.vacuum_forcefield_settings.hydrogen_mass = 1.0
 
     protocol = AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -474,11 +477,10 @@ def test_high_timestep(benzene_modifications, tmpdir):
 
 
 @pytest.fixture
-def benzene_solvation_dag(benzene_modifications):
-    s = AbsoluteSolvationProtocol.default_settings()
-
+def benzene_solvation_dag(benzene_modifications, protocol_dry_settings):
+    protocol_dry_settings.protocol_repeats = 3
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -672,20 +674,21 @@ def test_filenotfound_replica_states(self, protocolresult):
                          [[100 * offunit.picosecond, None],
                          [None, None],
                          [None, 100 * offunit.picosecond]])
-def test_dry_run_vacuum_write_frequency(benzene_modifications,
-                                        positions_write_frequency,
-                                        velocities_write_frequency,
-                                        tmpdir):
-    s = openmm_afe.AbsoluteSolvationProtocol.default_settings()
-    s.protocol_repeats = 1
-    s.solvent_output_settings.output_indices = "resname UNK"
-    s.solvent_output_settings.positions_write_frequency = positions_write_frequency
-    s.solvent_output_settings.velocities_write_frequency = velocities_write_frequency
-    s.vacuum_output_settings.positions_write_frequency = positions_write_frequency
-    s.vacuum_output_settings.velocities_write_frequency = velocities_write_frequency
+def test_dry_run_vacuum_write_frequency(
+    benzene_modifications,
+    positions_write_frequency,
+    velocities_write_frequency,
+    protocol_dry_settings,
+    tmpdir
+):
+    protocol_dry_settings.solvent_output_settings.output_indices = "resname UNK"
+    protocol_dry_settings.solvent_output_settings.positions_write_frequency = positions_write_frequency
+    protocol_dry_settings.solvent_output_settings.velocities_write_frequency = velocities_write_frequency
+    protocol_dry_settings.vacuum_output_settings.positions_write_frequency = positions_write_frequency
+    protocol_dry_settings.vacuum_output_settings.velocities_write_frequency = velocities_write_frequency
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({