Minimal patch for JSON deserialization of _SimpleMolecules (#1958)

* Minimal patch for JSON deserialization of `_SimpleMolecule`s * update releasenotes --------- Co-authored-by: Jeff Wagner <[email protected]>
openforcefield · Nov 7, 2024 · 9c122a2 · 9c122a2
1 parent 7f8652c
commit 9c122a2
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/docs/releasehistory.md b/docs/releasehistory.md
@@ -18,6 +18,8 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
 
 ### New features
 
+- [PR #1958](https://github.com/openforcefield/openff-toolkit/pull/1958): Allows serialization of `_SimpleMolecule`s (relevant to serialization of `Topology` objects created by `Interchange.from_openmm`)
+
 ### Improved documentation and warnings
 
 - [PR #1957](https://github.com/openforcefield/openff-toolkit/pull/1957): Fixed a misleading typo in an error message.

diff --git a/openff/toolkit/_tests/test_topology.py b/openff/toolkit/_tests/test_topology.py
@@ -2045,6 +2045,17 @@ def test_roundtrip(self, oleic_acid, with_conformers, n_molecules, format):
         assert next(iter(roundtrip.molecules)).n_conformers == n_conformers
 
 
+    @pytest.mark.parametrize("format", ["dict", "json"])
+    def test_roundtrip_simple_molecules(self, mixed_topology, format):
+        if format == "dict":
+            roundtrip = Topology.from_dict(mixed_topology.to_dict())
+        elif format == "json":
+            roundtrip = Topology.from_json(mixed_topology.to_json())
+
+        assert roundtrip.n_molecules == mixed_topology.n_molecules
+        assert roundtrip.n_atoms == mixed_topology.n_atoms
+
+
 @pytest.mark.parametrize(
     ("n_degrees", "num_pairs"),
     [

diff --git a/openff/toolkit/topology/topology.py b/openff/toolkit/topology/topology.py
@@ -1302,8 +1302,18 @@ def _initialize_from_dict(self, topology_dict):
             self.box_vectors = Quantity(box_vectors_unitless, box_vectors_unit)
 
         for molecule_dict in topology_dict["molecules"]:
-            new_mol = Molecule.from_dict(molecule_dict)
+            # the serialized representation doesn't store which molecule model
+            # each of these dicts is meant to be deserialized to. Usually it'll
+            # be Molecule, so try that first. If it's supposed to be a
+            # _SimpleMolecule, it'll KeyError because Molecule has more fields.
+            try:
+                new_mol = Molecule.from_dict(molecule_dict)
+            except KeyError:
+                # masks possible other ways KeyErrors could pop up
+                new_mol = _SimpleMolecule.from_dict(molecule_dict)
+
             self._add_molecule_keep_cache(new_mol)
+
         self._invalidate_cached_properties()
 
     @staticmethod