fix test

Author: TieuLongPhan

.github/workflows/publish-package.yml

@@ -3,7 +3,7 @@ name: PyPI publish
 on:
   push:
     branches:
-      - mech
+      - beta
   release:
     types:
       - published

Test/Vis/test_mtg_drawer.py

@@ -0,0 +1,106 @@
+import unittest
+
+import matplotlib
+import networkx as nx
+
+matplotlib.use("Agg")
+
+from synkit.Graph.ITS.its_construction import ITSConstruction  # noqa: E402
+from synkit.Graph.MTG.mtg import MTG  # noqa: E402
+from synkit.Vis.mtg_drawer import draw_mtg_graph, draw_mtg_steps  # noqa: E402
+
+
+class TestMTGDrawer(unittest.TestCase):
+    @staticmethod
+    def _atom(element, *, hcount=0, charge=0, lone_pairs=0, radical=0):
+        return {
+            "element": element,
+            "aromatic": False,
+            "hcount": hcount,
+            "charge": charge,
+            "lone_pairs": lone_pairs,
+            "radical": radical,
+            "valence_electrons": {"H": 1, "C": 4, "N": 5, "O": 6, "Cl": 7}[element],
+        }
+
+    @staticmethod
+    def _bond(graph, u, v, sigma=1.0, pi=0.0):
+        graph.add_edge(
+            u,
+            v,
+            order=sigma + pi,
+            kekule_order=sigma + pi,
+            sigma_order=sigma,
+            pi_order=pi,
+        )
+
+    def _graph(self, nodes, edges):
+        graph = nx.Graph()
+        for node, attrs in nodes.items():
+            graph.add_node(node, **attrs)
+        for edge in edges:
+            self._bond(graph, *edge)
+        return graph
+
+    def _mtg(self):
+        g0 = self._graph(
+            {
+                1: self._atom("C", hcount=3),
+                2: self._atom("Cl", lone_pairs=3),
+            },
+            [(1, 2, 1.0, 0.0)],
+        )
+        g1 = self._graph(
+            {
+                1: self._atom("C", hcount=3, radical=1),
+                2: self._atom("Cl", radical=1, lone_pairs=3),
+            },
+            [],
+        )
+        g2 = self._graph(
+            {
+                1: self._atom("C", hcount=3),
+                2: self._atom("Cl", lone_pairs=3),
+            },
+            [(1, 2, 1.0, 0.0)],
+        )
+        return MTG(
+            [ITSConstruction.construct(g0, g1), ITSConstruction.construct(g1, g2)],
+            mappings=[{1: 1, 2: 2}],
+        )
+
+    def test_draw_mtg_graph_accepts_mtg_object(self):
+        mtg = self._mtg()
+
+        fig, ax = draw_mtg_graph(mtg, title="radical rebound")
+
+        self.assertIs(fig, ax.figure)
+        self.assertEqual(ax.get_title(), "radical rebound")
+
+    def test_draw_mtg_graph_accepts_raw_graph_without_mutation(self):
+        graph = self._mtg().get_mtg()
+        before_nodes = dict(graph.nodes(data=True))
+        before_edges = list(graph.edges(data=True))
+
+        fig, ax = draw_mtg_graph(graph)
+
+        self.assertIs(fig, ax.figure)
+        self.assertEqual(dict(graph.nodes(data=True)), before_nodes)
+        self.assertEqual(list(graph.edges(data=True)), before_edges)
+
+    def test_draw_mtg_steps_draws_ordered_its_panels_and_composed_panel(self):
+        mtg = self._mtg()
+
+        fig, axes = draw_mtg_steps(mtg, include_composed=True, show_edge_labels=True)
+
+        self.assertIs(fig, axes[0].figure)
+        self.assertEqual(len(axes), 3)
+        self.assertEqual([ax.get_title() for ax in axes], ["Step 1", "Step 2", "Composed"])
+
+    def test_draw_mtg_steps_validates_indices(self):
+        with self.assertRaises(IndexError):
+            draw_mtg_steps(self._mtg(), steps=[2])
+
+
+if __name__ == "__main__":
+    unittest.main()

doc/api/vis.rst

