Trying to speed up CI a little bit

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -182,7 +182,7 @@ def _validate_alchemical_components(
     Parameters
     ----------
     alchemical_components : dict[str, list[Component]]
-      Dictionary contatining the alchemical components for
+      Dictionary containing the alchemical components for
       states A and B.
     mapping : Optional[Union[ComponentMapping, list[ComponentMapping]]]
       all mappings between transforming components.

openfe/tests/conftest.py

@@ -7,6 +7,7 @@
 from importlib import resources
 from rdkit import Chem
 from rdkit.Chem import AllChem
+from openff.toolkit import Molecule
 from openff.units import unit
 import urllib.request
 

openfe/tests/protocols/conftest.py

@@ -23,16 +23,23 @@ def available_platforms() -> set[str]:
 
 
 @pytest.fixture
-def benzene_vacuum_system(benzene_modifications):
+def charged_benzene(benzene_modifications):
+    benzene_offmol = benzene_modifications['benzene'].to_openff()
+    benzene_offmol.assign_partial_charges(partial_charge_method='gasteiger')
+    return openfe.SmallMoleculeComponent.from_openff(benzene_offmol)
+
+
+@pytest.fixture
+def benzene_vacuum_system(charged_benzene):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['benzene']},
+        {'ligand': charged_benzene},
     )
 
 
-@pytest.fixture(scope='session')
-def benzene_system(benzene_modifications):
+@pytest.fixture
+def benzene_system(charged_benzene):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['benzene'],
+        {'ligand': charged_benzene,
          'solvent': openfe.SolventComponent(
              positive_ion='Na', negative_ion='Cl',
              ion_concentration=0.15 * unit.molar)
@@ -41,9 +48,9 @@ def benzene_system(benzene_modifications):
 
 
 @pytest.fixture
-def benzene_complex_system(benzene_modifications, T4_protein_component):
+def benzene_complex_system(charged_benzene, T4_protein_component):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['benzene'],
+        {'ligand': charged_benzene,
          'solvent': openfe.SolventComponent(
              positive_ion='Na', negative_ion='Cl',
              ion_concentration=0.15 * unit.molar),
@@ -52,16 +59,23 @@ def benzene_complex_system(benzene_modifications, T4_protein_component):
 
 
 @pytest.fixture
-def toluene_vacuum_system(benzene_modifications):
+def charged_toluene(benzene_modifications):
+    offmol = benzene_modifications['toluene'].to_openff()
+    offmol.assign_partial_charges(partial_charge_method='gasteiger')
+    return openfe.SmallMoleculeComponent.from_openff(offmol)
+
+
+@pytest.fixture
+def toluene_vacuum_system(charged_toluene):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['toluene']},
+        {'ligand': charged_toluene},
     )
 
 
