Update code to match new upstream behaviours (#33)

* update the build environment with the latest pins

* deal with upstream gufe changes in gufe object ordering

* Update to deal with interchange v0.4 change in signature

* update change in engine assignment

* Address ismorphic charged molecules check in Interchange creation

* blacken

* collate=True for topology creation

environment.yml

@@ -5,7 +5,7 @@ dependencies:
   - pip
   - gufe >=1.0
   - openff-toolkit >0.16.0
-  - openff-interchange >0.3.28
+  - openff-interchange >=0.4
   - openff-nagl-base >=0.3.3
   # OpenFE stack deps
   - duecredit<0.10
@@ -16,13 +16,13 @@ dependencies:
   - rdkit
   - packaging
   - pip
-  - pydantic >=1.10.17
+  - pydantic >=2.0
   - pyyaml
   - coverage
   - cinnabar ~=0.4.0
   - click
   - typing-extensions
-  - openmm >=8.0.0,!=8.1.0
+  - openmm >=8.0.0,!=8.1.0,<8.2.0
   - openmmtools
   - openmmforcefields
   - plugcli

src/pontibus/protocols/solvation/base.py

@@ -195,7 +195,7 @@ def _get_omm_objects(
             )
 
         # Get omm objects back
-        omm_topology = interchange.to_openmm_topology()
+        omm_topology = interchange.to_openmm_topology(collate=True)
         omm_system = interchange.to_openmm_system(
             hydrogen_mass=settings["forcefield_settings"].hydrogen_mass
         )
@@ -288,7 +288,8 @@ def run(
         try:
             # 12. Get context caches
             energy_ctx_cache, sampler_ctx_cache = self._get_ctx_caches(
-                settings["engine_settings"]
+                settings["forcefield_settings"],
+                settings["engine_settings"],
             )
 
             # 13. Get integrator

src/pontibus/tests/protocols/solvation/test_tokenization.py

@@ -61,7 +61,7 @@ def protocol_result(afe_solv_transformation_json):
 
 class TestProtocol(GufeTokenizableTestsMixin):
     cls = ASFEProtocol
-    key = "ASFEProtocol-798d96f939ae6898c385e31e48caae6d"
+    key = "ASFEProtocol-a9fe65baa34fb42a281cf9064ba9afa0"
     repr = f"<{key}>"
 
     @pytest.fixture()
@@ -71,7 +71,7 @@ def instance(self, protocol):
 
 class TestSolventUnit(GufeTokenizableTestsMixin):
     cls = ASFESolventUnit
-    repr = "ASFESolventUnit(Absolute Solvation, benzene solvent leg: repeat 2 generation 0)"
+    repr = "ASFESolventUnit(Absolute Solvation, benzene solvent leg"
     key = None
 
     @pytest.fixture()
@@ -81,12 +81,17 @@ def instance(self, solvent_protocol_unit):
     def test_key_stable(self):
         pytest.skip()
 
+    def test_repr(self, instance):
+        """
+        Overwrites the base `test_repr` call.
+        """
+        assert isinstance(repr(instance), str)
+        assert self.repr in repr(instance)
+
 
 class TestVacuumUnit(GufeTokenizableTestsMixin):
     cls = ASFEVacuumUnit
-    repr = (
-        "ASFEVacuumUnit(Absolute Solvation, benzene vacuum leg: repeat 2 generation 0)"
-    )
+    repr = "ASFEVacuumUnit(Absolute Solvation, benzene vacuum leg"
     key = None
 
     @pytest.fixture()
@@ -96,10 +101,17 @@ def instance(self, vacuum_protocol_unit):
     def test_key_stable(self):
         pytest.skip()
 
+    def test_repr(self, instance):
+        """
+        Overwrites the base `test_repr` call.
+        """
+        assert isinstance(repr(instance), str)
+        assert self.repr in repr(instance)
+
 
 class TestProtocolResult(GufeTokenizableTestsMixin):
     cls = ASFEProtocolResult
-    key = "ASFEProtocolResult-f1172ed96a55d778bdfcc8d9ce0299f2"
+    key = "ASFEProtocolResult-e711f21656c3795ed9d545c326ec717a"
     repr = f"<{key}>"
 
     @pytest.fixture()

src/pontibus/tests/utils/test_interchange_packmol.py

@@ -5,6 +5,7 @@
 
 import pytest
 from gufe import SmallMoleculeComponent, SolventComponent
+import numpy as np
 from numpy.testing import assert_allclose, assert_equal
 from openff.interchange.interop.openmm import to_openmm_positions
 from openff.toolkit import ForceField, Molecule
@@ -24,7 +25,7 @@
 )
 from pontibus.utils.molecules import WATER
 from pontibus.utils.system_creation import (
-    _check_charged_mols,
+    _check_and_deduplicate_charged_mols,
     _check_library_charges,
     _get_offmol_resname,
     _set_offmol_resname,
@@ -108,17 +109,38 @@ def test_check_library_charges_fail(methanol):
 
 
 def test_check_charged_mols_pass(methanol):
-    _check_charged_mols([methanol])
+    _check_and_deduplicate_charged_mols([methanol])
+
+
+def test_check_deduplicate_charged_mols(smc_components_benzene_unnamed):
+    """
+    Base test case for deduplication. Same molecule, same partial charges,
+    different conformer.
+    """
+    benzene1 = list(smc_components_benzene_unnamed.values())[0]
+    benzene1.assign_partial_charges(partial_charge_method="gasteiger")
+    benzene2 = Molecule.from_smiles("c1ccccc1")
+    benzene2.generate_conformers(n_conformers=1)
+    benzene2.assign_partial_charges(partial_charge_method="gasteiger")
+
+    assert all(benzene1.partial_charges == benzene2.partial_charges)
+    assert np.any(benzene1.conformers[0] != benzene2.conformers[0])
+    assert benzene1.is_isomorphic_with(benzene2)
+
+    uniques = _check_and_deduplicate_charged_mols([benzene1, benzene2])
+
+    assert len(uniques) == 1
+    assert uniques[0] == benzene1
 
 
 def test_check_charged_mols_nocharge(water_off, methanol):
     with pytest.raises(ValueError, match="One or more"):
-        _check_charged_mols([water_off, methanol])
+        _check_and_deduplicate_charged_mols([water_off, methanol])
 
 
 def test_check_charged_mols(water_off_am1bcc, water_off_named_charged):
     with pytest.raises(ValueError, match="different charges"):
-        _check_charged_mols([water_off_am1bcc, water_off_named_charged])
+        _check_and_deduplicate_charged_mols([water_off_am1bcc, water_off_named_charged])
 
 
 def test_protein_component_fail(smc_components_benzene_named, T4_protein_component):
@@ -312,6 +334,12 @@ def test_nonwater_solvent(smc_components_benzene_named, smiles):
     solvent_offmol = Molecule.from_smiles(smiles)
     solvent_offmol.assign_partial_charges(partial_charge_method="gasteiger")
 
+    if smiles == "c1ccccc1":
+        ligand = list(smc_components_benzene_named.values())[0]
+        assert solvent_offmol.is_isomorphic_with(ligand)
+        assert_allclose(solvent_offmol.partial_charges, ligand.partial_charges)
+        assert all(solvent_offmol.partial_charges == ligand.partial_charges)
+
     interchange, _ = interchange_packmol_creation(
         ffsettings=InterchangeFFSettings(
             forcefields=[
@@ -355,7 +383,7 @@ def omm_system(self, interchange_system):
     def omm_topology(self, interchange_system):
         interchange, _ = interchange_system
 
-        return interchange.to_openmm_topology()
+        return interchange.to_openmm_topology(collate=True)
 
     @pytest.fixture(scope="class")
     def nonbonds(self, omm_system):

src/pontibus/utils/system_creation.py

@@ -107,19 +107,28 @@ def _check_library_charges(
         raise ValueError(errmsg)
 
 
-def _check_charged_mols(molecules: list[OFFMolecule]) -> None:
+def _check_and_deduplicate_charged_mols(
+    molecules: list[OFFMolecule],
+) -> list[OFFMolecule]:
     """
-    Checks list of molecules passed to Interchange for partial charge
-    assignment.
+    Checks list of molecules with charges and removes any isomorphic
+    duplicates so that it can be passed to Interchange for partial
+    charge assignment.
 
     Parameters
     ----------
     molecules : list[openff.toolkit.Molecule]
+      A list of molecules with charges.
+
+    Returns
+    -------
+    unique_mols : list[openff.toolkit.Molecule]
+      A list of ismorphically unique molecules with charges.
 
     Raises
     ------
     ValueError
-      If any molecules in the list are isomorphic.
+      If any molecules in the list are isomorphic with different charges.
       If any molecules in the last have no charges.
     """
     if any(m.partial_charges is None for m in molecules):
@@ -130,20 +139,31 @@ def _check_charged_mols(molecules: list[OFFMolecule]) -> None:
         )
         raise ValueError(errmsg)
 
-    for i, moli in enumerate(molecules):
-        for j, molj in enumerate(molecules):
-            if i == j:
-                continue
-            if moli.is_isomorphic_with(molj) and not all(
-                moli.partial_charges == molj.partial_charges
-            ):
-                errmsg = (
-                    f"Isomorphic molecules {moli} and {molj}"
-                    "have been passed for partial charge "
-                    "assignment with different charges. "
-                    "This is not currently allowed."
-                )
-                raise ValueError(errmsg)
+    unique_mols = []
+
+    for moli in molecules:
+        isomorphic_mols = [
+            molj for molj in unique_mols if moli.is_isomorphic_with(molj)
+        ]
+
+        if isomorphic_mols:
+            # If we have any cases where there are isomorphic mols
+            # either:
+            # 1. They have the same charge so we don't add a second entry
+            # 2. They have different charges and it's an error.
+            for molj in isomorphic_mols:
+                if not all(moli.partial_charges == molj.partial_charges):
+                    errmsg = (
+                        f"Isomorphic molecules {moli} and {molj}"
+                        "have been passed for partial charge "
+                        "assignment with different charges. "
+                        "This is not currently allowed."
+                    )
+                    raise ValueError(errmsg)
+        else:
+            unique_mols.append(moli)
+
+    return unique_mols
 
 
 def interchange_packmol_creation(
@@ -300,6 +320,7 @@ def interchange_packmol_creation(
         topology = solvate_topology_nonwater(
             topology=topology,
             solvent=solvent_offmol,
+            target_density=0.95 * offunit.grams / offunit.mL,
             padding=solvation_settings.solvent_padding,
             box_shape=box_shape,
             tolerance=solvation_settings.packing_tolerance,
@@ -322,11 +343,11 @@ def interchange_packmol_creation(
 
     # Run validation checks on inputs to Interchange
     # Examples: https://github.com/openforcefield/openff-interchange/issues/1058
-    _check_charged_mols(charged_mols)
+    unique_charged_mols = _check_and_deduplicate_charged_mols(charged_mols)
 
     interchange = force_field.create_interchange(
         topology=topology,
-        charge_from_molecules=charged_mols,
+        charge_from_molecules=unique_charged_mols,
     )
 
     return interchange, comp_resids