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:** + +*