torch-flowgraph 0.1.0__py3-none-any.whl

torch_flowgraph/clustering.py

@@ -0,0 +1,282 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""Clustering strategies for computation graphs.
+
+This module provides:
+
+* :class:`Clustering` — abstract base class for all clustering strategies.
+* :class:`ClanClustering` — re-clusters using the McCreary clan algorithm.
+* :class:`LevelClustering` — shows only cluster boxes at a specific depth.
+* :class:`NameClustering` — shows only cluster boxes with specific names.
+
+All concrete classes implement :meth:`~Clustering.build_clusters`, which
+receives a :class:`~torch_flowgraph.graph.GraphBase` and returns a new
+:class:`~torch_flowgraph.clusters.Cluster` root.  Use
+:meth:`~torch_flowgraph.graph.GraphBase.recluster` to apply a strategy::
+
+    from torch_flowgraph.clustering import LevelClustering
+    shallow = graph.recluster(LevelClustering(level=1))
+
+Backward-compatibility shim
+----------------------------
+:func:`apply_clan_clustering` is kept as a convenience wrapper so that
+existing code does not break::
+
+    from torch_flowgraph import apply_clan_clustering
+    clan_graph = apply_clan_clustering(graph)
+"""
+
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+
+from .clusters import Cluster
+
+
+class Clustering(ABC):
+    """Abstract base class for graph reclustering strategies.
+
+    Subclasses implement :meth:`build_clusters`, which returns a new
+    :class:`~torch_flowgraph.clusters.Cluster` root built from *graph*.
+    The graph's nodes and edges are *not* modified; only the cluster
+    hierarchy changes.
+    """
+
+    @abstractmethod
+    def build_clusters(self, graph) -> Cluster:
+        """Build and return a new cluster root for *graph*.
+
+        Parameters
+        ----------
+        graph:
+            Source :class:`~torch_flowgraph.graph.GraphBase`.
+
+        Returns
+        -------
+        Cluster
+            Root of the new cluster hierarchy.  The returned tree must
+            reference only node-ids that exist in *graph*.
+        """
+
+
+# ---------------------------------------------------------------------------
+# Clan clustering
+# ---------------------------------------------------------------------------
+
+
+class ClanClustering(Clustering):
+    """Re-cluster using the McCreary clan-based graph-parsing algorithm.
+
+    The resulting cluster hierarchy reflects the *dataflow* decomposition
+    of the graph (linear chains, independent sub-graphs, primitives) rather
+    than the original ``nn.Module`` code hierarchy.
+
+    Clan kinds assigned to clusters: ``LINEAR``, ``INDEPENDENT``,
+    ``PRIMITIVE``.
+
+    Trivial clusters are dissolved after parsing: single-node clusters
+    have their node promoted to the parent, and single-child wrapper
+    clusters (with no own nodes) are replaced by their only child.
+    """
+
+    def build_clusters(self, graph) -> Cluster:
+        from .parsing import Parser
+
+        G = graph.to_op_networkx()
+        parser = Parser(G)
+        parse_tree = parser.parse()
+        root = _parse_tree_to_cluster(parse_tree, parent=None)
+        root.update_cluster_nodes()
+        _dissolve_leaf_clusters(root)
+        return root
+
+
+def _parse_tree_to_cluster(node, parent) -> Cluster:
+    """Convert a :class:`~torch_flowgraph.parsing.ParseTree` node to a Cluster."""
+    from .parsing import Kind
+
+    kind_name = node.clan.kind.name if node.clan.kind != Kind.UNDEF else "UNKNOWN"
+    cluster = Cluster(name=kind_name, parent=parent, nodes=list(node.clan.nodes))
+    cluster.children = [_parse_tree_to_cluster(child, cluster) for child in node.children]
+    return cluster
+
+
+def _dissolve_leaf_clusters(cluster):
+    """Remove trivial clusters, attaching their contents to the parent.
+
+    A cluster is trivial when it contains only a single node (no
+    sub-clusters) or only a single sub-cluster (no direct nodes).
+    Runs repeatedly until no more trivial clusters remain.
+    """
+    changed = True
+    while changed:
+        changed = False
+        kept = []
+        for child in cluster.children:
+            _dissolve_leaf_clusters(child)
+            if not child.children and len(child.nodes) <= 1:
+                # Single-node (or empty) cluster — promote nodes to parent.
+                cluster.nodes.extend(child.nodes)
+                changed = True
+            elif len(child.children) == 1 and not child.nodes:
+                # Single-child wrapper with no own nodes — replace with the child.
+                grandchild = child.children[0]
+                grandchild.parent = cluster
+                kept.append(grandchild)
+                changed = True
+            else:
+                kept.append(child)
+        cluster.children = kept
+
+
+# ---------------------------------------------------------------------------
+# Selective clustering (level / name filters)
+# ---------------------------------------------------------------------------
+
+
+class _SelectiveClustering(Clustering):
+    """Base for strategies that select a subset of existing cluster boxes.
+
+    Matching clusters are kept as labeled boxes.  All other clusters are
+    dropped: their nodes are absorbed into the flat root so the complete
+    node set is still rendered — just without grouping boxes.
+    """
+
+    @abstractmethod
+    def _find_matching(self, root_cluster: Cluster) -> list[Cluster]:
+        """Return clusters from *root_cluster*'s tree that should be drawn."""
+
+    def build_clusters(self, graph) -> Cluster:
+        matching = self._find_matching(graph.clusters)
+
+        # Collect all node-ids claimed by matching cluster subtrees.
+        matching_nodes: set[int] = set()
+        for c in matching:
+            matching_nodes.update(c.all_nodes)
+
+        # Root gets every node *not* in a matching subtree (shown ungrouped).
+        root = Cluster(name=None)
+        root.nodes = [n for n in graph.nodes if n not in matching_nodes]
+        root.children = []
+
+        for c in matching:
+            child = Cluster(name=c.name, nodes=list(c.all_nodes), parent=root)
+            root.children.append(child)
+
+        return root
+
+
+class LevelClustering(_SelectiveClustering):
+    """Show only cluster boxes at one or more specific depths.
+
+    All nodes remain in the graph; those not inside a matching cluster
+    are rendered without a surrounding box.
+
+    Parameters
+    ----------
+    level : int or collection of int
+        Target cluster level(s), 0-indexed from the outermost real cluster.
+        ``0`` means the top-level cluster boxes (the root's direct children),
+        ``1`` their children, etc.  The nameless structural root is not a
+        selectable level.  Pass a single ``int`` or any iterable of ``int``
+        values to select multiple levels simultaneously.
+    """
+
+    def __init__(self, level) -> None:
+        if isinstance(level, int):
+            self.levels: frozenset[int] = frozenset([level])
+        else:
+            self.levels = frozenset(level)
+
+    def _find_matching(self, root_cluster: Cluster) -> list[Cluster]:
+        matching: list[Cluster] = []
+        # User-facing level L corresponds to internal cluster depth L + 1,
+        # since the root (depth 0) is a nameless structural container.
+        target_depths = frozenset(level + 1 for level in self.levels)
+        max_depth = max(target_depths) if target_depths else 0
+
+        def _walk(c: Cluster) -> None:
+            if c.depth in target_depths:
+                matching.append(c)
+                # Don't recurse: the whole subtree is kept in this box.
+            elif c.depth < max_depth:
+                for child in c.children:
+                    _walk(child)
+
+        _walk(root_cluster)
+        return matching
+
+
+class NameClustering(_SelectiveClustering):
+    """Show only cluster boxes whose name is in *names*.
+
+    All nodes remain in the graph; nodes not inside a named cluster are
+    rendered without a surrounding box.
+
+    Parameters
+    ----------
+    names : collection of str
+        Cluster names to keep.  When a matching cluster is found, its
+        children are not inspected (the whole subtree is included in the box).
+    """
+
+    def __init__(self, names) -> None:
+        self.names: frozenset[str] = frozenset(names)
+
+    def _find_matching(self, root_cluster: Cluster) -> list[Cluster]:
+        matching: list[Cluster] = []
+
+        def _walk(c: Cluster) -> None:
+            if c.name in self.names:
+                matching.append(c)
+            else:
+                for child in c.children:
+                    _walk(child)
+
+        _walk(root_cluster)
+        return matching
+
+
+# ---------------------------------------------------------------------------
+# No-op clustering — clears all cluster boxes
+# ---------------------------------------------------------------------------
+
+
+class NoClustering(Clustering):
+    """Remove all cluster boxes, rendering every node at the top level.
+
+    The returned cluster hierarchy is a single flat root containing all
+    nodes and no children.  The graph itself is unchanged; only the cluster
+    boxes disappear from the rendered output.
+    """
+
+    def build_clusters(self, graph) -> Cluster:
+        root = Cluster(name=None)
+        root.nodes = list(graph.nodes)
+        return root
+
+
+# ---------------------------------------------------------------------------
+# Backward-compatibility shim
+# ---------------------------------------------------------------------------
+
+
+def apply_clan_clustering(graph):
+    """Re-cluster *graph* using the McCreary clan algorithm.
+
+    Convenience wrapper around ``graph.recluster(ClanClustering())``.
+    Kept for backward compatibility with existing code.
+
+    Parameters
+    ----------
+    graph:
+        Source :class:`~torch_flowgraph.graph.GraphBase`.
+
+    Returns
+    -------
+    GraphBase
+        New graph with clan-derived cluster hierarchy.
+    """
+    return graph.recluster(ClanClustering())

