torch-flowgraph 0.1.0__py3-none-any.whl

torch_flowgraph/__init__.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""torch-flowgraph: PyTorch computation graph tracing and visualization.
+
+Public API
+----------
+:class:`GraphBase`
+    Base class for all computation graph representations.  JSON-deserialised
+    graphs are returned as :class:`GraphBase` instances.
+
+:class:`Clustering`
+    Abstract base class for reclustering strategies.
+
+:class:`ClanClustering`
+    Re-cluster using the McCreary clan-based graph-parsing algorithm.
+
+:class:`LevelClustering`
+    Keep only cluster boxes at a given depth; all nodes remain visible.
+
+:class:`NameClustering`
+    Keep only cluster boxes with specific names; all nodes remain visible.
+
+:func:`apply_clan_clustering`
+    Convenience wrapper around ``graph.recluster(ClanClustering())``.
+"""
+
+from . import patches as _patches  # noqa: F401 — apply torch meta-device patches
+from ._post_import import on_import
+from .clustering import ClanClustering
+from .clustering import Clustering
+from .clustering import LevelClustering
+from .clustering import NameClustering
+from .clustering import NoClustering
+from .clustering import apply_clan_clustering
+from .graph import GraphBase
+from .validation import validate_graph
+
+__all__ = [
+    "GraphBase",
+    "Clustering",
+    "ClanClustering",
+    "LevelClustering",
+    "NameClustering",
+    "NoClustering",
+    "apply_clan_clustering",
+    "validate_graph",
+    "on_import",
+]

torch_flowgraph/__main__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+
+from .cli import main
+
+sys.exit(main())

torch_flowgraph/_post_import.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""Post-import hooks.
+
+Provides :func:`on_import` to register a callback that fires the first time
+a given module is imported (or immediately, if already loaded).  This lets a
+package register patches that apply only when a third-party dependency is
+actually pulled in by some other code path.
+"""
+
+import builtins
+import sys
+
+_pending: dict = {}
+_done: set = set()
+_orig_import = None
+_depth = 0
+
+
+def on_import(module_name, callback):
+    """Run *callback(module)* the first time *module_name* is imported.
+
+    If the module is already in :data:`sys.modules`, *callback* is invoked
+    immediately.  Otherwise an import hook is installed that fires the
+    callback exactly once, the first time the module appears in
+    :data:`sys.modules` *after* the outermost ``import`` statement that
+    caused it to load has fully completed.  (Submodule imports during a
+    package's ``__init__.py`` execution don't fire callbacks for the parent
+    package, which would otherwise see a half-initialised module.)
+
+    Parameters
+    ----------
+    module_name:
+        Fully qualified module name (e.g. ``"torch_geometric.utils._subgraph"``).
+    callback:
+        Callable taking the imported module as its single argument.
+    """
+    if module_name in sys.modules and module_name not in _pending:
+        _done.add(module_name)
+        callback(sys.modules[module_name])
+        return
+
+    _pending.setdefault(module_name, []).append(callback)
+    _install_hook()
+
+
+def _install_hook():
+    global _orig_import
+    if _orig_import is not None:
+        return
+    _orig_import = builtins.__import__
+
+    def _wrapped_import(name, globals=None, locals=None, fromlist=(), level=0):
+        global _depth
+        _depth += 1
+        try:
+            result = _orig_import(name, globals, locals, fromlist, level)
+        finally:
+            _depth -= 1
+
+        # Only fire callbacks once we have returned to the outermost
+        # __import__ call.  Otherwise modules in sys.modules may still be
+        # mid-execution (a package whose __init__.py is partially run).
+        if _depth == 0 and _pending:
+            for pending_name in list(_pending):
+                if pending_name in sys.modules and pending_name not in _done:
+                    _done.add(pending_name)
+                    callbacks = _pending.pop(pending_name)
+                    mod = sys.modules[pending_name]
+                    for cb in callbacks:
+                        try:
+                            cb(mod)
+                        except Exception:
+                            pass
+        return result
+
+    builtins.__import__ = _wrapped_import

torch_flowgraph/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

