Merge pull request #369 from OpenFreeEnergy/graphml_annotations

Save edge annotations in `LigandNetwork.to_graphml`
OpenFreeEnergy · Nov 5, 2024 · f7eb1a7 · f7eb1a7
2 parents 7142e80 + 78d557d
commit f7eb1a7
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 11 deletions.
diff --git a/gufe/ligandnetwork.py b/gufe/ligandnetwork.py
@@ -10,7 +10,7 @@
 
 from gufe import SmallMoleculeComponent
 from .mapping import LigandAtomMapping
-from .tokenization import GufeTokenizable
+from .tokenization import GufeTokenizable, JSON_HANDLER
 
 
 class LigandNetwork(GufeTokenizable):
@@ -97,7 +97,8 @@ def _serializable_graph(self) -> nx.Graph:
             (
                 mol_to_label[edge.componentA],
                 mol_to_label[edge.componentB],
-                json.dumps(list(edge.componentA_to_componentB.items()))
+                json.dumps(list(edge.componentA_to_componentB.items())),
+                json.dumps(edge.annotations, cls=JSON_HANDLER.encoder),
             )
             for edge in self.edges
         ])
@@ -109,8 +110,8 @@ def _serializable_graph(self) -> nx.Graph:
                                         moldict=json.dumps(mol.to_dict(),
                                                            sort_keys=True))
 
-        for molA, molB, mapping in edge_data:
-            serializable_graph.add_edge(molA, molB, mapping=mapping)
+        for molA, molB, mapping, annotation in edge_data:
+            serializable_graph.add_edge(molA, molB, mapping=mapping, annotations=annotation)
 
         return serializable_graph
 
@@ -126,8 +127,10 @@ def _from_serializable_graph(cls, graph: nx.Graph):
         edges = [
             LigandAtomMapping(componentA=label_to_mol[node1],
                               componentB=label_to_mol[node2],
-                              componentA_to_componentB=dict(json.loads(mapping)))
-            for node1, node2, mapping in graph.edges(data='mapping')
+                              componentA_to_componentB=dict(json.loads(edge_data["mapping"])),
+                              annotations=json.loads(edge_data.get("annotations", 'null'), cls=JSON_HANDLER.decoder) # work around old graphml files with missing edge annotations
+                              )
+            for node1, node2, edge_data in graph.edges(data=True)
         ]
 
         return cls(edges=edges, nodes=label_to_mol.values())

diff --git a/gufe/mapping/ligandatommapping.py b/gufe/mapping/ligandatommapping.py
@@ -96,7 +96,7 @@ def componentB_unique(self):
     def annotations(self):
         """Any extra metadata, for example the score of a mapping"""
         # return a copy (including copy of nested)
-        return json.loads(json.dumps(self._annotations))
+        return json.loads(json.dumps(self._annotations, cls=JSON_HANDLER.encoder), cls=JSON_HANDLER.decoder)
 
     def _to_dict(self):
         """Serialize to dict"""

diff --git a/gufe/tests/data/ligand_network.graphml b/gufe/tests/data/ligand_network.graphml
@@ -1,4 +1,5 @@
 <graphml xmlns="http://graphml.graphdrawing.org/xmlns" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://graphml.graphdrawing.org/xmlns http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd">
+  <key id="d2" for="edge" attr.name="annotations" attr.type="string" />
   <key id="d1" for="edge" attr.name="mapping" attr.type="string" />
   <key id="d0" for="node" attr.name="moldict" attr.type="string" />
   <graph edgedefault="directed">
@@ -13,12 +14,15 @@
     </node>
     <edge source="mol0" target="mol2" id="0">
       <data key="d1">[[0, 0]]</data>
+      <data key="d2">{"score": 1.0}</data>
     </edge>
     <edge source="mol1" target="mol0" id="0">
       <data key="d1">[[0, 0], [1, 1]]</data>
+      <data key="d2">{"score": 0.0, "length": {"magnitude": 1.0, "unit": "angstrom", ":is_custom:": true, "pint_unit_registry": "openff_units"}}</data>
     </edge>
     <edge source="mol1" target="mol2" id="0">
       <data key="d1">[[0, 0], [2, 1]]</data>
+      <data key="d2">{"score": 0.5, "time": {"magnitude": 2.0, "unit": "second", ":is_custom:": true, "pint_unit_registry": "openff_units"}}</data>
     </edge>
   </graph>
 </graphml>
diff --git a/gufe/tests/test_ligand_network.py b/gufe/tests/test_ligand_network.py
@@ -7,6 +7,8 @@
 from gufe.tests.test_protocol import DummyProtocol
 from gufe import SmallMoleculeComponent, LigandNetwork, LigandAtomMapping
 
+from openff.units import unit
+
 from rdkit import Chem
 
 from networkx import NetworkXError
@@ -47,9 +49,9 @@ def mols():
 @pytest.fixture
 def std_edges(mols):
     mol1, mol2, mol3 = mols
-    edge12 = LigandAtomMapping(mol1, mol2, {0: 0, 1: 1})
-    edge23 = LigandAtomMapping(mol2, mol3, {0: 0})
-    edge13 = LigandAtomMapping(mol1, mol3, {0: 0, 2: 1})
+    edge12 = LigandAtomMapping(mol1, mol2, {0: 0, 1: 1}, {"score": 0.0, "length": 1.0 * unit.angstrom})
+    edge23 = LigandAtomMapping(mol2, mol3, {0: 0}, {"score": 1.0})
+    edge13 = LigandAtomMapping(mol1, mol3, {0: 0, 2: 1}, {"score": 0.5, "time": 2.0 * unit.second})
     return edge12, edge23, edge13
 
 
@@ -250,7 +252,7 @@ def test_enlarge_graph_add_duplicate_edge(self, mols, simple_network):
         # Adding a duplicate of an existing edge should create a new network
         # with the same edges and nodes as the previous one.
         mol1, _, mol3 = mols
-        duplicate = LigandAtomMapping(mol1, mol3, {0: 0, 2: 1})
+        duplicate = LigandAtomMapping(mol1, mol3, {0: 0, 2: 1}, {"score": 0.5, "time": 2.0 * unit.second})
         network = simple_network.network
 
         existing = network.edges

diff --git a/news/graphml_annotations.rst b/news/graphml_annotations.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* Fixed a bug where edge annotations were lost when converting a ``LigandNetwork`` to graphml, all JSON codec types are now supported.
+
+**Security:**
+
+* <news item>