-@pytest.fixture(scope='session')
-def toluene_system(benzene_modifications):
+@pytest.fixture
+def toluene_system(charged_toluene):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['toluene'],
+        {'ligand': charged_toluene,
          'solvent': openfe.SolventComponent(
              positive_ion='Na', negative_ion='Cl',
              ion_concentration=0.15 * unit.molar),
@@ -70,24 +84,20 @@ def toluene_system(benzene_modifications):
 
 
 @pytest.fixture
-def toluene_complex_system(benzene_modifications, T4_protein_component):
+def toluene_complex_system(charged_toluene, T4_protein_component):
     return openfe.ChemicalSystem(
-        {'ligand': benzene_modifications['toluene'],
+        {'ligand': charged_toluene,
          'solvent': openfe.SolventComponent(
              positive_ion='Na', negative_ion='Cl',
              ion_concentration=0.15 * unit.molar),
          'protein': T4_protein_component,}
     )
 
 
-@pytest.fixture(scope='session')
-def benzene_to_toluene_mapping(benzene_modifications):
+@pytest.fixture
+def benzene_to_toluene_mapping(charged_benzene, charged_toluene):
     mapper = openfe.setup.LomapAtomMapper(element_change=False)
-
-    molA = benzene_modifications['benzene']
-    molB = benzene_modifications['toluene']
-
-    return next(mapper.suggest_mappings(molA, molB))
+    return next(mapper.suggest_mappings(charged_benzene, charged_toluene))
 
 
 @pytest.fixture
@@ -142,10 +152,10 @@ def aniline_to_benzoic_mapping(benzene_charges):
 
 
 @pytest.fixture
-def benzene_many_solv_system(benzene_modifications):
+def benzene_many_solv_system(charged_benzene, charged_benzene_modifications):
 
-    rdmol_phenol = benzene_modifications['phenol'].to_rdkit()
-    rdmol_benzo = benzene_modifications['benzonitrile'].to_rdkit()
+    rdmol_phenol = charged_benzene_modifications['phenol'].to_rdkit()
+    rdmol_benzo = charged_benzene_modifications['benzonitrile'].to_rdkit()
 
     conf_phenol = rdmol_phenol.GetConformer()
     conf_benzo = rdmol_benzo.GetConformer()
@@ -167,18 +177,18 @@ def benzene_many_solv_system(benzene_modifications):
     )
 
     return openfe.ChemicalSystem(
-        {'whatligand': benzene_modifications['benzene'],
+        {'whatligand': charged_benzene,
          "foo": phenol,
          "bar": benzo,
          "solvent": openfe.SolventComponent()},
     )
 
 
 @pytest.fixture
-def toluene_many_solv_system(benzene_modifications):
+def toluene_many_solv_system(charged_toluene, charged_benzene_modifications):
 
-    rdmol_phenol = benzene_modifications['phenol'].to_rdkit()
-    rdmol_benzo = benzene_modifications['benzonitrile'].to_rdkit()
+    rdmol_phenol = charged_benzene_modifications['phenol'].to_rdkit()
+    rdmol_benzo = charged_benzene_modifications['benzonitrile'].to_rdkit()
 
     conf_phenol = rdmol_phenol.GetConformer()
     conf_benzo = rdmol_benzo.GetConformer()
@@ -199,7 +209,7 @@ def toluene_many_solv_system(benzene_modifications):
         rdmol_benzo, name='benzonitrile'
     )
     return openfe.ChemicalSystem(
-        {'whatligand': benzene_modifications['toluene'],
+        {'whatligand': charged_toluene,
          "foo": phenol,
          "bar": benzo,
          "solvent": openfe.SolventComponent()},

openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py

@@ -256,6 +256,7 @@ def test_dry_run_gaff_vacuum(benzene_vacuum_system, toluene_vacuum_system,
 
 
 @pytest.mark.slow
+## TODO: this is breaking because the components are not the same gufe tokenizables as in the mapping
 def test_dry_many_molecules_solvent(
     benzene_many_solv_system, toluene_many_solv_system,
     benzene_to_toluene_mapping, tmpdir
@@ -430,8 +431,11 @@ def test_confgen_mocked_fail(benzene_system, toluene_system,
 
     protocol = openmm_rfe.RelativeHybridTopologyProtocol(settings=settings)
 
-    dag = protocol.create(stateA=benzene_system, stateB=toluene_system,
-                          mapping=benzene_to_toluene_mapping)
+    dag = protocol.create(
+        stateA=benzene_system,
+        stateB=toluene_system,
+        mapping=benzene_to_toluene_mapping,
+    )
     dag_unit = list(dag.protocol_units)[0]
 
     with tmpdir.as_cwd():
@@ -443,12 +447,34 @@ def test_confgen_mocked_fail(benzene_system, toluene_system,
 
 @pytest.fixture(scope='session')
 def tip4p_hybrid_factory(
-    benzene_system, toluene_system,
-    benzene_to_toluene_mapping, tmp_path_factory,
+    benzene_modifications, tmp_path_factory,
 ):
     """
     Hybrid system with virtual sites in the environment (waters)
     """
+    # Session scoped, so we do things by hand here
+    # Generate the end state systems
+    benzene_offmol = benzene_modifications['benzene'].to_openff()
+    benzene_offmol.assign_partial_charges(partial_charge_method='gasteiger')
+    benzene = openfe.SmallMoleculeComponent.from_openff(benzene_offmol)
+
+    toluene_offmol = benzene_modifications['toluene'].to_openff()
+    toluene_offmol.assign_partial_charges(partial_charge_method='gasteiger')
+    toluene = openfe.SmallMoleculeComponent.from_openff(toluene_offmol)
+
+    solvent = openfe.SolventComponent(
+        positive_ion='Na',
+        negative_ion='Cl',
+        ion_concentration=0.15 * unit.molar
+    )
+
+    stateA = openfe.ChemicalSystem({'ligand': benzene, 'solvent': solvent})
+    stateB = openfe.ChemicalSystem({'ligand': toluene, 'solvent': solvent})
+
+    # Now the mapping
+    mapper = openfe.setup.LomapAtomMapper(element_change=False)
+    mapping = next(mapper.suggest_mappings(benzene, toluene))
+
     settings = openmm_rfe.RelativeHybridTopologyProtocol.default_settings()
     settings.forcefield_settings.forcefields = [
         "amber/ff14SB.xml",    # ff14SB protein force field
@@ -464,9 +490,9 @@ def tip4p_hybrid_factory(
             settings=settings,
     )
     dag = protocol.create(
-        stateA=benzene_system,
-        stateB=toluene_system,
-        mapping=benzene_to_toluene_mapping,
+        stateA=stateA,
+        stateB=stateB,
+        mapping=mapping,
     )
     dag_unit = list(dag.protocol_units)[0]
 
@@ -999,19 +1025,16 @@ def test_missing_ligand(benzene_system, benzene_to_toluene_mapping):
         )
 
 
-def test_vaccuum_PME_error(benzene_vacuum_system, benzene_modifications,
-                           benzene_to_toluene_mapping):
-    # state B doesn't have a solvent component (i.e. its vacuum)
-    stateB = openfe.ChemicalSystem({'ligand': benzene_modifications['toluene']})
-
+def test_vacuum_PME_error(benzene_vacuum_system, toluene_vacuum_system,
+                          benzene_to_toluene_mapping):
     p = openmm_rfe.RelativeHybridTopologyProtocol(
         settings=openmm_rfe.RelativeHybridTopologyProtocol.default_settings(),
     )
     errmsg = "PME cannot be used for vacuum transform"
     with pytest.raises(ValueError, match=errmsg):
         _ = p.create(
             stateA=benzene_vacuum_system,
-            stateB=stateB,
+            stateB=toluene_vacuum_system,
             mapping=benzene_to_toluene_mapping,
         )
 

openfe/tests/protocols/test_openmmutils.py

@@ -230,8 +230,9 @@ def get_settings():
 
 class TestFEAnalysis:
 
-    # Note: class scope _will_ cause this to segfault - the reporter has to close
-    @pytest.fixture(scope='function')
+    # Note: scope on this can sometimes cause segfault, may need to revert to
+    # function scope if it happens.
+    @pytest.fixture(scope='class')
     def reporter(self):
         with resources.as_file(resources.files('openfe.tests.data.openmm_rfe')) as d:
             ncfile = str(d / 'vacuum_nocoord.nc')
@@ -247,7 +248,7 @@ def reporter(self):
         finally:
             r.close()
 
-    @pytest.fixture()
+    @pytest.fixture(scope='class')
     def analyzer(self, reporter):
         return multistate_analysis.MultistateEquilFEAnalysis(
             reporter, sampling_method='repex',