@@ -19,6 +19,10 @@ Modern molecule/reaction/ITS renderers
    :members:
    :show-inheritance:
 
+.. automodule:: synkit.Vis.mtg_drawer
+   :members:
+   :show-inheritance:
+
 Diagnostic adapter layer
 ------------------------
 

doc/vis.rst

@@ -168,6 +168,38 @@ Use ``projection=True`` when you need to inspect how an ITS decomposes back into
 left and right molecular graphs. Use the default ITS-only view for reports and
 notebooks.
 
+MTG Timelines
+-------------
+
+Compact MTG visualization has two complementary views:
+
+* ``draw_mtg_graph`` shows the fused MTG as a timeline graph;
+* ``draw_mtg_steps`` reconstructs ordered ITS steps and draws each step with
+  the ITS renderer.
+
+.. code-block:: python
+
+   from synkit.Graph.MTG.mtg import MTG
+   from synkit.Vis import draw_mtg_graph, draw_mtg_steps
+
+   mtg = MTG([step_1_its, step_2_its])
+
+   fig, ax = draw_mtg_graph(
+       mtg,
+       title="MTG timeline",
+       mode="timeline",
+   )
+
+   fig, axes = draw_mtg_steps(
+       mtg,
+       include_composed=True,
+       show_edge_labels=True,
+   )
+
+Use the timeline graph to see transient bonds and electron-state paths across
+the mechanism. Use the step panels when you need to check each reconstructed
+ITS independently.
+
 Diagnostic Graph View
 ---------------------
 
@@ -198,4 +230,3 @@ The older visualization classes are still exported for compatibility:
 New code should use ``draw_molecule_graph``, ``draw_reaction_graph``, and
 ``draw_its_from_rsmi`` unless a legacy workflow specifically depends on the
 class-based API.
-

requirements.txt

@@ -2,6 +2,7 @@ scikit-learn>=1.4.0
 seaborn>=0.13.2
 rdkit>=2025.3.1
 pandas>=2.2.0
+networkx>=3.3
 requests>=2.32.3
 numpy>=2.2.0
 regex>=2024.11.6

synkit/Vis/__init__.py

@@ -21,6 +21,7 @@
     find_reaction_highlights,
 )
 from .its_drawer import draw_its_from_rsmi, draw_its_graph, draw_its_only
+from .mtg_drawer import draw_mtg_graph, draw_mtg_steps
 
 __all__ = [
     "GraphVisualizer",
@@ -44,4 +45,6 @@
     "draw_its_from_rsmi",
     "draw_its_graph",
     "draw_its_only",
+    "draw_mtg_graph",
+    "draw_mtg_steps",
 ]