torch_flowgraph/clusters.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""Cluster data structure for hierarchical graph grouping.
+
+A :class:`Cluster` represents a named group of graph nodes, arranged in a
+tree via parent/children references.  Clusters are used to generate Graphviz
+``subgraph cluster_*`` blocks and to control the collapsing/expansion of
+module hierarchies during rendering.
+"""
+
+from functools import cached_property
+
+
+class Cluster:
+    """A named, hierarchical group of graph nodes.
+
+    Clusters form a tree: each cluster has a *parent* and zero or more
+    *children* (sub-clusters).  The *nodes* list holds the integer node-ids
+    that belong directly to this cluster (not to its children).
+
+    Attributes
+    ----------
+    name : str or None
+        Human-readable cluster name shown in Graphviz DOT output as the
+        ``label`` attribute.
+    nodes : list[int]
+        Integer node-ids that are direct members of this cluster.
+    children : list[Cluster]
+        Sub-clusters nested inside this cluster.
+    parent : Cluster or None
+        Parent cluster, or ``None`` for the root.
+    depth : int
+        Depth of this cluster in the tree (set by :class:`ModelGraph`
+        during graph finalisation; ``-1`` until then).
+    visible : bool or None
+        Explicit visibility override.  ``None`` means visibility is
+        determined automatically by :meth:`show_cluster`.
+    kwargs : dict
+        Additional Graphviz DOT attributes passed through to
+        :attr:`dot_attributes`.
+    """
+
+    def __init__(self, *, name=None, parent=None, visible=None, nodes=None, children=None, **kwargs):
+        """Initialise a Cluster.
+
+        Parameters
+        ----------
+        name:
+            Display label for the cluster in DOT output.
+        parent:
+            Parent :class:`Cluster`, or ``None`` for the root.
+        visible:
+            Explicit visibility flag.  When ``None`` the cluster is always
+            shown (see :meth:`show_cluster`).
+        nodes:
+            Initial list of direct member node-ids.  Defaults to ``[]``.
+        children:
+            Initial list of child :class:`Cluster` instances.  Defaults to
+            ``[]``.
+        **kwargs:
+            Additional Graphviz DOT attributes (e.g. ``bgcolor``,
+            ``pencolor``).
+        """
+        self.name = name
+        self.nodes = [] if nodes is None else nodes
+        self.children = [] if children is None else children
+        self.parent = parent
+        self.depth = -1
+        self.visible = visible
+        self.is_repeated = False
+        self.kwargs = kwargs
+
+    @property
+    def dot_attributes(self):
+        """Graphviz DOT attribute dictionary for this cluster's subgraph block.
+
+        Starts from any extra keyword arguments supplied at construction time
+        and adds a ``"label"`` key when :attr:`name` is set.
+
+        Returns
+        -------
+        dict
+            Mapping of DOT attribute name to value.
+        """
+        attr = self.kwargs.copy()
+        if self.name is not None:
+            attr["label"] = self.name
+        return attr
+
+    def __len__(self):
+        """Return the number of direct member nodes."""
+        return len(self.nodes)
+
+    def __repr__(self):
+        return f"Cluster({self.name})"
+
+    def as_dict(self):
+        """Serialise the cluster tree to a JSON-compatible dictionary.
+
+        Returns
+        -------
+        dict
+            Dictionary with ``"name"``, ``"nodes"``, and ``"children"``
+            keys.  Children are recursively serialised.
+        """
+        return {
+            "name": self.name,
+            "nodes": self.nodes,
+            "children": [c.as_dict() for c in self.children],
+        }
+
+    @classmethod
+    def from_dict(cls, data, parent):
+        """Deserialise a cluster tree from a dictionary.
+
+        Handles both the current ``"children"`` key and the legacy
+        ``"clusters"`` key for backwards compatibility.
+
+        Parameters
+        ----------
+        data:
+            Dictionary as produced by :meth:`as_dict`.
+        parent:
+            Parent :class:`Cluster` to assign to the newly created cluster
+            (``None`` for the root).
+
+        Returns
+        -------
+        Cluster
+            Fully reconstructed cluster tree.
+        """
+        data = data.copy()
+        children = data.pop("children", data.get("clusters", []))
+        cluster = cls(parent=parent, **data)
+        cluster.children = [cls.from_dict(c, cluster) for c in children]
+        return cluster
+
+    @cached_property
+    def all_nodes(self):
+        """Set of all node-ids contained anywhere in this cluster subtree.
+
+        Computed once and cached.  Includes nodes from this cluster and all
+        descendant clusters.
+
+        Returns
+        -------
+        set[int]
+        """
+        result = set()
+
+        def _all_nodes(c):
+            for n in c.nodes:
+                result.add(n)
+            for c in c.children:
+                _all_nodes(c)
+
+        _all_nodes(self)
+        return result
+
+    def show_cluster(self, depth):
+        """Determine whether this cluster should be rendered as a subgraph.
+
+        Returns :attr:`visible` when it is explicitly set, otherwise always
+        returns ``True``.
+
+        Parameters
+        ----------
+        depth:
+            Current rendering depth (not used in the default implementation
+            but available for subclasses).
+
+        Returns
+        -------
+        bool
+        """
+        if self.visible is not None:
+            return self.visible
+
+        return True
+
+    def update_cluster_nodes(self):
+        """Redistribute nodes so each node appears in exactly one cluster.
+
+        Starting from the full set of nodes in this subtree (:attr:`all_nodes`),
+        recursively assigns each node to the deepest cluster that contains it,
+        removing it from all ancestor clusters' :attr:`nodes` lists.
+
+        Returns
+        -------
+        Cluster
+            *self*, mutated in-place.
+        """
+        nodes = set(self.all_nodes)
+
+        def _update_cluster_nodes(c):
+            nonlocal nodes
+
+            for c in c.children:
+                _update_cluster_nodes(c)
+
+            c.nodes = list(set(c.nodes).intersection(nodes))
+            nodes = nodes.difference(c.nodes)
+
+            return c
+
+        _update_cluster_nodes(self)
+
+        return self

torch_flowgraph/dot.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""High-level DOT export helper for :class:`ModelGraph`.
+
+This module provides :func:`save_dot`, a convenience wrapper around
+:func:`~torch_flowgraph.plotting.plot_graph` that accepts a
+:class:`~torch_flowgraph.model.graph.ModelGraph` and writes a Graphviz DOT
+file suitable for rendering with ``dot -Tpdf``.
+"""
+
+from torch_flowgraph.clusters import Cluster
+from torch_flowgraph.plotting import plot_graph
+
+
+def save_dot(
+    graph,
+    path,
+    graph_attributes=None,
+    clusters=None,
+    depth=None,
+    links=False,
+    show_clusters=True,
+    highlight_fn=None,
+    cluster_bg=True,
+):
+    """Write a :class:`ModelGraph` to a Graphviz DOT file.
+
+    Optionally replaces the graph's cluster hierarchy with *clusters* and/or
+    collapses the graph to *depth* levels before rendering.
+
+    Parameters
+    ----------
+    graph:
+        :class:`~torch_flowgraph.model.graph.ModelGraph` instance to export.
+    path:
+        Destination file path for the ``.dot`` output.
+    graph_attributes:
+        Optional dictionary of top-level Graphviz graph attributes (e.g.
+        ``{"rankdir": "LR"}``).
+    clusters:
+        Controls subgraph box rendering:
+
+        * ``None`` (default) — use the graph's own cluster hierarchy.
+        * A :class:`~torch_flowgraph.clusters.Cluster` instance — re-cluster
+          the graph with this hierarchy before rendering.
+        * ``[]`` (empty list) — suppress all subgraph boxes.
+    depth:
+        Maximum cluster depth to render.  When given, the graph is collapsed
+        to *depth* levels via
+        :meth:`~torch_flowgraph.model.graph.ModelGraph.collapsed`.
+    links:
+        When ``True``, function nodes include ``URL`` attributes linking to
+        the PyTorch documentation.  Defaults to ``False``.
+    show_clusters:
+        When ``False``, all cluster subgraph blocks are rendered as invisible
+        so they guide the layout engine without drawing visible boxes.
+        Forwarded to :func:`~torch_flowgraph.plotting.plot_graph`.
+        Defaults to ``True``.
+    cluster_bg:
+        When ``False``, cluster boxes are drawn without a background fill.
+        Forwarded to :func:`~torch_flowgraph.plotting.plot_graph`.
+        Defaults to ``True``.
+    """
+    # Re-cluster only when a real Cluster object is provided.
+    if isinstance(clusters, Cluster) and clusters is not graph.clusters:
+        graph = graph.use_clusters(clusters)
+
+    graph = graph.collapsed(depth) if depth is not None else graph
+
+    G = graph.to_op_networkx(links=links)
+    # None → show own clusters (skip root wrapper, use children directly);
+    # [] → no boxes; Cluster already applied above.
+    display_clusters = graph.clusters.children if clusters is None else clusters
+    plot_graph(
+        path,
+        G,
+        display_clusters,
+        graph_attributes=graph_attributes,
+        cluster_bg=cluster_bg,
+        show_clusters=show_clusters,
+        highlight_fn=highlight_fn,
+    )