Make CUDA the default compute platform

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,14 @@
 )
 
 
+@pytest.fixture()
+def protocol_dry_settings():
+    settings = AbsoluteSolvationProtocol.default_settings()
+    settings.vacuum_engine_settings.compute_platform = None
+    settings.solvent_engine_settings.compute_platform = None
+    return settings
+
+
 @pytest.fixture()
 def default_settings():
     return AbsoluteSolvationProtocol.default_settings()
@@ -55,14 +63,17 @@ 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.protocol_repeats = 1
+    protocol_dry_settings.vacuum_simulation_settings.sampler_method = method
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -98,13 +109,12 @@ 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_dry_settings.protocol_repeats = 1
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-        settings=s,
+        settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -134,13 +144,14 @@ 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.protocol_repeats = 1
+    protocol_dry_settings.solvent_output_settings.output_indices = "resname UNK"
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -179,24 +190,25 @@ 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.protocol_repeats = 1
+    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 +238,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 +287,18 @@ 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_dry_settings.protocol_repeats = 1
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     def assign_fictitious_charges(offmol):
@@ -385,19 +397,19 @@ 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.protocol_repeats = 1
+    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 +453,13 @@ 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.protocol_repeats = 1
+    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 +485,9 @@ 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 = openmm_afe.AbsoluteSolvationProtocol(
-            settings=s,
+            settings=protocol_dry_settings,
     )
 
     stateA = ChemicalSystem({
@@ -672,20 +681,22 @@ 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.protocol_repeats = 1
+    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({

openfe/tests/protocols/openmm_md/test_plain_md_protocol.py

@@ -26,6 +26,14 @@
 from openfe.tests.conftest import HAS_ESPALOMA
 
 
+@pytest.fixture()
+def vac_settings():
+    settings = PlainMDProtocol.default_settings()
+    settings.forcefield_settings.nonbonded_method = 'nocutoff'
+    settings.engine_settings.compute_platform = None
+    return settings
+
+
 def test_create_default_settings():
     settings = PlainMDProtocol.default_settings()
 
@@ -87,10 +95,7 @@ def test_create_independent_repeat_ids(benzene_system):
     assert len(repeat_ids) == 6
 
 
-def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir):
-
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
+def test_dry_run_default_vacuum(benzene_vacuum_system, vac_settings, tmpdir):
 
     protocol = PlainMDProtocol(
             settings=vac_settings,
@@ -112,10 +117,8 @@ def test_dry_run_default_vacuum(benzene_vacuum_system, tmpdir):
             protocol.settings.thermo_settings.temperature)).barostat is None
 
 
-def test_dry_run_logger_output(benzene_vacuum_system, tmpdir, caplog):
+def test_dry_run_logger_output(benzene_vacuum_system, vac_settings, tmpdir, caplog):
 
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
     vac_settings.simulation_settings.equilibration_length_nvt = 1 * unit.picosecond
     vac_settings.simulation_settings.equilibration_length = 1 * unit.picosecond
     vac_settings.simulation_settings.production_length = 1 * unit.picosecond
@@ -143,10 +146,8 @@ def test_dry_run_logger_output(benzene_vacuum_system, tmpdir, caplog):
         assert "running production phase" in messages
 
 
-def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir):
+def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, vac_settings, tmpdir):
 
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
     vac_settings.output_settings.forcefield_cache = None
 
     protocol = PlainMDProtocol(
@@ -166,9 +167,7 @@ def test_dry_run_ffcache_none_vacuum(benzene_vacuum_system, tmpdir):
         dag_unit.run(dry=True)['debug']['system']
 
 
-def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir):
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
+def test_dry_run_gaff_vacuum(benzene_vacuum_system, vac_settings, tmpdir):
     vac_settings.forcefield_settings.small_molecule_forcefield = 'gaff-2.11'
 
     protocol = PlainMDProtocol(
@@ -187,9 +186,7 @@ def test_dry_run_gaff_vacuum(benzene_vacuum_system, tmpdir):
 
 
 @pytest.mark.skipif(not HAS_ESPALOMA, reason='espaloma is not available')
-def test_dry_run_espaloma_vacuum(benzene_vacuum_system, tmpdir):
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
+def test_dry_run_espaloma_vacuum(benzene_vacuum_system, vac_settings, tmpdir):
     vac_settings.forcefield_settings.small_molecule_forcefield = 'espaloma-0.3.2'
 
     protocol = PlainMDProtocol(
@@ -230,10 +227,8 @@ def test_dry_run_espaloma_vacuum(benzene_vacuum_system, tmpdir):
     ),
 ])
 def test_dry_run_charge_backends(
-    CN_molecule, tmpdir, method, backend, ref_key, am1bcc_ref_charges
+    CN_molecule, tmpdir, method, backend, ref_key, vac_settings, am1bcc_ref_charges
 ):
-    vac_settings = PlainMDProtocol.default_settings()
-    vac_settings.forcefield_settings.nonbonded_method = 'nocutoff'
     vac_settings.partial_charge_settings.partial_charge_method = method
     vac_settings.partial_charge_settings.off_toolkit_backend = backend
     vac_settings.partial_charge_settings.nagl_model = "openff-gnn-am1bcc-0.1.0-rc.1.pt"
@@ -268,6 +263,7 @@ def test_dry_many_molecules_solvent(
     A basic test flushing "will it work if you pass multiple molecules"
     """
     settings = PlainMDProtocol.default_settings()
+    settings.engine_settings.compute_platform = None
 
     protocol = PlainMDProtocol(
             settings=settings,
@@ -357,6 +353,7 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir):
     environment (waters)
     """
     settings = PlainMDProtocol.default_settings()
+    settings.engine_settings.compute_platform = None
     settings.forcefield_settings.forcefields = [
         "amber/ff14SB.xml",    # ff14SB protein force field
         "amber/tip4pew_standard.xml",  # FF we are testsing with the fun VS
@@ -386,6 +383,7 @@ def test_dry_run_ligand_tip4p(benzene_system, tmpdir):
 def test_dry_run_complex(benzene_complex_system, tmpdir):
     # this will be very time consuming
     settings = PlainMDProtocol.default_settings()
+    settings.engine_settings.compute_platform = None
 
     protocol = PlainMDProtocol(
             settings=settings,