Merge branch 'main' of github.com:OpenFreeEnergy/gufe into shared-obj…

…ect-v2

MANIFEST.in

@@ -4,3 +4,4 @@ recursive-include gufe/tests/data/ *.sdf
 recursive-include gufe/tests/data/ *.cif
 recursive-include gufe/tests/data/ *.mol2
 recursive-include gufe/tests/data/ *.json
+recursive-include gufe/tests/data/ *.graphml

environment.yml

@@ -2,7 +2,6 @@ name: gufe
 channels:
   - jaimergp/label/unsupported-cudatoolkit-shim
   - conda-forge
-  - openeye
 dependencies:
   - coverage
   - networkx

gufe/__init__.py

@@ -37,5 +37,6 @@
 from .transformations import Transformation, NonTransformation
 
 from .network import AlchemicalNetwork
+from .ligandnetwork import LigandNetwork
 
 __version__ = version("gufe")

gufe/ligandnetwork.py

@@ -0,0 +1,180 @@
+# This code is part of gufe and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/gufe
+from __future__ import annotations
+
+import json
+import networkx as nx
+from typing import FrozenSet, Iterable, Optional
+
+from gufe import SmallMoleculeComponent
+from .mapping import LigandAtomMapping
+from .tokenization import GufeTokenizable
+
+
+class LigandNetwork(GufeTokenizable):
+    """A directed graph connecting many ligands according to their atom mapping
+
+    Parameters
+    ----------
+    edges : Iterable[LigandAtomMapping]
+        edges for this network
+    nodes : Iterable[SmallMoleculeComponent]
+        nodes for this network
+    """
+    def __init__(
+        self,
+        edges: Iterable[LigandAtomMapping],
+        nodes: Optional[Iterable[SmallMoleculeComponent]] = None
+    ):
+        if nodes is None:
+            nodes = []
+
+        self._edges = frozenset(edges)
+        edge_nodes = set.union(*[{edge.componentA, edge.componentB} for edge in edges])
+        self._nodes = frozenset(edge_nodes) | frozenset(nodes)
+        self._graph = None
+
+    @classmethod
+    def _defaults(cls):
+        return {}
+
+    def _to_dict(self) -> dict:
+        return {'graphml': self.to_graphml()}
+
+    @classmethod
+    def _from_dict(cls, dct: dict):
+        return cls.from_graphml(dct['graphml'])
+
+    @property
+    def graph(self) -> nx.Graph:
+        """NetworkX graph for this network"""
+        if self._graph is None:
+            graph = nx.MultiDiGraph()
+            # set iterator order depends on PYTHONHASHSEED, sorting ensures
+            # reproducibility
+            for node in sorted(self._nodes):
+                graph.add_node(node)
+            for edge in sorted(self._edges):
+                graph.add_edge(edge.componentA, edge.componentB, object=edge,
+                               **edge.annotations)
+
+            self._graph = nx.freeze(graph)
+
+        return self._graph
+
+    @property
+    def edges(self) -> FrozenSet[LigandAtomMapping]:
+        """A read-only view of the edges of the Network"""
+        return self._edges
+
+    @property
+    def nodes(self) -> FrozenSet[SmallMoleculeComponent]:
+        """A read-only view of the nodes of the Network"""
+        return self._nodes
+
+    def _serializable_graph(self) -> nx.Graph:
+        """
+        Create NetworkX graph with serializable attribute representations.
+
+        This enables us to use easily use different serialization
+        approaches.
+        """
+        # sorting ensures that we always preserve order in files, so two
+        # identical networks will show no changes if you diff their
+        # serialized versions
+        sorted_nodes = sorted(self.nodes, key=lambda m: (m.smiles, m.name))
+        mol_to_label = {mol: f"mol{num}"
+                        for num, mol in enumerate(sorted_nodes)}
+
+        edge_data = sorted([
+            (
+                mol_to_label[edge.componentA],
+                mol_to_label[edge.componentB],
+                json.dumps(list(edge.componentA_to_componentB.items()))
+            )
+            for edge in self.edges
+        ])
+
+        # from here, we just build the graph
+        serializable_graph = nx.MultiDiGraph()
+        for mol, label in mol_to_label.items():
+            serializable_graph.add_node(label,
+                                        moldict=json.dumps(mol.to_dict(),
+                                                           sort_keys=True))
+
+        for molA, molB, mapping in edge_data:
+            serializable_graph.add_edge(molA, molB, mapping=mapping)
+
+        return serializable_graph
+
+    @classmethod
+    def _from_serializable_graph(cls, graph: nx.Graph):
+        """Create network from NetworkX graph with serializable attributes.
+
+        This is the inverse of ``_serializable_graph``.
+        """
+        label_to_mol = {node: SmallMoleculeComponent.from_dict(json.loads(d))
+                        for node, d in graph.nodes(data='moldict')}
+
+        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')
+        ]
+
+        return cls(edges=edges, nodes=label_to_mol.values())
+
+    def to_graphml(self) -> str:
+        """Return the GraphML string representing this ``Network``.
+
+        This is the primary serialization mechanism for this class.
+
+        Returns
+        -------
+        str :
+            string representing this network in GraphML format
+        """
+        return "\n".join(nx.generate_graphml(self._serializable_graph()))
+
+    @classmethod
+    def from_graphml(cls, graphml_str: str):
+        """Create ``Network`` from GraphML string.
+
+        This is the primary deserialization mechanism for this class.
+
+        Parameters
+        ----------
+        graphml_str : str
+            GraphML string representation of a :class:`.Network`
+
+        Returns
+        -------
+        :class:`.Network`:
+            new network from the GraphML
+        """
+        return cls._from_serializable_graph(nx.parse_graphml(graphml_str))
+
+    def enlarge_graph(self, *, edges=None, nodes=None) -> LigandNetwork:
+        """
+        Create a new network with the given edges and nodes added
+
+        Parameters
+        ----------
+        edges : Iterable[:class:`.LigandAtomMapping`]
+            edges to append to this network
+        nodes : Iterable[:class:`.SmallMoleculeComponent`]
+            nodes to append to this network
+
+        Returns
+        -------
+        :class:`.Network :
+            a new network adding the given edges and nodes to this network
+        """
+        if edges is None:
+            edges = set([])
+
+        if nodes is None:
+            nodes = set([])
+
+        return LigandNetwork(self.edges | set(edges), self.nodes | set(nodes))

