Repository: openforcefield/openff-toolkit

Issue:
Add isomorphism checks for `_SimpleMolecule`
**Is your feature request related to a problem? Please describe.**
<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] -->

When I do something like `Interchange.from_openmm`, I end up with a topology of `_SimpleMolecule`. Exporting these to GROMACS loosely requires an isomorphism check that is currently unsupported by the toolkit. It's called from `Topology.identical_molecule_groups` or something similar, but ultimately fails on `Molecule.are_isomorphic` not knowing how to deal with this input.

**Describe the solution you'd like**
<!-- A clear and concise description of what you want to happen. -->

I'd like `Molecule.is_isomorphic_with` to work on `_SimpleMolecule`.

**Describe alternatives you've considered**
<!-- A clear and concise description of any alternative solutions or features you've considered. -->

I can monkey-patch this, but not easily or without copying sizable chunks of code.

Another backdoor is allowing GROMACS to not de-duplicate isomorphic molecules, but this solution is unworkable for large systems with many similar molecules, i.e. water or solvent.

**Additional context**
<!-- Add any other context or screenshots about the feature request here. -->


Discussion:
codecov[bot]: ## [Codecov](https://codecov.io/gh/openforcefield/openff-toolkit/pull/1581?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) Report
> Merging [#1581](https://codecov.io/gh/openforcefield/openff-toolkit/pull/1581?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) (e482faf) into [main](https://codecov.io/gh/openforcefield/openff-toolkit/commit/d5a1c1f572429324b526803b73d3cf4d1d5a03ba?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) (d5a1c1f) will **decrease** coverage by `0.22%`.
> The diff coverage is `100.00%`.

<details><summary>Additional details and impacted files</summary>


</details>
j-wags: Glorious!

Patch:
diff --git a/openff/toolkit/topology/_mm_molecule.py b/openff/toolkit/topology/_mm_molecule.py
index 3b272927..8e5fac62 100644
--- a/openff/toolkit/topology/_mm_molecule.py
+++ b/openff/toolkit/topology/_mm_molecule.py
@@ -9,7 +9,7 @@ TypedMolecule TODOs
   deserialize a Molecule or a TypedMolecule.
 
 """
-from typing import TYPE_CHECKING, Dict, List, NoReturn, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, NoReturn, Optional, Tuple, Union
 
 from openff.units import unit
 from openff.units.elements import MASSES, SYMBOLS
@@ -26,6 +26,8 @@ if TYPE_CHECKING:
     import networkx as nx
     from openff.units.unit import Quantity
 
+    from openff.toolkit.topology.molecule import FrozenMolecule
+
 
 class _SimpleMolecule:
     def __init__(self):
@@ -311,9 +313,11 @@ class _SimpleMolecule:
             "an OpenFF Molecule with sufficiently specified chemistry."
         )
 
-    def _is_isomorphic_with(self, other) -> bool:
+    def is_isomorphic_with(
+        self, other: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"], **kwargs
+    ) -> bool:
         """
-        Untrustworthy check for pseudo-isomorphism.
+        Check for pseudo-isomorphism.
 
         This currently checks that the two molecules have
         * The same number of atoms
@@ -323,18 +327,66 @@ class _SimpleMolecule:
         This currently does NOT checks that the two molecules have
         * Topologically identical bond graphs
         """
+        return _SimpleMolecule.are_isomorphic(
+            self,
+            other,
+            return_atom_map=False,
+        )[0]
+
+    @staticmethod
+    def are_isomorphic(
+        mol1: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"],
+        mol2: Union["FrozenMolecule", "_SimpleMolecule", "nx.Graph"],
+        return_atom_map: bool = False,
+    ) -> Tuple[bool, Optional[Dict[int, int]]]:
+        import networkx
+
+        _cls = _SimpleMolecule
+
+        if isinstance(mol1, networkx.Graph) and isinstance(mol2, networkx.Graph):
+            graph1 = _SimpleMolecule._from_subgraph(mol1).to_networkx()
+            graph2 = _SimpleMolecule._from_subgraph(mol2).to_networkx()
+
+        elif isinstance(mol1, networkx.Graph):
+            assert isinstance(mol2, _cls)
+            graph1 = _SimpleMolecule._from_subgraph(mol1).to_networkx()
+            graph2 = mol2.to_networkx()
+
+        elif isinstance(mol2, networkx.Graph):
+            assert isinstance(mol1, _cls)
+            graph1 = mol1.to_networkx()
+            graph2 = _SimpleMolecule._from_subgraph(mol2).to_networkx()
+
+        else:
+            # static methods (by definition) know nothing about their class,
+            # so the class to compare to must be hard-coded here
+            if not (isinstance(mol1, _cls) and isinstance(mol2, _cls)):
+                return False, None
+
+            graph1 = mol1.to_networkx()
+            graph2 = mol2.to_networkx()
 
-        if self.n_atoms != other.n_atoms:
-            return False
+        def node_match_func(node1, node2):
+            return node1["atomic_number"] == node2["atomic_number"]
 
-        if self.n_bonds != other.n_bonds:
-            return False
+        edge_match_func = None
 
-        for this_atom, other_atom in zip(self.atoms, other.atoms):
-            if this_atom.atomic_number != other_atom.atomic_number:
-                return False
+        matcher = networkx.algorithms.isomorphism.GraphMatcher(
+            graph1, graph2, node_match=node_match_func, edge_match=edge_match_func
+        )
+
+        if matcher.is_isomorphic():
+            if return_atom_map:
+                topology_atom_map = matcher.mapping
 
-        return True
+                return True, {
+                    key: topology_atom_map[key] for key in sorted(topology_atom_map)
+                }
+
+            else:
+                return True, None
+        else:
+            return False, None
 
     def generate_unique_atom_names(self):
         """Generate unique atom names. See `Molecule.generate_unique_atom_names`."""
diff --git a/openff/toolkit/topology/molecule.py b/openff/toolkit/topology/molecule.py
index f0229c4b..8d8b24cf 100644
--- a/openff/toolkit/topology/molecule.py
+++ b/openff/toolkit/topology/molecule.py
@@ -1893,8 +1893,8 @@ class FrozenMolecule(Serializable):
 
     @staticmethod
     def are_isomorphic(
-        mol1: Union["FrozenMolecule", nx.Graph],
-        mol2: Union["FrozenMolecule", nx.Graph],
+        mol1: Union["FrozenMolecule", "_SimpleMolecule", nx.Graph],
+        mol2: Union["FrozenMolecule", "_SimpleMolecule", nx.Graph],
         return_atom_map: bool = False,
         aromatic_matching: bool = True,
         formal_charge_matching: bool = True,
@@ -1903,7 +1903,7 @@ class FrozenMolecule(Serializable):
         bond_stereochemistry_matching: bool = True,
         strip_pyrimidal_n_atom_stereo: bool = True,
         toolkit_registry: TKR = GLOBAL_TOOLKIT_REGISTRY,
-    ):
+    ) -> Tuple[bool, Optional[Dict[int, int]]]:
         """
         Determine if ``mol1`` is isomorphic to ``mol2``.
 
@@ -1982,10 +1982,26 @@ class FrozenMolecule(Serializable):
             [Dict[int,int]] ordered by mol1 indexing {mol1_index: mol2_index}
             If molecules are not isomorphic given input arguments, will return None instead of dict.
         """
+        import networkx as nx
 
-        def _object_to_n_atoms(obj):
-            import networkx as nx
+        _cls = FrozenMolecule
+
+        if isinstance(mol1, nx.Graph) and isinstance(mol2, nx.Graph):
+            pass
+
+        elif isinstance(mol1, nx.Graph):
+            assert isinstance(mol2, _cls)
 
+        elif isinstance(mol2, nx.Graph):
+            assert isinstance(mol1, _cls)
+
+        else:
+            # static methods (by definition) know nothing about their class,
+            # so the class to compare to must be hard-coded here
+            if not (isinstance(mol1, _cls) and isinstance(mol2, _cls)):
+                return False, None
+
+        def _object_to_n_atoms(obj):
             if isinstance(obj, FrozenMolecule):
                 return obj.n_atoms
             elif isinstance(obj, nx.Graph):
@@ -2052,10 +2068,8 @@ class FrozenMolecule(Serializable):
             edge_match_func = None  # type: ignore
 
         # Here we should work out what data type we have, also deal with lists?
-        def to_networkx(data):
+        def to_networkx(data: Union[FrozenMolecule, nx.Graph]) -> nx.Graph:
             """For the given data type, return the networkx graph"""
-            import networkx as nx
-
             if strip_pyrimidal_n_atom_stereo:
                 SMARTS = "[N+0X3:1](-[*])(-[*])(-[*])"
 
@@ -2081,6 +2095,7 @@ class FrozenMolecule(Serializable):
 
         mol1_netx = to_networkx(mol1)
         mol2_netx = to_networkx(mol2)
+
         from networkx.algorithms.isomorphism import GraphMatcher  # type: ignore
 
         GM = GraphMatcher(
@@ -2101,7 +2116,9 @@ class FrozenMolecule(Serializable):
         else:
             return isomorphic, None
 
-    def is_isomorphic_with(self, other: Union["FrozenMolecule", nx.Graph], **kwargs):
+    def is_isomorphic_with(
+        self, other: Union["FrozenMolecule", "_SimpleMolecule", nx.Graph], **kwargs
+    ) -> bool:
         """
         Check if the molecule is isomorphic with the other molecule which can be an openff.toolkit.topology.Molecule
         or nx.Graph(). Full matching is done using the options described bellow.
@@ -3474,7 +3491,7 @@ class FrozenMolecule(Serializable):
         return self._hill_formula
 
     @staticmethod
-    def _object_to_hill_formula(obj: Union["Molecule", "nx.Graph"]) -> str:
+    def _object_to_hill_formula(obj: Union["FrozenMolecule", "nx.Graph"]) -> str:
         """Take a Molecule or NetworkX graph and generate its Hill formula.
         This provides a backdoor to the old functionality of Molecule.to_hill_formula, which
         was a static method that duck-typed inputs of Molecule or graph objects."""
diff --git a/openff/toolkit/topology/topology.py b/openff/toolkit/topology/topology.py
index 5d222052..b02b86c2 100644
--- a/openff/toolkit/topology/topology.py
+++ b/openff/toolkit/topology/topology.py
@@ -1175,9 +1175,13 @@ class Topology(Serializable):
                 if mol2_idx in already_matched_mols:
                     continue
                 mol2 = self.molecule(mol2_idx)
-                are_isomorphic, atom_map = Molecule.are_isomorphic(
-                    mol1, mol2, return_atom_map=True
-                )
+                if isinstance(mol1, type(mol2)) or isinstance(mol2, type(mol1)):
+                    are_isomorphic, atom_map = mol1.are_isomorphic(
+                        mol1, mol2, return_atom_map=True
+                    )
+                else:
+                    are_isomorphic = False
+
                 if are_isomorphic:
                     identity_maps[mol2_idx] = (
                         mol1_idx,
@@ -1417,12 +1421,12 @@ class Topology(Serializable):
                 )
                 if isomorphic:
                     # Take the first valid atom indexing map
-                    first_topology_atom_index = min(mapping.keys())
+                    first_topology_atom_index = min(mapping.keys())  # type: ignore[union-attr]
                     topology_molecules_to_add.append(
                         (
                             first_topology_atom_index,
                             unq_mol_G,
-                            mapping.items(),
+                            mapping.items(),  # type: ignore[union-attr]
                             omm_mol_G,
                         )
                     )


Test patch:
diff --git a/openff/toolkit/tests/test_mm_molecule.py b/openff/toolkit/tests/test_mm_molecule.py
index 00e86456..f8d38a78 100644
--- a/openff/toolkit/tests/test_mm_molecule.py
+++ b/openff/toolkit/tests/test_mm_molecule.py
@@ -1,40 +1,76 @@
 import numpy as np
 import pytest
 
+from openff.toolkit.tests.create_molecules import create_ethanol
 from openff.toolkit.topology._mm_molecule import _SimpleMolecule
 from openff.toolkit.topology.molecule import Molecule
 
 
-class TestMMMolecule:
-    @pytest.fixture()
-    def water(self):
-        water = _SimpleMolecule()
-        water.add_atom(atomic_number=8)
-        water.add_atom(atomic_number=1)
-        water.add_atom(atomic_number=1)
-        water.add_bond(0, 1)
-        water.add_bond(0, 2)
+@pytest.fixture()
+def water():
+    water = _SimpleMolecule()
+    water.add_atom(atomic_number=8)
+    water.add_atom(atomic_number=1)
+    water.add_atom(atomic_number=1)
+    water.add_bond(0, 1)
+    water.add_bond(0, 2)
 
-        return water
+    return water
 
-    @pytest.fixture()
-    def molecule_with_zero_atom(self):
-        molecule = _SimpleMolecule()
-        molecule.add_atom(atomic_number=6)
-        molecule.add_atom(atomic_number=0)
-        molecule.add_bond(0, 1)
 
-        return molecule
+@pytest.fixture()
+def methane():
+    methane = _SimpleMolecule()
+    methane.add_atom(atomic_number=6)
+    methane.add_atom(atomic_number=1)
+    methane.add_atom(atomic_number=1)
+    methane.add_atom(atomic_number=1)
+    methane.add_atom(atomic_number=1)
+    methane.add_bond(0, 1)
+    methane.add_bond(0, 2)
+    methane.add_bond(0, 3)
+    methane.add_bond(0, 4)
+
+    return methane
+
+
+@pytest.fixture()
+def methanol():
+    methanol = _SimpleMolecule()
+    methanol.add_atom(atomic_number=8)
+    methanol.add_atom(atomic_number=6)
+    methanol.add_atom(atomic_number=1)
+    methanol.add_atom(atomic_number=1)
+    methanol.add_atom(atomic_number=1)
+    methanol.add_bond(0, 1)
+    methanol.add_bond(1, 2)
+    methanol.add_bond(1, 3)
+    methanol.add_bond(1, 4)
+
+    return methanol
+
+
+@pytest.fixture()
+def molecule_with_zero_atom():
+    molecule = _SimpleMolecule()
+    molecule.add_atom(atomic_number=6)
+    molecule.add_atom(atomic_number=0)
+    molecule.add_bond(0, 1)
+
+    return molecule
 
-    @pytest.fixture()
-    def molecule_with_bogus_atom(self):
-        molecule = _SimpleMolecule()
-        molecule.add_atom(atomic_number=6)
-        molecule.add_atom(atomic_number=-1)
-        molecule.add_bond(0, 1)
 
-        return molecule
+@pytest.fixture()
+def molecule_with_bogus_atom():
+    molecule = _SimpleMolecule()
+    molecule.add_atom(atomic_number=6)
+    molecule.add_atom(atomic_number=-1)
+    molecule.add_bond(0, 1)
 
+    return molecule
+
+
+class TestMMMolecule:
     def test_create_water(self):
         water = _SimpleMolecule()
         water.add_atom(atomic_number=8)
@@ -112,3 +148,84 @@ class TestMMMolecule:
         for atom_index in range(converted.n_atoms):
             found = converted.atom(atom_index).atomic_number
             assert found == expected_atomic_numbers[atom_index]
+
+
+class TestIsomorphism:
+    @pytest.fixture()
+    def n_propanol(self):
+        return _SimpleMolecule.from_molecule(
+            Molecule.from_mapped_smiles(
+                "[H:5][C:1]([H:6])([H:7])[C:2]([H:8])([H:9])[C:3]([H:10])([H:11])[O:4][H:12]"
+            )
+        )
+
+    @pytest.fixture()
+    def iso_propanol(self):
+        return _SimpleMolecule.from_molecule(
+            Molecule.from_mapped_smiles(
+                "[H:5][C:1]([H:6])([H:7])[C:2]([H:8])([C:3]([H:10])([H:11])[H:12])[O:4][H:9]"
+            )
+        )
+
+    @pytest.fixture()
+    def o_dichlorobezene(self):
+        return _SimpleMolecule.from_molecule(
+            Molecule.from_mapped_smiles(
+                "[C:1]=1([Cl:7])[C:2]([H:9])=[C:3]([H:10])[C:4]([H:11])=[C:5]([H:12])[C:6]([Cl:8])=1"
+            )
+        )
+
+    @pytest.fixture()
+    def m_dichlorobezene(self):
+        return _SimpleMolecule.from_molecule(
+            Molecule.from_mapped_smiles(
+                "[C:1]([Cl:7])(=[C:2]1([H:9]))[C:3]([H:10])=[C:4]([H:11])[C:5]([H:12])=[C:6]([Cl:8])1"
+            )
+        )
+
+    @pytest.mark.parametrize("as_graphs", [True, False])
+    def test_are_isomorphic(self, water, methane, methanol, as_graphs):
+        if as_graphs:
+            water = water.to_networkx()
+            methane = methane.to_networkx()
+            methanol = methanol.to_networkx()
+
+        assert _SimpleMolecule.are_isomorphic(water, water)[0]
+        assert _SimpleMolecule.are_isomorphic(methane, methane)[0]
+        assert _SimpleMolecule.are_isomorphic(methanol, methanol)[0]
+        assert not _SimpleMolecule.are_isomorphic(water, methane)[0]
+        assert not _SimpleMolecule.are_isomorphic(water, methanol)[0]
+        assert not _SimpleMolecule.are_isomorphic(methane, methanol)[0]
+
+    def test_is_isomorphic_with(self, water, methane, methanol):
+        assert water.is_isomorphic_with(water)
+        assert methane.is_isomorphic_with(methane)
+        assert methanol.is_isomorphic_with(methanol)
+        assert not water.is_isomorphic_with(methane)
+        assert not water.is_isomorphic_with(methanol)
+        assert not methane.is_isomorphic_with(methanol)
+
+    def test_short_circuit_heterogeneous_input(self):
+        assert not _SimpleMolecule.are_isomorphic(
+            create_ethanol(),
+            _SimpleMolecule.from_molecule(create_ethanol()),
+        )[0]
+
+        assert not _SimpleMolecule.are_isomorphic(
+            _SimpleMolecule.from_molecule(create_ethanol()),
+            create_ethanol(),
+        )[0]
+
+    def test_graph_and_molecule_inputs(self, methanol):
+        graph = methanol.to_networkx()
+
+        assert _SimpleMolecule.are_isomorphic(methanol, graph)[0]
+        assert _SimpleMolecule.are_isomorphic(graph, methanol)[0]
+
+    def test_propanol_isopropanol_not_isomorphic(self, n_propanol, iso_propanol):
+        assert not _SimpleMolecule.are_isomorphic(n_propanol, iso_propanol)[0]
+
+    def test_o_m_dichlorobenzene_not_isomorphic(
+        self, o_dichlorobezene, m_dichlorobezene
+    ):
+        assert not _SimpleMolecule.are_isomorphic(o_dichlorobezene, m_dichlorobezene)[0]
diff --git a/openff/toolkit/tests/test_molecule.py b/openff/toolkit/tests/test_molecule.py
index 9189e021..329b99dd 100644
--- a/openff/toolkit/tests/test_molecule.py
+++ b/openff/toolkit/tests/test_molecule.py
@@ -1733,6 +1733,26 @@ class TestMolecule:
             "CCC[N@@](C)CC"
         )
 