synkit/Vis/mtg_drawer.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+"""MTG visualization helpers.
+
+The compact MTG view is a timeline diagnostic. Step panels reuse the molecule-
+like ITS renderer so each reconstructed ITS step is inspected with the same
+visual language as normal tuple ITS drawings.
+"""
+
+from typing import Any, Iterable, Optional
+
+import matplotlib.pyplot as plt
+import networkx as nx
+
+from synkit.Vis.its_drawer import draw_its_only
+from synkit.Vis.visual_drawer import draw_graph
+
+
+def draw_mtg_graph(
+    mtg: Any,
+    *,
+    ax: Optional[plt.Axes] = None,
+    title: Optional[str] = None,
+    mode: str = "timeline",
+    layout: str = "kamada_kawai",
+    show_atom_map: bool = True,
+    show_edge_labels: bool = True,
+    show_node_badges: bool = True,
+) -> tuple[plt.Figure, plt.Axes]:
+    """Draw a compact MTG timeline graph.
+
+    ``mtg`` may be a :class:`synkit.Graph.MTG.mtg.MTG` instance or a raw
+    compact MTG ``networkx.Graph`` from ``MTG.get_mtg()``.
+
+    :param mtg: MTG object or compact MTG graph.
+    :type mtg: Any
+    :param ax: Optional Matplotlib axes.
+    :type ax: Optional[plt.Axes]
+    :param title: Optional title.
+    :type title: Optional[str]
+    :param mode: Visual adapter mode. ``"timeline"`` is the recommended MTG
+        view; ``"sigma_pi"`` gives a shorter electron-bond diagnostic.
+    :type mode: str
+    :param layout: NetworkX layout name passed to ``draw_graph``.
+    :type layout: str
+    :returns: ``(figure, axes)``.
+    :rtype: tuple[plt.Figure, plt.Axes]
+    """
+
+    graph = _as_mtg_graph(mtg)
+    return draw_graph(
+        graph,
+        ax=ax,
+        mode=mode,
+        title=title or "MTG timeline",
+        show_atom_map=show_atom_map,
+        layout=layout,
+        show_edge_labels=show_edge_labels,
+        show_node_badges=show_node_badges,
+    )
+
+
+def draw_mtg_steps(
+    mtg: Any,
+    *,
+    steps: Optional[Iterable[int]] = None,
+    include_composed: bool = False,
+    title: Optional[str] = None,
+    max_columns: int = 3,
+    show_atom_map: bool = True,
+    label_mode: str = "hetero",
+    edge_label_mode: str = "kekule",
+    show_edge_labels: bool = False,
+    show_electron_labels: bool = False,
+    electron_label_mode: str = "charge",
+) -> tuple[plt.Figure, list[plt.Axes]]:
+    """Draw reconstructed MTG ITS steps as ordered panels.
+
+    :param mtg: MTG object exposing ``get_its_steps``.
+    :type mtg: Any
+    :param steps: Optional zero-based step indices to draw.
+    :type steps: Optional[Iterable[int]]
+    :param include_composed: Append the composed outer-state ITS panel.
+    :type include_composed: bool
+    :param title: Optional figure title.
+    :type title: Optional[str]
+    :param max_columns: Maximum subplot columns.
+    :type max_columns: int
+    :returns: ``(figure, axes)``.
+    :rtype: tuple[plt.Figure, list[plt.Axes]]
+    """
+
+    if not hasattr(mtg, "get_its_steps"):
+        raise TypeError("draw_mtg_steps expects an MTG object with get_its_steps().")
+
+    all_steps = list(mtg.get_its_steps())
+    selected = list(range(len(all_steps))) if steps is None else list(steps)
+    for step in selected:
+        if step < 0 or step >= len(all_steps):
+            raise IndexError(f"MTG step index out of range: {step}")
+
+    panels = [(f"Step {step + 1}", all_steps[step]) for step in selected]
+    if include_composed:
+        if not hasattr(mtg, "get_compose_its"):
+            raise TypeError("include_composed requires an MTG object with get_compose_its().")
+        panels.append(("Composed", mtg.get_compose_its()))
+
+    if not panels:
+        raise ValueError("No MTG steps selected for drawing.")
+
+    ncols = min(max(1, max_columns), len(panels))
+    nrows = (len(panels) + ncols - 1) // ncols
+    fig, axes_grid = plt.subplots(
+        nrows,
+        ncols,
+        figsize=(4.8 * ncols, 4.2 * nrows),
+        squeeze=False,
+        facecolor="white",
+    )
+    axes = [ax for row in axes_grid for ax in row]
+    if title:
+        fig.suptitle(title, fontsize=13, fontweight="bold")
+
+    for ax, (panel_title, its) in zip(axes, panels):
+        draw_its_only(
+            its,
+            ax=ax,
+            title=panel_title,
+            show_atom_map=show_atom_map,
+            label_mode=label_mode,
+            edge_label_mode=edge_label_mode,
+            show_edge_labels=show_edge_labels,
+            show_electron_labels=show_electron_labels,
+            electron_label_mode=electron_label_mode,
+        )
+
+    for ax in axes[len(panels):]:
+        ax.set_axis_off()
+
+    fig.tight_layout()
+    return fig, axes[: len(panels)]
+
+
+def _as_mtg_graph(mtg: Any) -> nx.Graph:
+    if isinstance(mtg, nx.Graph):
+        return mtg
+    if hasattr(mtg, "get_mtg"):
+        graph = mtg.get_mtg()
+        if isinstance(graph, nx.Graph):
+            return graph
+    raise TypeError("Expected an MTG object or a NetworkX compact MTG graph.")