gufe/network.py

@@ -1,7 +1,7 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/gufe
 
-from typing import FrozenSet, Iterable, Optional, Tuple
+from typing import Iterable, Optional
 
 import networkx as nx
 from .tokenization import GufeTokenizable
@@ -18,10 +18,10 @@ class AlchemicalNetwork(GufeTokenizable):
 
     Attributes
     ----------
-    edges : FrozenSet[Transformation]
+    edges : frozenset[Transformation]
         The edges of the network, given as a ``frozenset`` of
         :class:`.Transformation`\ s.
-    nodes : FrozenSet[ChemicalSystem]
+    nodes : frozenset[ChemicalSystem]
         The nodes of the network, given as a ``frozenset`` of
         :class:`.ChemicalSystem` \ s.
     name : Optional identifier for the network.
@@ -33,8 +33,8 @@ def __init__(
         nodes: Optional[Iterable[ChemicalSystem]] = None,
         name: Optional[str] = None,
     ):
-        self._edges: FrozenSet[Transformation] = frozenset(edges) if edges else frozenset()
-        self._nodes: FrozenSet[ChemicalSystem]
+        self._edges: frozenset[Transformation] = frozenset(edges) if edges else frozenset()
+        self._nodes: frozenset[ChemicalSystem]
 
         self._name = name
 
@@ -73,11 +73,11 @@ def graph(self):
         return self._graph.copy(as_view=True)
 
     @property
