Update tests to work without ambertools/with openeye

.github/workflows/ci.yaml

@@ -20,14 +20,14 @@ jobs:
       uses: pre-commit/action@v3.0.1
 
   test-matrix:
+    env:
+      OE_LICENSE: ${{ github.workspace }}/oe_license.txt
+
     needs: pre-commit
     strategy:
       matrix:
-        include:
-          - os: ubuntu-latest
-            pixi-platform: linux-64
-          - os: macos-latest
-            pixi-platform: osx-arm64
+        os: [ubuntu-latest, macos-latest]
+        pixi-environment: [default, pydantic-1-cpu, ambertools, openeye]
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
@@ -36,28 +36,34 @@ jobs:
     - name: Checkout
       uses: actions/checkout@v4
 
+    - name: Make oe_license.txt file from GH org secret "OE_LICENSE"
+      env:
+        OE_LICENSE_TEXT: ${{ secrets.OE_LICENSE }}
+      run: echo "${OE_LICENSE_TEXT}" > ${OE_LICENSE}
+
     - name: Setup pixi
       uses: prefix-dev/setup-pixi@v0.9.3
       with:
         pixi-version: v0.62.2
         cache: true
-        environments: default pydantic-1-cpu
-
-    - name: Run tests with pydantic >=2
-      run: |
-        pixi run test
+        # we use locked installs on macos-latest
+        # to avoid cross-platform resolution issues with cuda
+        locked: false
+        environments: ${{ matrix.pixi-environment }}
 
-    - name: Run tests with pydantic <2
+    - name: Run tests
       run: |
-        pixi run --environment pydantic-1-cpu test
+        pixi run --environment ${{ matrix.pixi-environment }} test
 
-    - name: Test examples with pydantic >=2
+    - name: Test examples (default environment only)
+      if: matrix.pixi-environment == 'default'
       run: |
-        pixi run test-examples
+        pixi run --environment ${{ matrix.pixi-environment }} test-examples
 
-    - name: Build docs
+    - name: Build docs (default environment only)
+      if: matrix.pixi-environment == 'default'
       run: |
-        pixi run docs
+        pixi run --environment ${{ matrix.pixi-environment }} docs
 
     - name: CodeCov
       uses: codecov/codecov-action@v4.1.1

pyproject.toml

@@ -60,6 +60,7 @@ pip = "*"
 openff-units = "*"
 openff-toolkit-base = ">=0.9.2"
 openff-interchange-base = ">=0.3.17"
+openff-models = "*"
 nnpops = "*"
 pydantic-units = "*"
 networkx = "*"
@@ -86,14 +87,22 @@ msgpack-python = "*"
 mdtraj = "*"
 absolv = ">=1.0.1"
 
+[tool.pixi.feature.openeye]
+channels = ["openeye"]
+
+[tool.pixi.feature.openeye.dependencies]
+openeye-toolkits = "*"
+
+[tool.pixi.feature.ambertools.dependencies]
+ambertools = "*"
+
 [tool.pixi.feature.examples.dependencies]
 jupyter = "*"
 nbconvert = "*"
 nglview = "*"
 smirnoff-plugins = "*"
 
 [tool.pixi.feature.dev.dependencies]
-ambertools = "*"
 scipy = "*"
 setuptools_scm = ">=8"
 pre-commit = "*"
@@ -123,11 +132,15 @@ pydantic = ">=2.0"
 pydantic = "<2.0"
 
 [tool.pixi.environments]
-default = { features = ["mm", "fe", "examples", "dev", "docs", "pydantic-2"] }
-cu12 = { features = ["cu12", "mm", "fe", "examples", "dev", "docs", "pydantic-2"] }
+default = { features = ["mm", "fe", "examples", "dev", "docs", "pydantic-2", "openeye"] }
+ambertools = { features = ["mm", "fe", "examples", "dev", "docs", "pydantic-2", "ambertools"] }
+# we get rid of examples here until smirnoff-plugins can be installed without ambertools
+# we also get rid of fe as absolv also installs ambertools
+openeye = { features = ["mm", "dev", "docs", "pydantic-2", "openeye"] }
+cu12 = { features = ["cu12", "mm", "fe", "examples", "dev", "docs", "pydantic-2", "ambertools"] }
 # The fe feature is excluded when pydantic < 2 as femto (used by absolv) needs pydantic >= 2
-pydantic-1-cpu = { features = ["mm", "examples", "dev", "docs", "pydantic-1"] }
-pydantic-1-cu12 = { features = ["cu12", "mm", "examples", "dev", "docs", "pydantic-1"] }
+pydantic-1-cpu = { features = ["mm", "examples", "dev", "docs", "pydantic-1", "ambertools"] }
+pydantic-1-cu12 = { features = ["cu12", "mm", "examples", "dev", "docs", "pydantic-1", "ambertools"] }
 
 [tool.pixi.tasks]
 lint = "ruff check smee examples"

smee/mm/_fe.py

@@ -8,7 +8,6 @@
 import numpy
 import openff.toolkit
 import openmm.unit
-import parmed.openmm
 import torch
 import yaml
 
