move network plotting from openfe

environment.yaml

@@ -14,7 +14,6 @@ dependencies:
   - nbval
   #Typing
   - gufe
-  - typing-extensions
   # Sci
   - numpy <2.3  # TypeError: assert_array_almost_equal() got an unexpected keyword argument 'x'
   - scipy
@@ -26,6 +25,7 @@ dependencies:
   - rdkit
   # Vis
   - ipycytoscape
+  - pygraphviz  # TODO: optional dependency?
   # docs
   - myst-parser
   - pydata-sphinx-theme

news/add_plotting.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* Added ``konnektor.visualization.atommapping_network_plotting``, which was previously only available in ``openfe`` (`PR #182 <https://github.com/OpenFreeEnergy/konnektor/pull/182>`_).
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/konnektor/tests/conftest.py

@@ -1,9 +1,15 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/openfe
 import importlib
+from collections.abc import Iterable
+from typing import NamedTuple
 
 import pytest
-from gufe import SmallMoleculeComponent
+from gufe import LigandAtomMapping, LigandNetwork, SmallMoleculeComponent
 from rdkit import Chem
 
+from .network_planners.conf import mol_from_smiles
+
 
 @pytest.fixture(scope="session")
 def atom_mapping_basic_test_files():
@@ -33,3 +39,43 @@ def toluene_vs_others(atom_mapping_basic_test_files):
     others = [v for (k, v) in atom_mapping_basic_test_files.items() if k != central_ligand_name]
     toluene = atom_mapping_basic_test_files[central_ligand_name]
     return toluene, others
+
+
+@pytest.fixture
+def mols():
+    mol1 = SmallMoleculeComponent(mol_from_smiles("CCO"))
+    mol2 = SmallMoleculeComponent(mol_from_smiles("CC"))
+    mol3 = SmallMoleculeComponent(mol_from_smiles("CO"))
+    return mol1, mol2, mol3
+
+
+@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})
+    return edge12, edge23, edge13
+
+
+class _NetworkTestContainer(NamedTuple):
+    """Container to facilitate network testing"""
+
+    network: LigandNetwork
+    nodes: Iterable[SmallMoleculeComponent]
+    edges: Iterable[LigandAtomMapping]
+    n_nodes: int
+    n_edges: int
+
+
+@pytest.fixture
+def simple_network(mols, std_edges):
+    """Network with no edges duplicated and all nodes in edges"""
+    network = LigandNetwork(std_edges)
+    return _NetworkTestContainer(
+        network=network,
+        nodes=mols,
+        edges=std_edges,
+        n_nodes=3,
+        n_edges=3,
+    )

src/konnektor/tests/test_atommapping_network_plotting.py