-    def edges(self) -> FrozenSet[Transformation]:
+    def edges(self) -> frozenset[Transformation]:
         return self._edges
 
     @property
-    def nodes(self) -> FrozenSet[ChemicalSystem]:
+    def nodes(self) -> frozenset[ChemicalSystem]:
         return self._nodes
 
     @property

gufe/protocols/protocol.py

@@ -28,7 +28,6 @@ class ProtocolResult(GufeTokenizable):
     The following methods should be implemented in any subclass:
     - `get_estimate`
     - `get_uncertainty`
-    - `get_rate_of_convergence`
 
     Attributes
     ----------
@@ -65,10 +64,6 @@ def get_estimate(self) -> Quantity:
     def get_uncertainty(self) -> Quantity:
         ...
 
-    @abc.abstractmethod
-    def get_rate_of_convergence(self):
-        ...
-
 
 class Protocol(GufeTokenizable):
     """A protocol that implements an alchemical transformation.

gufe/tests/conftest.py

@@ -21,8 +21,6 @@
     HAS_INTERNET = True
 
 
-## helper functions
-
 class URLFileLike:
     def __init__(self, url, encoding='utf-8'):
         self.url = url
@@ -45,9 +43,6 @@ def get_test_filename(filename):
         return str(file)
 
 
-## PDBs for input/output testing
-
-
 _benchmark_pdb_names = [
         "cmet_protein",
         "hif2a_protein",
@@ -79,8 +74,6 @@ def get_test_filename(filename):
 ALL_PDB_LOADERS = dict(**PDB_BENCHMARK_LOADERS, **PDB_FILE_LOADERS)
 
 
-## data file paths
-
 @pytest.fixture
 def ethane_sdf():
     with importlib.resources.path("gufe.tests.data", "ethane.sdf") as f:
@@ -145,7 +138,6 @@ def PDBx_181L_openMMClean_path():
                                   '181l_openmmClean.cif') as f:
         yield str(f)
 
-## RDKit molecules
 
 @pytest.fixture(scope='session')
 def benzene_modifications():
@@ -164,9 +156,6 @@ def benzene_transforms(benzene_modifications):
             for k, v in benzene_modifications.items()}
 
 
-## Components
-
-
 @pytest.fixture
 def benzene(benzene_modifications):
     return gufe.SmallMoleculeComponent(benzene_modifications["benzene"])
@@ -229,9 +218,6 @@ def phenol_ligand_comp(benzene_modifications):
     yield gufe.SmallMoleculeComponent.from_rdkit(benzene_modifications["phenol"])
 
 
-## ChemicalSystems
-
-
 @pytest.fixture
 def solvated_complex(prot_comp, solv_comp, toluene_ligand_comp):
     return gufe.ChemicalSystem(
@@ -253,9 +239,6 @@ def vacuum_ligand(toluene_ligand_comp):
     )
 
 
-## Transformations
-
-
 @pytest.fixture
 def absolute_transformation(solvated_ligand, solvated_complex):
     return gufe.Transformation(
@@ -274,9 +257,6 @@ def complex_equilibrium(solvated_complex):
     )
 
 
-## Alchemical Networks
-
-
 @pytest.fixture
 def benzene_variants_star_map(
     benzene,

gufe/tests/data/ligand_network.graphml

@@ -0,0 +1,24 @@
+<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="d1" for="edge" attr.name="mapping" attr.type="string" />
+  <key id="d0" for="node" attr.name="moldict" attr.type="string" />
+  <graph edgedefault="directed">
+    <node id="mol0">
+      <data key="d0">{"__module__": "gufe.components.smallmoleculecomponent", "__qualname__": "SmallMoleculeComponent", "atoms": [[6, 0, 0, false, 0, 0, {}], [6, 0, 0, false, 0, 0, {}]], "bonds": [[0, 1, 1, 0, {}]], "conformer": ["\u0093NUMPY\u0001\u0000v\u0000{'descr': '&lt;f8', 'fortran_order': False, 'shape': (2, 3), }                                                          \n\u0000\u0000\u0000\u0000\u0000\u0000\u00e8\u00bf\u0000\u0000\u0000\u0000\u0000\u0000\u0090&lt;\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u00e8?\u0000\u0000\u0000\u0000\u0000\u0000\u0090\u00bc\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000", {}], "molprops": {"ofe-name": ""}}</data>
+    </node>
+    <node id="mol1">
+      <data key="d0">{"__module__": "gufe.components.smallmoleculecomponent", "__qualname__": "SmallMoleculeComponent", "atoms": [[6, 0, 0, false, 0, 0, {}], [6, 0, 0, false, 0, 0, {}], [8, 0, 0, false, 0, 0, {}]], "bonds": [[0, 1, 1, 0, {}], [1, 2, 1, 0, {}]], "conformer": ["\u0093NUMPY\u0001\u0000v\u0000{'descr': '&lt;f8', 'fortran_order': False, 'shape': (3, 3), }                                                          \n\u00809B.\u00dc\u00c8\u00f4\u00bf\u00f5\u00ff\u00ff\u00ff\u00ff\u00ff\u00cf\u00bf\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0001\u0000\u0000\u0000\u0000\u0000\u00e0?\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u00809B.\u00dc\u00c8\u00f4?\u0006\u0000\u0000\u0000\u0000\u0000\u00d0\u00bf\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000", {}], "molprops": {"ofe-name": ""}}</data>
+    </node>
+    <node id="mol2">
+      <data key="d0">{"__module__": "gufe.components.smallmoleculecomponent", "__qualname__": "SmallMoleculeComponent", "atoms": [[6, 0, 0, false, 0, 0, {}], [8, 0, 0, false, 0, 0, {}]], "bonds": [[0, 1, 1, 0, {}]], "conformer": ["\u0093NUMPY\u0001\u0000v\u0000{'descr': '&lt;f8', 'fortran_order': False, 'shape': (2, 3), }                                                          \n\u0000\u0000\u0000\u0000\u0000\u0000\u00e8\u00bf\u0000\u0000\u0000\u0000\u0000\u0000\u0090&lt;\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u00e8?\u0000\u0000\u0000\u0000\u0000\u0000\u0090\u00bc\u0000\u0000\u0000\u0000\u0000\u0000\u0000\u0000", {}], "molprops": {"ofe-name": ""}}</data>
+    </node>
+    <edge source="mol0" target="mol2" id="0">
+      <data key="d1">[[0, 0]]</data>
+    </edge>
+    <edge source="mol1" target="mol0" id="0">
+      <data key="d1">[[0, 0], [1, 1]]</data>
+    </edge>
+    <edge source="mol1" target="mol2" id="0">
+      <data key="d1">[[0, 0], [2, 1]]</data>
+    </edge>
+  </graph>
+</graphml>

gufe/tests/dev/serialization_test_templates.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+# creates ligand_network.graphml
+
+
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from gufe import SmallMoleculeComponent, LigandNetwork, LigandAtomMapping
+
+
+def mol_from_smiles(smiles: str) -> Chem.Mol:
+    m = Chem.MolFromSmiles(smiles)
+    AllChem.Compute2DCoords(m)
+
+    return m
+
+
+# network_template.graphml
+mol1 = SmallMoleculeComponent(mol_from_smiles("CCO"))
+mol2 = SmallMoleculeComponent(mol_from_smiles("CC"))
+mol3 = SmallMoleculeComponent(mol_from_smiles("CO"))
+
+edge12 = LigandAtomMapping(mol1, mol2, {0: 0, 1: 1})
+edge23 = LigandAtomMapping(mol2, mol3, {0: 0})
+edge13 = LigandAtomMapping(mol1, mol3, {0: 0, 2: 1})
+
+network = LigandNetwork([edge12, edge23, edge13])
+
+with open("ligand_network.graphml", "w") as fn:
+    fn.write(network.to_graphml())