@@ -112,6 +111,7 @@ def generate_dg_solv_data(
     import absolv.runner
     import femto.md.config
     import femto.md.system
+    import parmed.openmm
 
     output_dir = pathlib.Path.cwd() if output_dir is None else output_dir
 

smee/mm/_reporters.py

@@ -5,7 +5,6 @@
 import os
 import typing
 
-import msgpack
 import numpy
 import openmm.app
 import openmm.unit
@@ -72,6 +71,8 @@ def describeNextReport(self, simulation: openmm.app.Simulation):
         return steps, True, False, False, True
 
     def report(self, simulation: openmm.app.Simulation, state: openmm.State):
+        import msgpack
+
         potential_energy = state.getPotentialEnergy()
         kinetic_energy = state.getKineticEnergy()
 
@@ -107,6 +108,7 @@ def unpack_frames(
     file: typing.BinaryIO,
 ) -> typing.Generator[tuple[torch.Tensor, torch.Tensor, float], None, None]:
     """Unpack frames saved by a ``TensorReporter``."""
+    import msgpack
 
     unpacker = msgpack.Unpacker(file, object_hook=_decoder)
 

smee/tests/conftest.py

@@ -17,6 +17,21 @@
 _E = openff.units.unit.elementary_charge
 
 
+@pytest.fixture(scope="module")
+def toolkit_registry_rdkit_first():
+    """Returns a toolkit registry with RDKit as the first toolkit."""
+    return openff.toolkit.utils.ToolkitRegistry(
+        toolkit_precedence=[
+            openff.toolkit.utils.nagl_wrapper.NAGLToolkitWrapper,
+            openff.toolkit.utils.rdkit_wrapper.RDKitToolkitWrapper,
+            openff.toolkit.utils.openeye_wrapper.OpenEyeToolkitWrapper,
+            openff.toolkit.utils.ambertools_wrapper.AmberToolsToolkitWrapper,
+            openff.toolkit.utils.builtin_wrapper.BuiltInToolkitWrapper,
+        ],
+        exception_if_unavailable=False,
+    )
+
+
 @pytest.fixture
 def tmp_cwd(tmp_path, monkeypatch) -> pathlib.Path:
     monkeypatch.chdir(tmp_path)
@@ -59,11 +74,12 @@ def ethanol_conformer(ethanol) -> torch.Tensor:
 
 
 @pytest.fixture(scope="module")
-def ethanol_interchange(ethanol, default_force_field) -> openff.interchange.Interchange:
+def ethanol_interchange(
+    ethanol, default_force_field, toolkit_registry_rdkit_first
+) -> openff.interchange.Interchange:
     """Returns a parameterized system of ethanol."""
-
-    return openff.interchange.Interchange.from_smirnoff(
-        default_force_field, ethanol.to_topology()
+    return default_force_field.create_interchange(
+        ethanol.to_topology(), toolkit_registry=toolkit_registry_rdkit_first
     )
 
 
@@ -87,12 +103,12 @@ def formaldehyde_conformer(formaldehyde) -> torch.Tensor:
 
 @pytest.fixture(scope="module")
 def formaldehyde_interchange(
-    formaldehyde, default_force_field
+    formaldehyde, default_force_field, toolkit_registry_rdkit_first
 ) -> openff.interchange.Interchange:
     """Returns a parameterized system of formaldehyde."""
 
-    return openff.interchange.Interchange.from_smirnoff(
-        default_force_field, formaldehyde.to_topology()
+    return default_force_field.create_interchange(
+        formaldehyde.to_topology(), toolkit_registry=toolkit_registry_rdkit_first
     )
 
 

smee/tests/convertors/openff/test_nonbonded.py

@@ -1,3 +1,7 @@
+import importlib.util
+
+import pytest
+
 import openff.interchange
 import openff.interchange.models
 import openff.toolkit
@@ -12,6 +16,8 @@
     convert_vdw,
 )
 
+SMIRNOFF_PLUGINS_AVAILABLE = importlib.util.find_spec("smirnoff_plugins") is not None
+
 
 def test_convert_electrostatics_am1bcc(ethanol, ethanol_interchange):
     charge_collection = ethanol_interchange.collections["Electrostatics"]
@@ -57,7 +63,7 @@ def test_convert_electrostatics_am1bcc(ethanol, ethanol_interchange):
     assert parameter_map.exclusion_scale_idxs.shape == (n_expected_exclusions, 1)
 
 
-def test_convert_electrostatics_v_site():
+def test_convert_electrostatics_v_site(toolkit_registry_rdkit_first):
     force_field = openff.toolkit.ForceField()
     force_field.get_parameter_handler("Electrostatics")
     force_field.get_parameter_handler("vdW")
@@ -85,8 +91,10 @@ def test_convert_electrostatics_v_site():
 
     molecule = openff.toolkit.Molecule.from_mapped_smiles("[Cl:2]-[H:1]")
 
-    interchange = openff.interchange.Interchange.from_smirnoff(
-        force_field, molecule.to_topology(), allow_nonintegral_charges=True
+    interchange = force_field.create_interchange(
+        molecule.to_topology(),
+        toolkit_registry=toolkit_registry_rdkit_first,
+        allow_nonintegral_charges=True,
     )
     charge_collection = interchange.collections["Electrostatics"]
 
@@ -160,16 +168,18 @@ def test_convert_electrostatics_v_site():
     assert torch.allclose(parameter_map.exclusion_scale_idxs, expected_scales)
 
 
-def test_convert_electrostatics_tip4p():
+def test_convert_electrostatics_tip4p(toolkit_registry_rdkit_first):
     """Explicitly test the case of TIP4P (FB) water to make sure v-site charges are
     correct.
     """
 
     force_field = openff.toolkit.ForceField("tip4p_fb.offxml")
     molecule = openff.toolkit.Molecule.from_mapped_smiles("[H:2][O:1][H:3]")
 
-    interchange = openff.interchange.Interchange.from_smirnoff(
-        force_field, molecule.to_topology(), allow_nonintegral_charges=True
+    interchange = force_field.create_interchange(
+        molecule.to_topology(),
+        toolkit_registry=toolkit_registry_rdkit_first,
+        allow_nonintegral_charges=True,
     )
 
     tensor_top: smee.TensorTopology
@@ -186,7 +196,7 @@ def test_convert_electrostatics_tip4p():
     assert torch.allclose(charges, expected_charges)
 
 
-def test_convert_bci_and_vsite():
+def test_convert_bci_and_vsite(toolkit_registry_rdkit_first):
     ff_off = openff.toolkit.ForceField()
     ff_off.get_parameter_handler("Electrostatics")
     ff_off.get_parameter_handler("vdW")
@@ -213,8 +223,9 @@ def test_convert_bci_and_vsite():
     mol = openff.toolkit.Molecule.from_mapped_smiles("[O:1]([H:2])[H:3]")
     mol.assign_partial_charges(charge_handler.partial_charge_method)
 
-    interchange = openff.interchange.Interchange.from_smirnoff(
-        ff_off, mol.to_topology()
+    interchange = ff_off.create_interchange(
+        mol.to_topology(),
+        toolkit_registry=toolkit_registry_rdkit_first,
     )
 
     expected_charges = [
@@ -293,13 +304,18 @@ def test_convert_vdw(ethanol, ethanol_interchange):
     assert potential.fn == smee.EnergyFn.VDW_LJ
 
 
-def test_convert_dexp(ethanol, test_data_dir):
+@pytest.mark.skipif(
+    not SMIRNOFF_PLUGINS_AVAILABLE,
+    reason="Requires SMIRNOFF plugins to be installed.",
+)
+def test_convert_dexp(ethanol, test_data_dir, toolkit_registry_rdkit_first):
     ff = openff.toolkit.ForceField(
         str(test_data_dir / "de-ff.offxml"), load_plugins=True
     )
 
-    interchange = openff.interchange.Interchange.from_smirnoff(
-        ff, ethanol.to_topology()
+    interchange = ff.create_interchange(
+        ethanol.to_topology(),
+        toolkit_registry=toolkit_registry_rdkit_first,
     )
     vdw_collection = interchange.collections["DoubleExponential"]
 

smee/tests/convertors/openff/test_openff.py

@@ -2,6 +2,7 @@
 
 import openff.interchange.models
 import openff.toolkit
+from openff.toolkit.utils.toolkit_registry import toolkit_registry_manager
 import openff.units
 import pytest
 import torch
@@ -190,12 +191,13 @@ def test_convert_interchange_multiple(
     ethanol_interchange,
     formaldehyde_conformer,
     formaldehyde_interchange,
+    toolkit_registry_rdkit_first,
 ):
-    force_field, topologies = convert_interchange(
-        [ethanol_interchange, formaldehyde_interchange]
-    )
-    assert len(topologies) == 2
-
+    with toolkit_registry_manager(toolkit_registry_rdkit_first):
+        force_field, topologies = convert_interchange(
+            [ethanol_interchange, formaldehyde_interchange]
+        )
+        assert len(topologies) == 2
     expected_potentials = {
         "Angles",
         "Bonds",

smee/tests/convertors/openmm/test_openmm.py

@@ -1,3 +1,5 @@
+import importlib.util
+
 import numpy.random
 import openff.interchange
 import openff.toolkit
@@ -17,6 +19,8 @@
     create_openmm_system,
 )
 
+SMIRNOFF_PLUGINS_AVAILABLE = importlib.util.find_spec("smirnoff_plugins") is not None
+
 
 def _compute_energy(
     system: openmm.System,
@@ -252,6 +256,10 @@ def test_convert_lj_potential_with_exceptions(with_exception):
     )
 
 
+@pytest.mark.skipif(
+    not SMIRNOFF_PLUGINS_AVAILABLE,
+    reason="Requires SMIRNOFF plugins to be installed.",
+)
 def test_convert_to_openmm_system_dexp_periodic(test_data_dir):
     ff = openff.toolkit.ForceField(
         str(test_data_dir / "de-ff.offxml"), load_plugins=True