torch_flowgraph/attributes.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""Graphviz DOT attribute name lists.
+
+This module enumerates the valid attribute names for different element types
+in Graphviz DOT files.  The lists are used by the plotting functions to
+filter out unknown attributes when building DOT output, ensuring that only
+recognised attributes are emitted.
+
+Attribute hierarchy
+-------------------
+``COMMON_ATTRIBUTES``
+    Attributes valid for all element types (nodes, edges, graphs, clusters).
+``NODE_EDGE_ATTRIBUTES``
+    Attributes valid for both nodes and edges (extends ``COMMON_ATTRIBUTES``).
+``NODE_ATTRIBUTES``
+    Node-specific attributes (extends ``NODE_EDGE_ATTRIBUTES``).
+``EDGE_ATTRIBUTES``
+    Edge-specific attributes (extends ``NODE_EDGE_ATTRIBUTES``).
+``CLUSTER_GRAPH_ATTRIBUTES``
+    Attributes valid for both clusters and graphs (extends
+    ``COMMON_ATTRIBUTES``).
+``CLUSTER_ATTRIBUTES``
+    Cluster-specific attributes (extends ``CLUSTER_GRAPH_ATTRIBUTES``).
+``GRAPH_ATTRIBUTES``
+    Graph-specific attributes (extends ``CLUSTER_GRAPH_ATTRIBUTES``).
+
+Reference: https://graphviz.org/doc/info/attrs.html
+"""
+
+# Attributes shared across all Graphviz element types.
+COMMON_ATTRIBUTES = [
+    "URL",
+    "class",
+    "colorscheme",
+    "fontcolor",
+    "fontname",
+    "fontsize",
+    "href",
+    "id",
+    "label",
+    "nojustify",
+    "style",
+    "target",
+    "tooltip",
+]
+
+# Attributes valid for both nodes and edges (but not graphs/clusters).
+NODE_EDGE_ATTRIBUTES = COMMON_ATTRIBUTES + [
+    "color",
+    "comment",
+    "fillcolor",
+    "layer",
+    "penwidth",
+    "pos",
+    "showboxes",
+    "xlabel",
+    "xlp",
+]
+
+# Attributes valid only for nodes.
+NODE_ATTRIBUTES = NODE_EDGE_ATTRIBUTES + [
+    "area",
+    "distortion",
+    "fixedsize",
+    "gradientangle",
+    "group",
+    "height",
+    "image",
+    "imagepos",
+    "imagescale",
+    "labelloc",
+    "margin",
+    "ordering",
+    "orientation",
+    "peripheries",
+    "pin",
+    "rects",
+    "regular",
+    "root",
+    "samplepoints",
+    "shape",
+    "shapefile",
+    "sides",
+    "skew",
+    "sortv",
+    "vertices",
+    "width",
+    "z",
+]
+
+# Attributes valid only for edges.
+EDGE_ATTRIBUTES = NODE_EDGE_ATTRIBUTES + [
+    "arrowhead",
+    "arrowsize",
+    "arrowtail",
+    "constraint",
+    "decorate",
+    "dir",
+    "edgeURL",
+    "edgehref",
+    "edgetarget",
+    "edgetooltip",
+    "headURL",
+    "head_lp",
+    "headclip",
+    "headhref",
+    "headlabel",
+    "headport",
+    "headtarget",
+    "headtooltip",
+    "labelURL",
+    "labelangle",
+    "labeldistance",
+    "labelfloat",
+    "labelfontcolor",
+    "labelfontname",
+    "labelfontsize",
+    "labelhref",
+    "labeltarget",
+    "labeltooltip",
+    "len",
+    "lhead",
+    "lp",
+    "ltail",
+    "minlen",
+    "samehead",
+    "sametail",
+    "tailURL",
+    "tail_lp",
+    "tailclip",
+    "tailhref",
+    "taillabel",
+    "tailport",
+    "tailtarget",
+    "tailtooltip",
+    "weight",
+]
+
+# Attributes valid for both clusters and top-level graphs.
+CLUSTER_GRAPH_ATTRIBUTES = COMMON_ATTRIBUTES + [
+    "K",
+    "bb",
+    "bgcolor",
+    "gradientangle",
+    "labeljust",
+    "labelloc",
+    "lheight",
+    "lp",
+    "lwidth",
+    "margin",
+    "sortv",
+]
+
+# Attributes valid only for clusters (subgraphs whose name begins with
+# ``cluster_``).
+CLUSTER_ATTRIBUTES = CLUSTER_GRAPH_ATTRIBUTES + [
+    "area",
+    "cluster",
+    "color",
+    "fillcolor",
+    "layer",
+    "pencolor",
+    "penwidth",
+    "peripheries",
+]
+
+# Attributes valid only for the top-level graph.
+GRAPH_ATTRIBUTES = CLUSTER_GRAPH_ATTRIBUTES + [
+    "Damping",
+    "TBbalance",
+    "_background",
+    "beautify",
+    "center",
+    "charset",
+    "clusterrank",
+    "comment",
+    "compound",
+    "concentrate",
+    "defaultdist",
+    "dim",
+    "dimen",
+    "diredgeconstraints",
+    "dpi",
+    "epsilon",
+    "esep",
+    "fontnames",
+    "fontpath",
+    "forcelabels",
+    "imagepath",
+    "inputscale",
+    "label_scheme",
+    "landscape",
+    "layerlistsep",
+    "layers",
+    "layerselect",
+    "layersep",
+    "layout",
+    "levels",
+    "levelsgap",
+    "linelength",
+    "maxiter",
+    "mclimit",
+    "mindist",
+    "mode",
+    "model",
+    "newrank",
+    "nodesep",
+    "normalize",
+    "notranslate",
+    "nslimit",
+    "nslimit1",
+    "oneblock",
+    "ordering",
+    "orientation",
+    "outputorder",
+    "overlap",
+    "overlap_scaling",
+    "overlap_shrink",
+    "pack",
+    "packmode",
+    "pad",
+    "page",
+    "pagedir",
+    "quadtree",
+    "quantum",
+    "rankdir",
+    "ranksep",
+    "ratio",
+    "remincross",
+    "repulsiveforce",
+    "resolution",
+    "root",
+    "rotate",
+    "rotation",
+    "scale",
+    "searchsize",
+    "sep",
+    "showboxes",
+    "size",
+    "smoothing",
+    "splines",
+    "start",
+    "stylesheet",
+    "truecolor",
+    "viewport",
+    "voro_margin",
+    "xdotversion",
+]
+
+
+if __name__ == "__main__":
+    # Helper script to verify that the attribute group boundaries are correct.
+    common = (set(CLUSTER_ATTRIBUTES).intersection(GRAPH_ATTRIBUTES)) - set(COMMON_ATTRIBUTES)
+    print(f"CLUSTER_GRAPH_ATTRIBUTES = COMMON_ATTRIBUTES + {repr(sorted(common- set(COMMON_ATTRIBUTES)))}")
+    print(
+        f"CLUSTER_ATTRIBUTES =  CLUSTER_GRAPH_ATTRIBUTES+{repr(sorted(set(CLUSTER_ATTRIBUTES)-common- set(COMMON_ATTRIBUTES)))}"
+    )
+    print(
+        f"GRAPH_ATTRIBUTES = CLUSTER_GRAPH_ATTRIBUTES + {repr(sorted(set(GRAPH_ATTRIBUTES)-common- set(COMMON_ATTRIBUTES)))}"
+    )

torch_flowgraph/cli.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: 2026 European Centre for Medium-Range Weather Forecasts (ECMWF)
+# SPDX-License-Identifier: Apache-2.0
+
+"""Command-line interface for torch-flowgraph.
+
+Usage::
+
+    torch-flowgraph validate path/to/graph.json
+    python -m torch_flowgraph validate path/to/graph.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from .graph import GraphBase
+from .validation import validate_graph
+
+
+def _cmd_validate(args) -> int:
+    graph = GraphBase.from_json_file(args.path)
+    ok = validate_graph(graph, verbose=not args.quiet)
+    return 0 if ok else 1
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="torch-flowgraph",
+        description="Inspect and validate torch-flowgraph JSON graphs.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    validate = sub.add_parser("validate", help="Check a graph JSON file for inconsistencies.")
+    validate.add_argument("path", help="Path to a graph JSON file.")
+    validate.add_argument("-q", "--quiet", action="store_true", help="Suppress the report; only set the exit code.")
+    validate.set_defaults(func=_cmd_validate)
+
+    args = parser.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())