@@ -0,0 +1,187 @@
+import inspect
+from unittest import mock
+
+import matplotlib
+import pytest
+from matplotlib import pyplot as plt
+
+from konnektor.visualization.atommapping_network_plotting import (
+    AtomMappingNetworkDrawing,
+    LigandNode,
+    plot_atommapping_network,
+)
+
+from .test_network_plotting import mock_event
+
+
+def bound_args(func, args, kwargs):
+    """Return a dictionary mapping parameter name to value.
+
+    Parameters
+    ----------
+    func : Callable
+        this must be inspectable; mocks will require a spec
+    args : List
+        args list
+    kwargs : Dict
+        kwargs Dict
+
+    Returns
+    -------
+    Dict[str, Any] :
+        mapping of string name of function parameter to the value it would
+        be bound to
+    """
+    sig = inspect.Signature.from_callable(func)
+    bound = sig.bind(*args, **kwargs)
+    return bound.arguments
+
+
+@pytest.fixture
+def network_drawing(simple_network):
+    nx_graph = simple_network.network.graph
+    node_dict = {node.smiles: node for node in nx_graph.nodes}
+    positions = {
+        node_dict["CC"]: (0.0, 0.0),
+        node_dict["CO"]: (0.5, 0.0),
+        node_dict["CCO"]: (0.25, 0.25),
+    }
+    graph = AtomMappingNetworkDrawing(nx_graph, positions)
+    graph.ax.set_xlim(0, 1)
+    graph.ax.set_ylim(0, 1)
+    yield graph
+    plt.close(graph.fig)
+
+
+@pytest.fixture
+def default_edge(network_drawing):
+    node_dict = {node.smiles: node for node in network_drawing.graph.nodes}
+    yield network_drawing.edges[node_dict["CC"], node_dict["CO"]]
+
+
+@pytest.fixture
+def default_node(network_drawing):
+    node_dict = {node.smiles: node for node in network_drawing.graph.nodes}
+    yield LigandNode(node_dict["CC"], 0.5, 0.5, 0.1, 0.1)
+
+
+class TestAtomMappingEdge:
+    def test_draw_mapped_molecule(self, default_edge):
+        assert len(default_edge.artist.axes.images) == 0
+        im = default_edge._draw_mapped_molecule(
+            (0.05, 0.45, 0.5, 0.9),
+            default_edge.node_artists[0].node,
+            default_edge.node_artists[1].node,
+            {0: 0},
+        )
+        # maybe add something about im itself? not sure what to test here
+        assert len(default_edge.artist.axes.images) == 1
+        assert default_edge.artist.axes.images[0] == im
+
+    def test_get_image_extents(self, default_edge):
+        left_extent, right_extent = default_edge._get_image_extents()
+        assert left_extent == (0.05, 0.45, 0.5, 0.9)
+        assert right_extent == (0.55, 0.95, 0.5, 0.9)
+
+    def test_select(self, default_edge, network_drawing):
+        assert not default_edge.picked
+        assert len(default_edge.artist.axes.images) == 0
+
+        event = mock_event("mouseup", 0.25, 0.0, network_drawing.fig)
+        default_edge.select(event, network_drawing)
+
+        assert default_edge.picked
+        assert len(default_edge.artist.axes.images) == 2
+
+    @pytest.mark.parametrize(
+        "edge_str,left_right,molA_to_molB",
+        [
+            (("CCO", "CC"), ("CC", "CCO"), {0: 0, 1: 1}),
+            (("CC", "CO"), ("CC", "CO"), {0: 0}),
+            (("CCO", "CO"), ("CCO", "CO"), {0: 0, 2: 1}),
+        ],
+    )
+    def test_select_mock_drawing(self, edge_str, left_right, molA_to_molB, network_drawing):
+        # this tests that we call _draw_mapped_molecule with the correct
+        # kwargs -- in particular, it ensures that we get the left and right
+        # molecules correctly
+        node_dict = {node.smiles: node for node in network_drawing.graph.nodes}
+        edge_tuple = tuple(node_dict[node] for node in edge_str)
+        edge = network_drawing.edges[edge_tuple]
+        left, right = [network_drawing.nodes[node_dict[node]] for node in left_right]
+        # ensure that we have them labelled correctly
+        assert left.xy[0] < right.xy[0]
+        func = edge._draw_mapped_molecule  # save for bound_args
+        edge._draw_mapped_molecule = mock.Mock()
+
+        event = mock_event("mouseup", 0.25, 0.0, network_drawing.fig)
+        edge.select(event, network_drawing)
+
+        arg_dicts = [
+            bound_args(func, call.args, call.kwargs)
+            for call in edge._draw_mapped_molecule.mock_calls
+        ]
+        expected_left = {
+            "extent": (0.05, 0.45, 0.5, 0.9),
+            "molA": left.node,
+            "molB": right.node,
+            "molA_to_molB": molA_to_molB,
+        }
+        expected_right = {
+            "extent": (0.55, 0.95, 0.5, 0.9),
+            "molA": right.node,
+            "molB": left.node,
+            "molA_to_molB": {v: k for k, v in molA_to_molB.items()},
+        }
+        assert len(arg_dicts) == 2
+        assert expected_left in arg_dicts
+        assert expected_right in arg_dicts
+
+    def test_unselect(self, default_edge, network_drawing):
+        # start by selecting; hard to be sure we mocked all the side effects
+        # of select
+        event = mock_event("mouseup", 0.25, 0.0, network_drawing.fig)
+        default_edge.select(event, network_drawing)
+        assert default_edge.picked
+        assert len(default_edge.artist.axes.images) == 2
+        assert default_edge.right_image is not None
+        assert default_edge.left_image is not None
+
+        default_edge.unselect()
+
+        assert not default_edge.picked
+        assert len(default_edge.artist.axes.images) == 0
+        assert default_edge.right_image is None
+        assert default_edge.left_image is None
+
+
+class TestLigandNode:
+    def setup_method(self):
+        self.fig, self.ax = plt.subplots()
+
+    def teardown_method(self):
+        plt.close(self.fig)
+
+    def test_register_artist(self, default_node):
+        assert len(self.ax.texts) == 0
+        default_node.register_artist(self.ax)
+        assert len(self.ax.texts) == 1
+        assert self.ax.texts[0] == default_node.artist
+
+    def test_extent(self, default_node):
+        default_node.register_artist(self.ax)
+        xmin, xmax, ymin, ymax = default_node.extent
+        assert xmin == pytest.approx(0.5)
+        assert ymin == pytest.approx(0.5)
+        # can't do anything about upper bounds
+
+    def test_xy(self, default_node):
+        # default_node.register_artist(self.ax)
+        x, y = default_node.xy
+        assert x == pytest.approx(0.5)
+        assert y == pytest.approx(0.5)
+
+
+def test_plot_atommapping_network(simple_network):
+    fig = plot_atommapping_network(simple_network.network)
+    assert isinstance(fig, matplotlib.figure.Figure)