diff --git a/gufe/ligandnetwork.py b/gufe/ligandnetwork.py
index 177aa61e..e92fa812 100644
--- 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
index 9a5b198b..cc8ada71 100644
--- 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
index 07adb79d..31fe9e56 100644
--- a/gufe/tests/data/ligand_network.graphml
+++ b/gufe/tests/data/ligand_network.graphml
@@ -1,4 +1,5 @@
+
@@ -13,12 +14,15 @@
[[0, 0]]
+ {"score": 1.0}
[[0, 0], [1, 1]]
+ {"score": 0.0, "length": {"magnitude": 1.0, "unit": "angstrom", ":is_custom:": true, "pint_unit_registry": "openff_units"}}
[[0, 0], [2, 1]]
+ {"score": 0.5, "time": {"magnitude": 2.0, "unit": "second", ":is_custom:": true, "pint_unit_registry": "openff_units"}}
\ No newline at end of file
diff --git a/gufe/tests/test_ligand_network.py b/gufe/tests/test_ligand_network.py
index ff6f4d48..9b3be8c4 100644
--- 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
new file mode 100644
index 00000000..b142b305
--- /dev/null
+++ b/news/graphml_annotations.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+*
+
+**Changed:**
+
+*
+
+**Deprecated:**
+
+*
+
+**Removed:**
+
+*
+
+**Fixed:**
+
+* Fixed a bug where edge annotations were lost when converting a ``LigandNetwork`` to graphml, all JSON codec types are now supported.
+
+**Security:**
+
+*