+    def test_short_circuit_heterogeneous_input(self):
+        from openff.toolkit.topology._mm_molecule import _SimpleMolecule
+
+        assert not Molecule.are_isomorphic(
+            create_ethanol(),
+            _SimpleMolecule.from_molecule(create_ethanol()),
+        )[0]
+
+        assert not Molecule.are_isomorphic(
+            _SimpleMolecule.from_molecule(create_ethanol()),
+            create_ethanol(),
+        )[0]
+
+    def test_graph_and_molecule_inputs(self):
+        molecule = create_ethanol()
+        graph = molecule.to_networkx()
+
+        assert Molecule.are_isomorphic(molecule, graph)[0]
+        assert Molecule.are_isomorphic(graph, molecule)[0]
+
     class TestRemap:
         """Tests for the ``Molecule.remap()`` method"""
 
diff --git a/openff/toolkit/tests/test_topology.py b/openff/toolkit/tests/test_topology.py
index 229b8cf6..0fbd6039 100644
--- a/openff/toolkit/tests/test_topology.py
+++ b/openff/toolkit/tests/test_topology.py
@@ -42,6 +42,7 @@ from openff.toolkit.topology import (
     Topology,
     ValenceDict,
 )
+from openff.toolkit.topology._mm_molecule import _SimpleMolecule
 from openff.toolkit.utils import (
     BASIC_CHEMINFORMATICS_TOOLKITS,
     OPENEYE_AVAILABLE,
@@ -107,6 +108,18 @@ def test_cheminformatics_toolkit_is_installed():
         raise Exception(msg)
 
 
+@pytest.fixture()
+def mixed_topology():
+    return Topology.from_molecules(
+        [
+            create_ethanol(),
+            create_ethanol(),
+            _SimpleMolecule.from_molecule(create_ethanol()),
+            _SimpleMolecule.from_molecule(create_ethanol()),
+        ]
+    )
+
+
 # TODO: Refactor this to pytest
 class TestTopology:
     def test_empty(self):
@@ -1523,6 +1536,9 @@ class TestTopology:
         assert_cyclohexane_is_grouped_correctly(groupings)
         assert_last_ethanol_is_grouped_correctly(groupings)
 
+    def test_identical_molecule_groups_mixed_topology(self, mixed_topology):
+        assert len(mixed_topology.identical_molecule_groups) == 2
+
     @requires_openeye
     def test_chemical_environments_matches_OE(self):
         """Test Topology.chemical_environment_matches"""