From e024ea010d8e53cb643b8f55048553f75ecaf269 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 16 Oct 2026 15:11:47 -0400 Subject: [PATCH 1/4] Add export_graphviz for DOT export of fitted trees Edge labels reuse the export_text routing rules, so both exports agree. Nodes show samples, value and class; filled=True shades them like sklearn. Closes #69. --- ROADMAP.md | 2 +- docs/api/index.rst | 1 + docs/api/plotting.rst | 20 +++- docs/roadmap.rst | 2 +- sgtlearn/__init__.py | 3 +- sgtlearn/_export.py | 215 ++++++++++++++++++++++++++++++++------ tests/test_export_text.py | 95 ++++++++++++++++- 7 files changed, 299 insertions(+), 39 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index cb4ae8f..7ed616c 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -11,7 +11,7 @@ ## v0.2.0 - [x] Superset branching on categorical features - [x] More plotting options -- [ ] Exporting to Graphviz, planned for v0.4.0 ([#69](https://github.com/optimal-uoft/sgtlearn/issues/69)) +- [x] Exporting to Graphviz with `export_graphviz` ([#69](https://github.com/optimal-uoft/sgtlearn/issues/69)) - [x] Feature importances matching scikit-learn's API for SGTs and SGT$_K$ - [x] Adding TAO Refinement - [x] **Sklearn-style NaN support**: split search uses finite values only, and each univariate node learns which child receives missing values. diff --git a/docs/api/index.rst b/docs/api/index.rst index 8158d9a..1ddfd50 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -15,6 +15,7 @@ The public API is re-exported from the top-level ``sgtlearn`` package. plot_tree configure_feature_dict export_text + export_graphviz make_plus tao.TAO_refine diff --git a/docs/api/plotting.rst b/docs/api/plotting.rst index 2c36070..447bf10 100644 --- a/docs/api/plotting.rst +++ b/docs/api/plotting.rst @@ -3,7 +3,7 @@ Plotting & Export .. currentmodule:: sgtlearn -Visualize a fitted SGT, or export it as indented text rules. +Visualize a fitted SGT, or export it as indented text rules or Graphviz DOT. plot_tree --------- @@ -36,3 +36,21 @@ apply to ``X`` rounded to float32: ``1.5 + 1e-9`` rounds to ``1.5`` and follows | |--- value: 0.00 |--- x > 1.50 | |--- value: 1.00 + +export_graphviz +--------------- + +.. autofunction:: export_graphviz + +Edge labels use the same routing rules as ``export_text``, so the two exports +cannot disagree. ``export_graphviz`` writes plain DOT text and needs no extra +package. To render it, install the ``graphviz`` Python package and the Graphviz +binaries: + +.. code-block:: python + + import graphviz + from sgtlearn import export_graphviz + + dot = export_graphviz(model, feature_names=feature_names, filled=True) + graphviz.Source(dot).render("tree", format="png") diff --git a/docs/roadmap.rst b/docs/roadmap.rst index 0452d76..57d6ab7 100644 --- a/docs/roadmap.rst +++ b/docs/roadmap.rst @@ -18,7 +18,7 @@ v0.2.0 - ✅ Superset branching on categorical features - ✅ More plotting options -- ⬜ Exporting to Graphviz, planned for v0.4.0 (see +- ✅ Exporting to Graphviz with ``export_graphviz`` (see `issue #69 `_) - ✅ Feature importances matching scikit-learn's API for SGTs and :math:`\mathrm{SGT}_K` - ✅ TAO refinement diff --git a/sgtlearn/__init__.py b/sgtlearn/__init__.py index c354463..f7770bb 100644 --- a/sgtlearn/__init__.py +++ b/sgtlearn/__init__.py @@ -6,7 +6,7 @@ """ from sgtlearn import tao -from sgtlearn._export import export_text, plot_tree +from sgtlearn._export import export_graphviz, export_text, plot_tree from sgtlearn.base import ( BaseShapeCART, ProcessedFeatures, @@ -25,6 +25,7 @@ "SGTClassifier", "SGTRegressor", "configure_feature_dict", + "export_graphviz", "export_text", "make_plus", "plot_tree", diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index e295302..a951757 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -5,7 +5,8 @@ bins colored by the destination child partition (a two-feature heatmap for pair nodes) plus histograms of ``X`` when it is supplied, and every leaf is drawn as colored text with the predicted class / value. -``export_text`` returns the same tree as indented routing rules. +``export_text`` returns the same tree as indented routing rules, and +``export_graphviz`` as Graphviz DOT source. """ from __future__ import annotations @@ -16,7 +17,7 @@ import numpy as np from matplotlib.patches import FancyArrowPatch, Rectangle -__all__ = ["export_text", "plot_tree"] +__all__ = ["export_graphviz", "export_text", "plot_tree"] import matplotlib.pyplot as plt from sklearn.utils.validation import check_is_fitted @@ -62,32 +63,11 @@ def export_text( str One line per child and leaf, indented by depth. """ - _check_fitted_sgt(estimator, "export_text") - n_features = estimator.n_features_in_ or 0 - if feature_names is not None and len(feature_names) != n_features: - raise ValueError( - f"feature_names has {len(feature_names)} entries, but the estimator " - f"was fit on {n_features} features" - ) - feat_names = _resolve_feature_names(estimator, feature_names) - tree = estimator.tree_export() + tree, feat_names, group_label, names_per_output = _prepare_export( + estimator, feature_names, class_names, "export_text" + ) nodes = {n["id"]: n for n in tree["nodes"]} - def group_label(cols: Sequence[int]) -> str: - if len(cols) == 1: - return _column_label(int(cols[0]), feat_names) - return _logical_feature_label( - estimator, - {"features": list(cols)}, - feat_names, - prefer_logical_name=feature_names is None, - ) - - names: Any = class_names - if names is None: - names = getattr(estimator, "classes_", None) - names_per_output: list[Any] = [names] if tree["num_outputs"] == 1 else names - def leaf_text(node: dict) -> str: if not isinstance(estimator, SGTClassifier): value = node["value"] @@ -95,9 +75,7 @@ def leaf_text(node: dict) -> str: return f"value: {_fmt_list(value, decimals)}" return f"value: {value:.{decimals}f}" counts = node["class_counts"] - text = "class: " + _per_output( - [str(names_per_output[o][int(np.argmax(c))]) for o, c in enumerate(counts)] - ) + text = "class: " + _majority_class(counts, names_per_output) if show_weights: fmt = [_fmt_list(c, decimals) for c in counts] text = f"weights: {_per_output(fmt)} {text}" @@ -121,6 +99,183 @@ def walk(nid: int, depth: int) -> None: return "\n".join(lines) + "\n" +def export_graphviz( + estimator: Any, + out_file: Any = None, + *, + feature_names: Sequence[str] | None = None, + class_names: Sequence[Any] | None = None, + filled: bool = False, + decimals: int = 2, +) -> str | None: + """Export the fitted tree in Graphviz DOT format. + + Every node shows its routing feature (internal nodes), ``samples`` and + ``value``; classifier nodes also show the majority ``class``. Regressor + internal nodes carry no ``value``, so they show only ``samples``. An + internal node with ``K`` children has ``K`` out-edges, each labelled with + the inputs routed to that child, using the same rules as ``export_text``. + + Producing DOT needs no extra package. To render it, install the + ``graphviz`` Python package and binaries and use ``graphviz.Source(dot)``. + + Parameters + ---------- + estimator : SGTClassifier or SGTRegressor + Fitted estimator to export. + out_file : str, path-like or file object, optional + Where to write the DOT source. When omitted, it is returned instead. + feature_names : sequence of str, optional + Names for the input columns, as in ``export_text``. + class_names : sequence, optional + Classifier only: class names, as in ``export_text``. + filled : bool, default=False + Fill nodes by majority class (classifier, shaded by purity) or by + value (regressor, shaded from the smallest to the largest value). + decimals : int, default=2 + Digits shown for thresholds and values. + + Returns + ------- + str or None + The DOT source when ``out_file`` is None, else None. + """ + tree, feat_names, group_label, names_per_output = _prepare_export( + estimator, feature_names, class_names, "export_graphviz" + ) + is_classifier = isinstance(estimator, SGTClassifier) + values = [v for n in tree["nodes"] if (v := _first_value(n)) is not None] + v_lo, v_hi = (min(values), max(values)) if values else (0.0, 0.0) + + def fill(node: dict) -> str: + """sklearn-style shading; multi-output trees use the first output.""" + if is_classifier: + counts = np.asarray(node["class_counts"][0], dtype=float) + if counts.sum() <= 0: + return "#ffffff" + p = np.sort(counts / counts.sum())[::-1] + second = p[1] if len(p) > 1 else 0.0 + color = _DEFAULT_PALETTE_COLORS[ + int(np.argmax(counts)) % len(_DEFAULT_PALETTE_COLORS) + ] + return _blend_white(color, (p[0] - second) / (1 - second)) + value = _first_value(node) + if value is None: + return "#ffffff" + alpha = 1.0 if v_hi == v_lo else (value - v_lo) / (v_hi - v_lo) + return _blend_white(_DEFAULT_PALETTE_COLORS[0], alpha) + + style = '"filled, rounded"' if filled else '"rounded"' + lines = [ + "digraph Tree {", + f'node [shape=box, style={style}, color="black", fontname="helvetica"] ;', + 'edge [fontname="helvetica"] ;', + ] + for node in tree["nodes"]: + label = ( + [] + if node["is_leaf"] + else [_node_feature_label(node, feat_names, group_label)] + ) + label.append(f"samples = {node['n_samples']}") + if is_classifier: + counts = node["class_counts"] + label.append( + "value = " + _per_output([_fmt_list(c, decimals) for c in counts]) + ) + label.append("class = " + _majority_class(counts, names_per_output)) + elif isinstance(node["value"], (list, tuple)): + label.append(f"value = {_fmt_list(node['value'], decimals)}") + elif node["value"] is not None: + label.append(f"value = {node['value']:.{decimals}f}") + attrs = f'label="{_dot_escape(chr(10).join(label))}"' + if filled: + attrs += f', fillcolor="{fill(node)}"' + lines.append(f"{node['id']} [{attrs}] ;") + if node["is_leaf"]: + continue + for region, cid in zip(_child_regions(node), node["children"]): + text = _format_region(region, feat_names, group_label, decimals) + lines.append(f'{node["id"]} -> {cid} [label="{_dot_escape(text)}"] ;') + lines.append("}") + dot = "\n".join(lines) + "\n" + if out_file is None: + return dot + if hasattr(out_file, "write"): + out_file.write(dot) + else: + with open(out_file, "w", encoding="utf-8") as f: + f.write(dot) + return None + + +def _prepare_export( + estimator: Any, + feature_names: Sequence[str] | None, + class_names: Sequence[Any] | None, + caller: str, +) -> tuple[dict, list[str], Callable[[Sequence[int]], str], list[Any]]: + """Validate inputs shared by the exporters and resolve display names.""" + _check_fitted_sgt(estimator, caller) + n_features = estimator.n_features_in_ or 0 + if feature_names is not None and len(feature_names) != n_features: + raise ValueError( + f"feature_names has {len(feature_names)} entries, but the estimator " + f"was fit on {n_features} features" + ) + feat_names = _resolve_feature_names(estimator, feature_names) + tree = estimator.tree_export() + + def group_label(cols: Sequence[int]) -> str: + if len(cols) == 1: + return _column_label(int(cols[0]), feat_names) + return _logical_feature_label( + estimator, + {"features": list(cols)}, + feat_names, + prefer_logical_name=feature_names is None, + ) + + names: Any = class_names + if names is None: + names = getattr(estimator, "classes_", None) + names_per_output: list[Any] = [names] if tree["num_outputs"] == 1 else names + return tree, feat_names, group_label, names_per_output + + +def _majority_class(counts: Sequence[Sequence[float]], names_per_output: list) -> str: + return _per_output( + [str(names_per_output[o][int(np.argmax(c))]) for o, c in enumerate(counts)] + ) + + +def _node_feature_label( + node: dict, feat_names: list[str], group_label: Callable[[Sequence[int]], str] +) -> str: + if _is_pair_node(node): + return " & ".join(group_label(a["columns"]) for a in node["pair_axes"]) + if _is_categorical_node(node): + return group_label(node["features"]) + return _column_label(int(node["feature"]), feat_names) + + +def _first_value(node: dict) -> float | None: + value = node.get("value") + if isinstance(value, (list, tuple)): + return float(value[0]) + return None if value is None else float(value) + + +def _blend_white(color: str, alpha: float) -> str: + """Mix a ``#rrggbb`` color with white; ``alpha=1`` keeps the color.""" + rgb = [int(color[i : i + 2], 16) for i in (1, 3, 5)] + return "#" + "".join(f"{round(alpha * c + (1 - alpha) * 255):02x}" for c in rgb) + + +def _dot_escape(text: str) -> str: + return text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") + + def _per_output(items: Sequence[str]) -> str: """One item as-is; several (one per output) as ``[a, b]``.""" return items[0] if len(items) == 1 else "[" + ", ".join(items) + "]" @@ -162,7 +317,7 @@ def _child_regions(node: dict) -> list[list[list[tuple]]]: - ``("none", cols)``: no column of ``cols`` is active (missing category) The rules replay the native router on float32 input, so ``export_text`` - (and ``export_graphviz``, #69) describe exactly what ``predict`` does. + and ``export_graphviz`` describe exactly what ``predict`` does. """ regions: list[list[list[tuple]]] = [[] for _ in node["children"]] if _is_pair_node(node): diff --git a/tests/test_export_text.py b/tests/test_export_text.py index 5546763..fe9645a 100644 --- a/tests/test_export_text.py +++ b/tests/test_export_text.py @@ -1,14 +1,16 @@ -"""Contracts for ``sgtlearn.export_text`` and the shared child-region rules.""" +"""Contracts for ``export_text``, ``export_graphviz`` and the shared child-region rules.""" from __future__ import annotations +import re + import numpy as np import pytest from sklearn.exceptions import NotFittedError from sklearn.tree import DecisionTreeClassifier import sgtlearn -from sgtlearn import SGTClassifier, SGTRegressor, export_text +from sgtlearn import SGTClassifier, SGTRegressor, export_graphviz, export_text from sgtlearn._export import ( _active_onehot_column, _child_regions, @@ -242,6 +244,89 @@ def test_export_text_validates_inputs(): export_text(model, feature_names=["only_one"]) -def test_export_graphviz_is_not_exported_until_implemented(): - assert not hasattr(sgtlearn, "export_graphviz") - assert "export_graphviz" not in sgtlearn.__all__ +_DOT_NODE = re.compile( + r'^(\d+) \[label="((?:[^"\\]|\\.)*)"(?:, fillcolor="(#[0-9a-f]{6})")?\] ;$' +) +_DOT_EDGE = re.compile(r'^(\d+) -> (\d+) \[label="((?:[^"\\]|\\.)*)"\] ;$') + + +def _parse_dot(dot: str) -> tuple[dict, list]: + """Parse ``export_graphviz`` output; every body line must be a node or edge.""" + lines = dot.splitlines() + assert lines[0] == "digraph Tree {" and lines[-1] == "}" + nodes, edges = {}, [] + for line in lines[3:-1]: + if m := _DOT_NODE.match(line): + nodes[int(m[1])] = (m[2].replace("\\n", "\n"), m[3]) + else: + m = _DOT_EDGE.match(line) + assert m, line + edges.append((int(m[1]), int(m[2]), m[3].replace('\\"', '"'))) + return nodes, edges + + +@pytest.mark.parametrize( + "make", + [_univariate_regressor, _categorical_classifier, _pair_classifier, _pair_regressor], +) +def test_export_graphviz_matches_export_text(make): + model, _ = make() + tree = model.tree_export() + nodes, edges = _parse_dot(export_graphviz(model, filled=True)) + assert set(nodes) == {n["id"] for n in tree["nodes"]} + assert all(color for _, color in nodes.values()) + assert sorted((a, b) for a, b, _ in edges) == sorted( + (n["id"], c) for n in tree["nodes"] for c in n["children"] + ) + child_lines = [ + line.split("|--- ", 1)[1] + for line in export_text(model).splitlines() + if "|--- class:" not in line and "|--- value:" not in line + ] + by_id = {n["id"]: n for n in tree["nodes"]} + dfs: list[tuple[int, int]] = [] + + def walk(nid: int) -> None: # export_text's line order + for cid in by_id[nid]["children"]: + dfs.append((nid, cid)) + walk(cid) + + walk(tree["root_index"]) + labels = {(a, b): label for a, b, label in edges} + assert [labels[e] for e in dfs] == child_lines + + +def test_export_graphviz_node_labels_and_out_file(tmp_path): + X = np.arange(4.0)[:, None] + model = SGTRegressor(max_leaf_nodes=2, tao_n_runs=0, random_state=0).fit( + X, np.array([0.0, 0.0, 1.0, 1.0]) + ) + dot = export_graphviz(model, feature_names=['x "q"']) + nodes, edges = _parse_dot(dot) + assert nodes[0] == ('x \\"q\\"\nsamples = 4', None) + assert sorted(label for label, _ in list(nodes.values())[1:]) == [ + "samples = 2\nvalue = 0.00", + "samples = 2\nvalue = 1.00", + ] + assert [label for *_, label in edges] == [ + 'x "q" <= 1.50 or x "q" is missing', + 'x "q" > 1.50', + ] + path = tmp_path / "tree.dot" + assert export_graphviz(model, out_file=path, feature_names=['x "q"']) is None + assert path.read_text() == dot + + clf = SGTClassifier(max_leaf_nodes=2, tao_n_runs=0, random_state=0).fit( + X, ["a", "a", "b", "b"] + ) + root = _parse_dot(export_graphviz(clf, class_names=["neg", "pos"]))[0][0][0] + assert root == "X[0]\nsamples = 4\nvalue = [2.00, 2.00]\nclass = neg" + + +def test_export_graphviz_is_public_and_validates_inputs(): + assert "export_graphviz" in sgtlearn.__all__ + with pytest.raises(TypeError): + export_graphviz(DecisionTreeClassifier().fit([[0], [1]], [0, 1])) + model = SGTRegressor(tao_n_runs=0).fit(np.zeros((4, 2)), np.arange(4.0)) + with pytest.raises(ValueError, match="feature_names"): + export_graphviz(model, feature_names=["only_one"]) From 6627ac17b9e65c4e69fe3c45edcb6af83953d27c Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 16 Oct 2026 16:10:51 -0400 Subject: [PATCH 2/4] Draw plot_tree-style panels in export_graphviz nodes Internal nodes now show their shape function shaded by destination child, with an optional X histogram (pair nodes: routing heatmap with marginal histograms). Edges are coloured by child and carry the export_text rule as a tooltip instead of a label. --- docs/_static/export_graphviz_iris.svg | 349 ++++++++++++++++++++ docs/api/plotting.rst | 32 +- sgtlearn/_export.py | 450 ++++++++++++++++++++++---- tests/test_export_text.py | 153 ++++++--- 4 files changed, 883 insertions(+), 101 deletions(-) create mode 100644 docs/_static/export_graphviz_iris.svg diff --git a/docs/_static/export_graphviz_iris.svg b/docs/_static/export_graphviz_iris.svg new file mode 100644 index 0000000..f6d1780 --- /dev/null +++ b/docs/_static/export_graphviz_iris.svg @@ -0,0 +1,349 @@ + + + + + + +Tree + + + +0 + +petal width +samples = 150 +value = [50.00, 50.00, 50.00] +class = setosa + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0.10  |  0.80, 1.65, 2.05  |  2.50 + + + + +1 + +samples = 50 +value = [50.00, 0.00, 0.00] +class = setosa + + + +0->1 + + + + + + + + +2 + +petal length +samples = 75 +value = [0.00, 48.00, 27.00] +class = versicolor + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +3.00  |  4.45, 4.55, 4.95  |  6.90 + + + + +0->2 + + + + + + + + +3 + +sepal length +samples = 25 +value = [0.00, 2.00, 23.00] +class = virginica + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +4.90 to 7.90 + + + + +0->3 + + + + + + + + +4 + +samples = 7 +value = [0.00, 7.00, 0.00] +class = versicolor + + + +2->4 + + + + + + + + +5 + +samples = 40 +value = [0.00, 40.00, 0.00] +class = versicolor + + + +2->5 + + + + + + + + +6 + +samples = 28 +value = [0.00, 1.00, 27.00] +class = virginica + + + +2->6 + + + + + + + + +7 + +samples = 25 +value = [0.00, 2.00, 23.00] +class = virginica + + + +3->7 + + + + + + + + +8 + +samples = 0 +value = [0.00, 0.00, 0.00] +class = setosa + + + +3->8 + + + + + + + + diff --git a/docs/api/plotting.rst b/docs/api/plotting.rst index 447bf10..0d3653f 100644 --- a/docs/api/plotting.rst +++ b/docs/api/plotting.rst @@ -42,15 +42,31 @@ export_graphviz .. autofunction:: export_graphviz -Edge labels use the same routing rules as ``export_text``, so the two exports -cannot disagree. ``export_graphviz`` writes plain DOT text and needs no extra -package. To render it, install the ``graphviz`` Python package and the Graphviz -binaries: +``export_graphviz`` draws the same panels as ``plot_tree`` with Graphviz +HTML-like labels. Each internal node shows its shape function, shaded by the +child each region routes to, and with ``X`` a histogram of the samples that +reach it (a heatmap with marginal histograms for pair nodes). Edges are colored +by child; hover over an edge in a rendered SVG to see the ``export_text`` rule +for that child. The output is plain DOT text and needs no extra package. To +render it, install the ``graphviz`` Python package and the Graphviz binaries: .. code-block:: python import graphviz - from sgtlearn import export_graphviz - - dot = export_graphviz(model, feature_names=feature_names, filled=True) - graphviz.Source(dot).render("tree", format="png") + from sklearn.datasets import load_iris + from sgtlearn import SGTClassifier, export_graphviz + + iris = load_iris() + model = SGTClassifier(num_partitions=3, max_depth=2, random_state=0) + model.fit(iris.data, iris.target) + dot = export_graphviz( + model, + X=iris.data, + feature_names=[n.replace(" (cm)", "") for n in iris.feature_names], + class_names=iris.target_names, + ) + graphviz.Source(dot).render("tree", format="svg") + +.. image:: ../_static/export_graphviz_iris.svg + :alt: An SGT on iris rendered from export_graphviz, with a shape-function + histogram in each internal node. diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index a951757..a9cd39d 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -6,11 +6,12 @@ pair nodes) plus histograms of ``X`` when it is supplied, and every leaf is drawn as colored text with the predicted class / value. ``export_text`` returns the same tree as indented routing rules, and -``export_graphviz`` as Graphviz DOT source. +``export_graphviz`` as Graphviz DOT source with the same panels. """ from __future__ import annotations +import html from collections.abc import Callable, Sequence from typing import Any, cast @@ -103,20 +104,25 @@ def export_graphviz( estimator: Any, out_file: Any = None, *, + X: Any = None, feature_names: Sequence[str] | None = None, class_names: Sequence[Any] | None = None, filled: bool = False, decimals: int = 2, ) -> str | None: - """Export the fitted tree in Graphviz DOT format. - - Every node shows its routing feature (internal nodes), ``samples`` and - ``value``; classifier nodes also show the majority ``class``. Regressor - internal nodes carry no ``value``, so they show only ``samples``. An - internal node with ``K`` children has ``K`` out-edges, each labelled with - the inputs routed to that child, using the same rules as ``export_text``. - - Producing DOT needs no extra package. To render it, install the + """Export the fitted tree in Graphviz DOT format, with ``plot_tree`` panels. + + Each internal node is drawn as a panel of its shape function, like + ``plot_tree``: the routing bins are shaded by the child they route to + (a heatmap for pair nodes). When ``X`` is given, the panel also shows a + histogram of the samples that reach the node, with each bar colored by + its child (marginal histograms for pair nodes), plus a ``NaN`` column + when any of them are missing. Edges are colored by child. Each edge's + tooltip holds the inputs routed to that child, using the same rules as + ``export_text``. + + Panels are drawn with Graphviz HTML-like labels, so the DOT source is + self-contained and needs no extra package. To render it, install the ``graphviz`` Python package and binaries and use ``graphviz.Source(dot)``. Parameters @@ -125,12 +131,14 @@ def export_graphviz( Fitted estimator to export. out_file : str, path-like or file object, optional Where to write the DOT source. When omitted, it is returned instead. + X : array-like of shape (n_samples, n_features), optional + Samples to histogram in each internal node, as in ``plot_tree``. feature_names : sequence of str, optional Names for the input columns, as in ``export_text``. class_names : sequence, optional Classifier only: class names, as in ``export_text``. filled : bool, default=False - Fill nodes by majority class (classifier, shaded by purity) or by + Fill leaves by majority class (classifier, shaded by purity) or by value (regressor, shaded from the smallest to the largest value). decimals : int, default=2 Digits shown for thresholds and values. @@ -143,6 +151,8 @@ def export_graphviz( tree, feat_names, group_label, names_per_output = _prepare_export( estimator, feature_names, class_names, "export_graphviz" ) + X_arr = _validate_X(estimator, X) + reach = _route_samples(tree, X_arr) if X_arr is not None else {} is_classifier = isinstance(estimator, SGTClassifier) values = [v for n in tree["nodes"] if (v := _first_value(n)) is not None] v_lo, v_hi = (min(values), max(values)) if values else (0.0, 0.0) @@ -155,48 +165,51 @@ def fill(node: dict) -> str: return "#ffffff" p = np.sort(counts / counts.sum())[::-1] second = p[1] if len(p) > 1 else 0.0 - color = _DEFAULT_PALETTE_COLORS[ - int(np.argmax(counts)) % len(_DEFAULT_PALETTE_COLORS) - ] - return _blend_white(color, (p[0] - second) / (1 - second)) + return _blend_white( + _partition_color(int(np.argmax(counts))), (p[0] - second) / (1 - second) + ) value = _first_value(node) if value is None: return "#ffffff" alpha = 1.0 if v_hi == v_lo else (value - v_lo) / (v_hi - v_lo) - return _blend_white(_DEFAULT_PALETTE_COLORS[0], alpha) + return _blend_white(_partition_color(0), alpha) - style = '"filled, rounded"' if filled else '"rounded"' - lines = [ - "digraph Tree {", - f'node [shape=box, style={style}, color="black", fontname="helvetica"] ;', - 'edge [fontname="helvetica"] ;', - ] - for node in tree["nodes"]: - label = ( - [] - if node["is_leaf"] - else [_node_feature_label(node, feat_names, group_label)] - ) - label.append(f"samples = {node['n_samples']}") + def stats(node: dict) -> list[str]: + lines = [f"samples = {node['n_samples']}"] if is_classifier: counts = node["class_counts"] - label.append( + lines.append( "value = " + _per_output([_fmt_list(c, decimals) for c in counts]) ) - label.append("class = " + _majority_class(counts, names_per_output)) + lines.append("class = " + _majority_class(counts, names_per_output)) elif isinstance(node["value"], (list, tuple)): - label.append(f"value = {_fmt_list(node['value'], decimals)}") + lines.append(f"value = {_fmt_list(node['value'], decimals)}") elif node["value"] is not None: - label.append(f"value = {node['value']:.{decimals}f}") - attrs = f'label="{_dot_escape(chr(10).join(label))}"' - if filled: - attrs += f', fillcolor="{fill(node)}"' - lines.append(f"{node['id']} [{attrs}] ;") + lines.append(f"value = {node['value']:.{decimals}f}") + return lines + + leaf_style = '"filled, rounded"' if filled else '"rounded"' + lines = [ + "digraph Tree {", + 'node [fontname="helvetica"] ;', + "edge [penwidth=2] ;", + ] + for node in tree["nodes"]: + nid = node["id"] if node["is_leaf"]: + attrs = f'shape=box, style={leaf_style}, label="{_dot_escape(chr(10).join(stats(node)))}"' + if filled: + attrs += f', fillcolor="{fill(node)}"' + lines.append(f"{nid} [{attrs}] ;") continue - for region, cid in zip(_child_regions(node), node["children"]): - text = _format_region(region, feat_names, group_label, decimals) - lines.append(f'{node["id"]} -> {cid} [label="{_dot_escape(text)}"] ;') + rows = X_arr[reach[nid]] if X_arr is not None else None + label = _panel_label(node, rows, feat_names, group_label, stats(node), decimals) + lines.append(f"{nid} [shape=plain, label=<{label}>] ;") + for k, (region, cid) in enumerate(zip(_child_regions(node), node["children"])): + tip = _dot_escape(_format_region(region, feat_names, group_label, decimals)) + lines.append( + f'{nid} -> {cid} [color="{_partition_color(k)}", tooltip="{tip}"] ;' + ) lines.append("}") dot = "\n".join(lines) + "\n" if out_file is None: @@ -209,6 +222,345 @@ def fill(node: dict) -> str: return None +#: Panel geometry for ``export_graphviz`` (pixels and cell counts). +_PANEL_BAR_ROWS = 10 # vertical resolution of a histogram; bars use 8, 2 are headroom +_PANEL_ROW_PX = 5 +_PANEL_FINE_COLS = 80 # columns across a continuous axis +_PANEL_FINE_PX = 2 +_PANEL_HIST_BINS = 20 # histogram bars across a continuous axis +_PANEL_CAT_PX = 16 # one categorical column or row +_PANEL_PAIR_CELLS = 24 # cells across a continuous pair axis +_PANEL_PAIR_PX = 4 +_PANEL_MARGIN_CELLS = 6 # depth of a pair node's marginal histogram +_PANEL_NAN_PX = 10 +_PANEL_GAP_PX = 4 +_PANEL_TINT = 0.4 # shading strength of the routing background +_PANEL_MARGIN_COLOR = "#9e9e9e" + +#: One panel cell: (width px, height px, fill color or None, text or None). +_Cell = tuple[int, int, "str | None", "str | None"] + + +def _partition_color(k: int) -> str: + return _DEFAULT_PALETTE_COLORS[int(k) % len(_DEFAULT_PALETTE_COLORS)] + + +def _panel_label( + node: dict, + rows: np.ndarray | None, + feat_names: list[str], + group_label: Callable[[Sequence[int]], str], + stats: list[str], + decimals: int, +) -> str: + """HTML-like label: feature name, node stats, shape-function panel, captions.""" + if _is_pair_node(node): + grid, captions = _pair_panel(node, rows, group_label, decimals) + elif _is_categorical_node(node): + grid, captions = _categorical_panel(node, rows, feat_names) + else: + grid, captions = _univariate_panel(node, rows, decimals) + title = html.escape(_node_feature_label(node, feat_names, group_label)) + parts = [f"{title}"] + parts += [_html_text_row(s, 10) for s in stats] + parts.append(f"{_html_grid(grid)}") + parts += [_html_text_row(c, 9) for c in captions] + return ( + '' + "".join(parts) + "
" + ) + + +def _html_text_row(text: str, size: int) -> str: + return f'{html.escape(text)}' + + +def _html_grid(grid: list[list[_Cell]]) -> str: + """Render rows of cells as a borderless Graphviz table. + + Every row must have the same column widths. Runs of same-colored cells + become one spanning cell; a 1px ruler row pins each column's width, since + Graphviz otherwise spreads spanning cells unevenly. + """ + width = sum(w for w, *_ in grid[0]) + height = sum(row[0][1] for row in grid) + 1 + out = [ + ( + '' + ) + ] + for row in grid: + out.append("") + j = 0 + while j < len(row): + w, h, color, text = row[j] + if text is not None: + body = f'{html.escape(text)}' + out.append(f'') + j += 1 + continue + # Merge a run of same-colored cells into one spanning cell. + span = 1 + while j + span < len(row) and row[j + span][2:] == (color, None): + w += row[j + span][0] + span += 1 + bg = f' BGCOLOR="{color}"' if color else "" + out.append(f'') + j += span + out.append("") + ruler = "".join(f'' for w, *_ in grid[0]) + out.append(f"{ruler}
{body}
") + return "".join(out) + + +def _bar_rows( + widths: list[int], parts: list[int | None], counts: list[int] | None +) -> list[list[_Cell]]: + """Histogram over a shaded routing background, one column per entry. + + ``parts[j]`` is the child of column ``j`` (None leaves a blank gap). With + ``counts`` None, only the background is drawn. + """ + top = max(counts or [0]) or 1 + heights = [int(np.ceil(c / top * (_PANEL_BAR_ROWS - 2))) for c in counts or []] + if counts is None: # background only: one row of full panel height + h = _PANEL_BAR_ROWS * _PANEL_ROW_PX + return [[(w, h, _tint(k), None) for w, k in zip(widths, parts)]] + grid = [] + for r in range(_PANEL_BAR_ROWS, 0, -1): + row: list[_Cell] = [] + for w, k, n in zip(widths, parts, heights): + color = _partition_color(k) if k is not None and n >= r else _tint(k) + row.append((w, _PANEL_ROW_PX, color, None)) + grid.append(row) + return grid + + +def _tint(k: int | None) -> str | None: + """Light routing-background shade of child ``k`` (None for a gap).""" + return None if k is None else _blend_white(_partition_color(k), _PANEL_TINT) + + +def _axis_extent(finite: np.ndarray, thresholds: list[float]) -> tuple[float, float]: + """Data range when available, else the thresholds padded by one step.""" + if finite.size: + lo, hi = float(finite.min()), float(finite.max()) + elif thresholds: + step = (thresholds[-1] - thresholds[0]) / max(len(thresholds) - 1, 1) or 1.0 + lo, hi = thresholds[0] - step, thresholds[-1] + step + else: + lo, hi = 0.0, 1.0 + return (lo, hi) if hi > lo else (lo - 0.5, lo + 0.5) + + +def _univariate_panel( + node: dict, rows: np.ndarray | None, decimals: int +) -> tuple[list[list[_Cell]], list[str]]: + thresholds = list(node["thresholds"]) + b2p = _finite_routing_bins(thresholds, list(node["bin_to_partition"])) + values = rows[:, int(node["feature"])] if rows is not None else np.empty(0) + finite = values[np.isfinite(values)] + lo, hi = _axis_extent(finite, thresholds) + edges = np.linspace(lo, hi, _PANEL_FINE_COLS + 1) + parts: list[int | None] = [ + int(b2p[int(np.searchsorted(thresholds, m, side="left"))]) + for m in (edges[:-1] + edges[1:]) / 2 + ] + widths = [_PANEL_FINE_PX] * _PANEL_FINE_COLS + counts: list[int] | None = None + n_missing = values.size - finite.size + if rows is not None: + hist, _ = np.histogram(finite, bins=_PANEL_HIST_BINS, range=(lo, hi)) + per_bar = _PANEL_FINE_COLS // _PANEL_HIST_BINS + counts = [int(hist[j // per_bar]) for j in range(_PANEL_FINE_COLS)] + if n_missing: + parts += [None, int(node.get("nan_prediction_partition", 0))] + widths += [_PANEL_GAP_PX, _PANEL_NAN_PX] + counts += [0, int(n_missing)] + breaks = [ + t for t, a, b in zip(thresholds, b2p[:-1], b2p[1:]) if a != b and lo < t < hi + ] + middle = ", ".join(f"{t:.{decimals}f}" for t in breaks) + caption = ( + f"{lo:.{decimals}f} | {middle} | " if breaks else f"{lo:.{decimals}f} to " + ) + caption += f"{hi:.{decimals}f}" + (" NaN" if n_missing else "") + return _bar_rows(widths, parts, counts), [caption] + + +def _categorical_panel( + node: dict, rows: np.ndarray | None, feat_names: list[str] +) -> tuple[list[list[_Cell]], list[str]]: + """One column per category bin, plus a NaN column when ``rows`` has missing. + + A bin with no categories is the missing-value bin; it is drawn as the NaN + column instead. + """ + bins = _bin_categories_for_node(node) + keep = [b for b, cats in enumerate(bins) if cats] + b2p = [int(node["bin_to_partition"][b]) for b in keep] + all_labels = _bin_labels_for_categorical_node(node, feat_names) + labels: list[str | None] = [all_labels[b] for b in keep] + parts: list[int | None] = list(b2p) + widths = [_PANEL_CAT_PX] * len(b2p) + counts: list[int] | None = None + if rows is not None: + cols = [int(c) for c in node.get("features") or []] + counts = [0] * len(b2p) + n_missing = 0 + for row in rows: + active = _active_onehot_column(row, cols) + if active is None: + n_missing += 1 + else: + counts[keep.index(_categorical_bin_index(active, bins))] += 1 + if n_missing: + parts += [None, int(node.get("nan_prediction_partition", 0))] + widths += [_PANEL_GAP_PX, _PANEL_NAN_PX] + counts += [0, n_missing] + labels += [None, "NaN"] + # Widen columns to fit their 8pt labels (~5px per character). + widths = [max(w, 5 * len(t or "") + 4) for w, t in zip(widths, labels)] + grid = _bar_rows(widths, parts, counts) + grid.append([(w, 12, None, text) for w, text in zip(widths, labels)]) + return grid, [] + + +def _pair_axis_layout( + node: dict, a: int, values: np.ndarray | None, rows: np.ndarray | None +) -> tuple[list[object], list[int], tuple[float, float] | None]: + """Cell values along one pair axis, then ``"gap"`` and ``None`` (missing) if shown.""" + axis = node["pair_axes"][a] + if axis["kind"] == "categorical": + cells: list[object] = [int(c) for c in axis["categories"]] + sizes = [_PANEL_CAT_PX] * len(cells) + extent = None + else: + cuts = _pair_axis_cells(node, a, axis, values, include_missing=False) + extent = (cuts[0][0], cuts[-1][1]) + edges = np.linspace(*extent, _PANEL_PAIR_CELLS + 1) + cells = list((edges[:-1] + edges[1:]) / 2) + sizes = [_PANEL_PAIR_PX] * len(cells) + has_missing = rows is None or any(_pair_axis_missing(r, axis) for r in rows) + if has_missing: + cells += ["gap", None] + sizes += [_PANEL_GAP_PX, _PANEL_NAN_PX] + return cells, sizes, extent + + +def _pair_cell_index(row: np.ndarray, axis: dict, cells: list[object]) -> int: + """Index in ``cells`` of the cell that ``row`` falls in along ``axis``.""" + if _pair_axis_missing(row, axis): + return cells.index(None) + cols = [int(c) for c in axis["columns"]] + if axis["kind"] == "categorical": + return cells.index(_active_onehot_column(row, cols)) + mids = np.array([c for c in cells if c not in ("gap", None)], dtype=float) + return int(np.argmin(np.abs(mids - row[cols[0]]))) + + +def _pair_panel( + node: dict, + rows: np.ndarray | None, + group_label: Callable[[Sequence[int]], str], + decimals: int, +) -> tuple[list[list[_Cell]], list[str]]: + """Routing heatmap (x = first axis, y = second, high values on top). + + Cells holding samples are drawn in full color; with ``rows``, gray + marginal histograms sit above (x) and to the right (y). + """ + x_axis, y_axis = node["pair_axes"] + + def axis_values(axis: dict) -> np.ndarray | None: + if rows is None or axis["kind"] != "continuous": + return None + return rows[:, int(axis["columns"][0])] + + xs, x_px, x_ext = _pair_axis_layout(node, 0, axis_values(x_axis), rows) + ys, y_px, y_ext = _pair_axis_layout(node, 1, axis_values(y_axis), rows) + if y_axis["kind"] == "continuous": # high values on top, missing row last + n = _PANEL_PAIR_CELLS + ys[:n], y_px[:n] = ys[:n][::-1], y_px[:n][::-1] + x_counts, y_counts = [0] * len(xs), [0] * len(ys) + occupied = set() + for row in rows if rows is not None else []: + i = _pair_cell_index(row, x_axis, xs) + j = _pair_cell_index(row, y_axis, ys) + x_counts[i] += 1 + y_counts[j] += 1 + occupied.add((i, j)) + margin = rows is not None and len(rows) > 0 + x_bars = _margin_lengths(x_counts) + y_bars = _margin_lengths(y_counts) + grid: list[list[_Cell]] = [] + if margin: + for r in range(_PANEL_MARGIN_CELLS, 0, -1): + cells: list[_Cell] = [ + (w, _PANEL_PAIR_PX, _PANEL_MARGIN_COLOR if n >= r else None, None) + for w, n in zip(x_px, x_bars) + ] + cells += [(_PANEL_PAIR_PX, _PANEL_PAIR_PX, None, None)] * ( + _PANEL_MARGIN_CELLS + 1 + ) + grid.append(cells) + for j, (yv, h) in enumerate(zip(ys, y_px)): + cells = [] + for i, (xv, w) in enumerate(zip(xs, x_px)): + color = None + if "gap" not in (xv, yv): + k = _route_pair_partition( + node, _pair_cell_row(node, x_axis, xv, y_axis, yv) + ) + color = _partition_color(k) if (i, j) in occupied else _tint(k) + cells.append((w, h, color, None)) + if margin: + cells.append((_PANEL_PAIR_PX, h, None, None)) + cells += [ + ( + _PANEL_PAIR_PX, + h, + _PANEL_MARGIN_COLOR if m < y_bars[j] else None, + None, + ) + for m in range(_PANEL_MARGIN_CELLS) + ] + grid.append(cells) + + def caption(side: str, axis: dict, ext: tuple[float, float] | None) -> str: + name = group_label(axis["columns"]) + if ext is None: + cats = ", ".join(group_label([c]) for c in axis["categories"]) + return f"{side}-axis: {name} in {{{cats}}}" + return f"{side}-axis: {name}, {ext[0]:.{decimals}f} to {ext[1]:.{decimals}f}" + + return grid, [caption("x", x_axis, x_ext), caption("y", y_axis, y_ext)] + + +def _margin_lengths(counts: list[int]) -> list[int]: + """Marginal bar lengths in cells; any nonzero count gets at least one.""" + top = max(counts, default=0) or 1 + return [int(np.ceil(c / top * _PANEL_MARGIN_CELLS)) for c in counts] + + +def _validate_X(estimator: Any, X: Any) -> np.ndarray | None: + if X is None: + return None + X_arr = np.asarray(X, dtype=np.float64) + if X_arr.ndim != 2: + raise ValueError( + f"X must be 2-D (n_samples, n_features); got shape {X_arr.shape}" + ) + n_features = estimator.n_features_in_ or 0 + if X_arr.shape[1] != n_features: + raise ValueError( + f"X has {X_arr.shape[1]} features, but estimator was fit on " + f"{n_features}" + ) + return X_arr + + def _prepare_export( estimator: Any, feature_names: Sequence[str] | None, @@ -1606,22 +1958,8 @@ class index (``0``, ``1``, ...). For multi-output classifiers, pass tree = estimator.tree_export() - reach: dict[int, Any] = {} - if X is not None: - X_arr = np.asarray(X, dtype=np.float64) - if X_arr.ndim != 2: - raise ValueError( - f"X must be 2-D (n_samples, n_features); got shape {X_arr.shape}" - ) - n_features = estimator.n_features_in_ or 0 - if X_arr.shape[1] != n_features: - raise ValueError( - f"X has {X_arr.shape[1]} features, but estimator was fit on " - f"{n_features}" - ) - reach = _route_samples(tree, X_arr) - else: - X_arr = None + X_arr = _validate_X(estimator, X) + reach: dict[int, Any] = _route_samples(tree, X_arr) if X_arr is not None else {} if ax is None: n_leaves_est = sum(1 for n in tree["nodes"] if n["is_leaf"]) or 1 diff --git a/tests/test_export_text.py b/tests/test_export_text.py index fe9645a..0caa916 100644 --- a/tests/test_export_text.py +++ b/tests/test_export_text.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from xml.etree import ElementTree import numpy as np import pytest @@ -12,6 +13,7 @@ import sgtlearn from sgtlearn import SGTClassifier, SGTRegressor, export_graphviz, export_text from sgtlearn._export import ( + _DEFAULT_PALETTE_COLORS, _active_onehot_column, _child_regions, _format_region, @@ -244,25 +246,44 @@ def test_export_text_validates_inputs(): export_text(model, feature_names=["only_one"]) -_DOT_NODE = re.compile( - r'^(\d+) \[label="((?:[^"\\]|\\.)*)"(?:, fillcolor="(#[0-9a-f]{6})")?\] ;$' +_DOT_LEAF = re.compile( + r'^(\d+) \[shape=box, style="[a-z, ]+", label="((?:[^"\\]|\\.)*)"' + r'(?:, fillcolor="(#[0-9a-f]{6})")?\] ;$' ) -_DOT_EDGE = re.compile(r'^(\d+) -> (\d+) \[label="((?:[^"\\]|\\.)*)"\] ;$') +_DOT_PANEL = re.compile(r"^(\d+) \[shape=plain, label=<(.*)>\] ;$") +_DOT_EDGE = re.compile( + r'^(\d+) -> (\d+) \[color="(#[0-9A-Fa-f]{6})", tooltip="((?:[^"\\]|\\.)*)"\] ;$' +) +_PALETTE = [c.lower() for c in _DEFAULT_PALETTE_COLORS] + +def _parse_dot(dot: str) -> tuple[dict, dict, list]: + """Parse ``export_graphviz`` output; every body line must be a node or edge. -def _parse_dot(dot: str) -> tuple[dict, list]: - """Parse ``export_graphviz`` output; every body line must be a node or edge.""" + Returns leaves ``{id: (label, fillcolor)}``, panels ``{id: XML element}`` + and edges ``[(parent, child, color, tooltip)]``. + """ lines = dot.splitlines() assert lines[0] == "digraph Tree {" and lines[-1] == "}" - nodes, edges = {}, [] + leaves, panels, edges = {}, {}, [] for line in lines[3:-1]: - if m := _DOT_NODE.match(line): - nodes[int(m[1])] = (m[2].replace("\\n", "\n"), m[3]) + if m := _DOT_LEAF.match(line): + leaves[int(m[1])] = (m[2].replace("\\n", "\n"), m[3]) + elif m := _DOT_PANEL.match(line): + panels[int(m[1])] = ElementTree.fromstring(m[2]) # well-formed HTML label else: m = _DOT_EDGE.match(line) assert m, line - edges.append((int(m[1]), int(m[2]), m[3].replace('\\"', '"'))) - return nodes, edges + edges.append((int(m[1]), int(m[2]), m[3].lower(), m[4].replace('\\"', '"'))) + return leaves, panels, edges + + +def _panel_colors(panel: ElementTree.Element) -> set[str]: + return {td.get("BGCOLOR").lower() for td in panel.iter("TD") if td.get("BGCOLOR")} + + +def _panel_text(panel: ElementTree.Element) -> list[str]: + return [t.strip() for t in panel.itertext() if t.strip()] @pytest.mark.parametrize( @@ -270,20 +291,20 @@ def _parse_dot(dot: str) -> tuple[dict, list]: [_univariate_regressor, _categorical_classifier, _pair_classifier, _pair_regressor], ) def test_export_graphviz_matches_export_text(make): - model, _ = make() + model, X = make() tree = model.tree_export() - nodes, edges = _parse_dot(export_graphviz(model, filled=True)) - assert set(nodes) == {n["id"] for n in tree["nodes"]} - assert all(color for _, color in nodes.values()) - assert sorted((a, b) for a, b, _ in edges) == sorted( + by_id = {n["id"]: n for n in tree["nodes"]} + leaves, panels, edges = _parse_dot(export_graphviz(model, X=X, filled=True)) + assert set(leaves) == {n["id"] for n in tree["nodes"] if n["is_leaf"]} + assert set(panels) == {n["id"] for n in tree["nodes"] if not n["is_leaf"]} + assert all(color for _, color in leaves.values()) + for parent, child, color, _ in edges: + k = by_id[parent]["children"].index(child) + assert color == _PALETTE[k % len(_PALETTE)] + assert sorted((a, b) for a, b, *_ in edges) == sorted( (n["id"], c) for n in tree["nodes"] for c in n["children"] ) - child_lines = [ - line.split("|--- ", 1)[1] - for line in export_text(model).splitlines() - if "|--- class:" not in line and "|--- value:" not in line - ] - by_id = {n["id"]: n for n in tree["nodes"]} + dfs: list[tuple[int, int]] = [] def walk(nid: int) -> None: # export_text's line order @@ -292,35 +313,91 @@ def walk(nid: int) -> None: # export_text's line order walk(cid) walk(tree["root_index"]) - labels = {(a, b): label for a, b, label in edges} - assert [labels[e] for e in dfs] == child_lines + tips = {(a, b): tip for a, b, _, tip in edges} + child_lines = [ + line.split("|--- ", 1)[1] + for line in export_text(model).splitlines() + if "|--- class:" not in line and "|--- value:" not in line + ] + assert [tips[e] for e in dfs] == child_lines -def test_export_graphviz_node_labels_and_out_file(tmp_path): - X = np.arange(4.0)[:, None] +@pytest.mark.parametrize( + "make", + [_univariate_regressor, _categorical_classifier, _pair_classifier, _pair_regressor], +) +def test_export_graphviz_histograms_only_with_X(make): + """Full child colors mark histogram bars or occupied cells; tints the routing.""" + model, X = make() + _, with_x, _ = _parse_dot(export_graphviz(model, X=X)) + _, without_x, _ = _parse_dot(export_graphviz(model)) + reached = {n["id"] for n in model.tree_export()["nodes"] if n["n_samples"]} + for nid, panel in with_x.items(): + if nid in reached: + assert _panel_colors(panel) & set(_PALETTE), nid + for panel in without_x.values(): + colors = _panel_colors(panel) + assert colors and not colors & set(_PALETTE) + assert "#9e9e9e" not in colors # no pair marginals + + +def test_export_graphviz_panels_and_out_file(tmp_path): + X = np.array([[0.0], [1.0], [2.0], [3.0], [np.nan]]) model = SGTRegressor(max_leaf_nodes=2, tao_n_runs=0, random_state=0).fit( - X, np.array([0.0, 0.0, 1.0, 1.0]) + X, np.array([0.0, 0.0, 1.0, 1.0, 0.0]) ) - dot = export_graphviz(model, feature_names=['x "q"']) - nodes, edges = _parse_dot(dot) - assert nodes[0] == ('x \\"q\\"\nsamples = 4', None) - assert sorted(label for label, _ in list(nodes.values())[1:]) == [ - "samples = 2\nvalue = 0.00", + dot = export_graphviz(model, X=X, feature_names=['x "q"']) + leaves, panels, edges = _parse_dot(dot) + assert _panel_text(panels[0]) == [ + 'x "q"', + "samples = 5", + "0.00 | 1.50 | 3.00 NaN", + ] + assert sorted(label for label, _ in leaves.values()) == [ "samples = 2\nvalue = 1.00", + "samples = 3\nvalue = 0.00", ] - assert [label for *_, label in edges] == [ + assert [tip for *_, tip in edges] == [ 'x "q" <= 1.50 or x "q" is missing', 'x "q" > 1.50', ] + assert ( + _panel_text(_parse_dot(export_graphviz(model))[1][0])[-1] + == "0.50 | 1.50 | 2.50" + ) path = tmp_path / "tree.dot" - assert export_graphviz(model, out_file=path, feature_names=['x "q"']) is None - assert path.read_text() == dot + out = export_graphviz(model, out_file=path, X=X, feature_names=['x "q"']) + assert out is None and path.read_text() == dot clf = SGTClassifier(max_leaf_nodes=2, tao_n_runs=0, random_state=0).fit( - X, ["a", "a", "b", "b"] + X[:4], ["a", "a", "b", "b"] ) - root = _parse_dot(export_graphviz(clf, class_names=["neg", "pos"]))[0][0][0] - assert root == "X[0]\nsamples = 4\nvalue = [2.00, 2.00]\nclass = neg" + root = _parse_dot(export_graphviz(clf, class_names=["neg", "pos"]))[1][0] + assert _panel_text(root)[:4] == [ + "X[0]", + "samples = 4", + "value = [2.00, 2.00]", + "class = neg", + ] + + +def test_export_graphviz_categorical_and_pair_panels(): + model, X = _pair_classifier() + panels = _parse_dot( + export_graphviz(model, X=X, feature_names=["x", "red", "blue"]) + )[1] + text = _panel_text(panels[model.tree_export()["root_index"]]) + assert text[0] == "x & [red, blue]" + assert text[-2:] == [ + "x-axis: x, -1.00 to 1.00", + "y-axis: [red, blue] in {red, blue}", + ] + + model, X = _categorical_classifier() + root = _parse_dot(export_graphviz(model, X=X))[1][0] + labels = _panel_text(root)[4:] + assert labels.count("NaN") == 1 # the empty category bin is the NaN column + assert len(labels) == 5 def test_export_graphviz_is_public_and_validates_inputs(): @@ -330,3 +407,5 @@ def test_export_graphviz_is_public_and_validates_inputs(): model = SGTRegressor(tao_n_runs=0).fit(np.zeros((4, 2)), np.arange(4.0)) with pytest.raises(ValueError, match="feature_names"): export_graphviz(model, feature_names=["only_one"]) + with pytest.raises(ValueError, match="features"): + export_graphviz(model, X=np.zeros((3, 3))) From 082c57e26f08d99cb49dd5692560789ec7320b0f Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 16 Oct 2026 16:19:11 -0400 Subject: [PATCH 3/4] Align export_graphviz axis labels with the panel Each colour switch in a univariate panel now gets a tick, with its threshold label centred under it (overlapping labels drop to a further row), and the NaN label sits under the NaN column. --- docs/_static/export_graphviz_iris.svg | 470 +++++++++++++------------- sgtlearn/_export.py | 87 ++++- tests/test_export_text.py | 63 +++- 3 files changed, 375 insertions(+), 245 deletions(-) diff --git a/docs/_static/export_graphviz_iris.svg b/docs/_static/export_graphviz_iris.svg index f6d1780..b4a3534 100644 --- a/docs/_static/export_graphviz_iris.svg +++ b/docs/_static/export_graphviz_iris.svg @@ -4,259 +4,273 @@ - - + + Tree - + 0 - -petal width -samples = 150 -value = [50.00, 50.00, 50.00] -class = setosa - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -0.10  |  0.80, 1.65, 2.05  |  2.50 - + +petal width +samples = 150 +value = [50.00, 50.00, 50.00] +class = setosa + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0.10 +0.80 +1.65 +2.05 + 1 - -samples = 50 -value = [50.00, 0.00, 0.00] -class = setosa + +samples = 50 +value = [50.00, 0.00, 0.00] +class = setosa 0->1 - - + + 2 - -petal length -samples = 75 -value = [0.00, 48.00, 27.00] -class = versicolor - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -3.00  |  4.45, 4.55, 4.95  |  6.90 - + +petal length +samples = 75 +value = [0.00, 48.00, 27.00] +class = versicolor + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +3.00 +4.45 +6.90 +4.55 +4.95 + 0->2 - - + + 3 - -sepal length -samples = 25 -value = [0.00, 2.00, 23.00] -class = virginica - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -4.90 to 7.90 - + +sepal length +samples = 25 +value = [0.00, 2.00, 23.00] +class = virginica + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +4.90 +7.90 + 0->3 - - + + @@ -272,8 +286,8 @@ 2->4 - - + + @@ -289,8 +303,8 @@ 2->5 - - + + @@ -306,8 +320,8 @@ 2->6 - - + + @@ -323,8 +337,8 @@ 3->7 - - + + @@ -340,8 +354,8 @@ 3->8 - - + + diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index a9cd39d..d3e0c14 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -232,12 +232,14 @@ def stats(node: dict) -> list[str]: _PANEL_PAIR_CELLS = 24 # cells across a continuous pair axis _PANEL_PAIR_PX = 4 _PANEL_MARGIN_CELLS = 6 # depth of a pair node's marginal histogram -_PANEL_NAN_PX = 10 +_PANEL_NAN_PX = 18 # fits an 8pt "NaN" label _PANEL_GAP_PX = 4 _PANEL_TINT = 0.4 # shading strength of the routing background _PANEL_MARGIN_COLOR = "#9e9e9e" +_PANEL_TICK_COLOR = "#555555" #: One panel cell: (width px, height px, fill color or None, text or None). +#: Text ``""`` continues the text cell to its left, which then spans it. _Cell = tuple[int, int, "str | None", "str | None"] @@ -296,9 +298,14 @@ def _html_grid(grid: list[list[_Cell]]) -> str: while j < len(row): w, h, color, text = row[j] if text is not None: - body = f'{html.escape(text)}' - out.append(f'{body}') - j += 1 + span = 1 + while j + span < len(row) and row[j + span][3] == "": + w += row[j + span][0] + span += 1 + lines = "
".join(html.escape(t) for t in text.split("\n")) + body = f'{lines}' if text else "" + out.append(f'{body}') + j += span continue # Merge a run of same-colored cells into one spanning cell. span = 1 @@ -378,15 +385,71 @@ def _univariate_panel( parts += [None, int(node.get("nan_prediction_partition", 0))] widths += [_PANEL_GAP_PX, _PANEL_NAN_PX] counts += [0, int(n_missing)] - breaks = [ - t for t, a, b in zip(thresholds, b2p[:-1], b2p[1:]) if a != b and lo < t < hi + # Tick and label each visible color switch, with the threshold that + # causes it (the one closest to the column boundary). + breaks = [t for t, a, b in zip(thresholds, b2p[:-1], b2p[1:]) if a != b] + cols = [j for j in range(1, _PANEL_FINE_COLS) if parts[j] != parts[j - 1]] + labels = [ + (j, f"{min(breaks, key=lambda t: abs(t - edges[j])):.{decimals}f}") + for j in cols ] - middle = ", ".join(f"{t:.{decimals}f}" for t in breaks) - caption = ( - f"{lo:.{decimals}f} | {middle} | " if breaks else f"{lo:.{decimals}f} to " - ) - caption += f"{hi:.{decimals}f}" + (" NaN" if n_missing else "") - return _bar_rows(widths, parts, counts), [caption] + rows_out = _axis_rows(labels, f"{lo:.{decimals}f}", f"{hi:.{decimals}f}") + tick: list[_Cell] = [ + (_PANEL_FINE_PX, 3, _PANEL_TICK_COLOR if j in cols else None, None) + for j in range(_PANEL_FINE_COLS) + ] + grid = _bar_rows(widths, parts, counts) + [tick] + rows_out + if len(widths) > _PANEL_FINE_COLS: # NaN column: blank tick, label under it + tick += [(_PANEL_GAP_PX, 3, None, None), (_PANEL_NAN_PX, 3, None, None)] + for i, row in enumerate(rows_out): + nan = "NaN" if i == 0 else None + row += [ + (_PANEL_GAP_PX, row[0][1], None, None), + (_PANEL_NAN_PX, row[0][1], None, nan), + ] + return grid, [] + + +def _axis_rows(labels: list[tuple[int, str]], lo: str, hi: str) -> list[list[_Cell]]: + """Label rows under a continuous axis, each label centered on its column. + + Labels that would overlap move to a further row. The axis ends ``lo`` and + ``hi`` are added at the edges only where they fit in an existing row. + """ + n = _PANEL_FINE_COLS + + def width(text: str) -> int: # columns needed by an 8pt label (~5px/char) + return min(n, -(-(5 * len(text) + 4) // _PANEL_FINE_PX)) + + placed: list[list[tuple[int, int, str]]] = [] # per row: (start, stop, text) + + def place(start: int, text: str, new_row: bool) -> None: + stop = start + width(text) + for row in placed: + if all(stop + 1 <= a or start >= b + 1 for a, b, _ in row): + row.append((start, stop, text)) + return + if new_row: + placed.append([(start, stop, text)]) + + for col, text in labels: + start = min(max(col - width(text) // 2, 0), n - width(text)) + place(start, text, new_row=True) + if not placed: + placed.append([]) + place(0, lo, new_row=False) + place(n - width(hi), hi, new_row=False) + + out = [] + for row in placed: + cells: list[_Cell] = [(_PANEL_FINE_PX, 12, None, None)] * n + for start, stop, text in row: + cells[start] = (_PANEL_FINE_PX, 12, None, text) + cells[start + 1 : stop] = [(_PANEL_FINE_PX, 12, None, "")] * ( + stop - start - 1 + ) + out.append(cells) + return out def _categorical_panel( diff --git a/tests/test_export_text.py b/tests/test_export_text.py index 0caa916..41957cb 100644 --- a/tests/test_export_text.py +++ b/tests/test_export_text.py @@ -351,7 +351,10 @@ def test_export_graphviz_panels_and_out_file(tmp_path): assert _panel_text(panels[0]) == [ 'x "q"', "samples = 5", - "0.00 | 1.50 | 3.00 NaN", + "0.00", + "1.50", + "3.00", + "NaN", ] assert sorted(label for label, _ in leaves.values()) == [ "samples = 2\nvalue = 1.00", @@ -361,10 +364,8 @@ def test_export_graphviz_panels_and_out_file(tmp_path): 'x "q" <= 1.50 or x "q" is missing', 'x "q" > 1.50', ] - assert ( - _panel_text(_parse_dot(export_graphviz(model))[1][0])[-1] - == "0.50 | 1.50 | 2.50" - ) + no_x = _panel_text(_parse_dot(export_graphviz(model))[1][0]) + assert no_x[-3:] == ["0.50", "1.50", "2.50"] path = tmp_path / "tree.dot" out = export_graphviz(model, out_file=path, X=X, feature_names=['x "q"']) assert out is None and path.read_text() == dot @@ -381,6 +382,58 @@ def test_export_graphviz_panels_and_out_file(tmp_path): ] +@pytest.mark.parametrize("make", [_univariate_regressor, _categorical_classifier]) +def test_export_graphviz_axis_labels_line_up(make): + """Ticks sit on color switches, threshold labels center on their tick, and + "NaN" sits under the NaN column.""" + model, X = make() + tree = model.tree_export() + _, panels, _ = _parse_dot(export_graphviz(model, X=X)) + checked = 0 + for node in tree["nodes"]: + if node["is_leaf"] or _is_categorical_node(node) or _is_pair_node(node): + continue + grid = next(t for t in panels[node["id"]].iter("TABLE") if t.get("FIXEDSIZE")) + rows = grid.findall("TR") + tick_rows = [ + i + for i, r in enumerate(rows) + if "#555555" in ElementTree.tostring(r, "unicode") + ] + if not tick_rows: # no visible switch, so nothing to line up + continue + tick_i = tick_rows[0] + top, tick, labels = rows[0], rows[tick_i], rows[tick_i + 1 : -1] + + def cells(row): # (first column, last column, td) + col, out = 0, [] + for td in row.findall("TD"): + n = int(td.get("COLSPAN", 1)) + out.append((col, col + n - 1, td)) + col += n + return out + + def per_column(row): + return [ + td.get("BGCOLOR") for a, b, td in cells(row) for _ in range(a, b + 1) + ] + + n_fine = 80 + bg = per_column(top)[:n_fine] + switches = {j for j in range(1, n_fine) if bg[j] != bg[j - 1]} + ticks = {j for j, c in enumerate(per_column(tick)) if c == "#555555"} + assert ticks == switches + for row_i, row in enumerate(labels): + for a, b, td in cells(row): + text = "".join(td.itertext()) + if text == "NaN": + assert row_i == 0 and (a, b) == cells(top)[-1][:2] + elif text and 0 < a and b < n_fine - 1: # not an axis-end label + assert min(abs((a + b + 1) / 2 - j) for j in ticks) <= 1, text + checked += 1 + assert checked + + def test_export_graphviz_categorical_and_pair_panels(): model, X = _pair_classifier() panels = _parse_dot( From a3f7fd8c025e0763a4915f8f958c21fafacb3558 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 16 Oct 2026 16:25:40 -0400 Subject: [PATCH 4/4] Tick and label pair heatmap axes at their switches Continuous pair axes get a tick at each colour switch (including switches visible only in the other axis's missing margin) with the threshold label beside it: x below the heatmap, y to its left. Categorical axes and missing margins are labelled per cell. --- sgtlearn/_export.py | 245 +++++++++++++++++++++++++++++++------- tests/test_export_text.py | 47 +++++++- 2 files changed, 243 insertions(+), 49 deletions(-) diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index d3e0c14..3b1fb0d 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -257,15 +257,15 @@ def _panel_label( ) -> str: """HTML-like label: feature name, node stats, shape-function panel, captions.""" if _is_pair_node(node): - grid, captions = _pair_panel(node, rows, group_label, decimals) + chart, captions = _pair_panel(node, rows, group_label, decimals) elif _is_categorical_node(node): - grid, captions = _categorical_panel(node, rows, feat_names) + chart, captions = _categorical_panel(node, rows, feat_names) else: - grid, captions = _univariate_panel(node, rows, decimals) + chart, captions = _univariate_panel(node, rows, decimals) title = html.escape(_node_feature_label(node, feat_names, group_label)) parts = [f"{title}"] parts += [_html_text_row(s, 10) for s in stats] - parts.append(f"{_html_grid(grid)}") + parts.append(f"{chart}") parts += [_html_text_row(c, 9) for c in captions] return ( ' tuple[float, fl def _univariate_panel( node: dict, rows: np.ndarray | None, decimals: int -) -> tuple[list[list[_Cell]], list[str]]: +) -> tuple[str, list[str]]: thresholds = list(node["thresholds"]) b2p = _finite_routing_bins(thresholds, list(node["bin_to_partition"])) values = rows[:, int(node["feature"])] if rows is not None else np.empty(0) @@ -393,7 +393,13 @@ def _univariate_panel( (j, f"{min(breaks, key=lambda t: abs(t - edges[j])):.{decimals}f}") for j in cols ] - rows_out = _axis_rows(labels, f"{lo:.{decimals}f}", f"{hi:.{decimals}f}") + rows_out = _axis_rows( + labels, + f"{lo:.{decimals}f}", + f"{hi:.{decimals}f}", + _PANEL_FINE_COLS, + _PANEL_FINE_PX, + ) tick: list[_Cell] = [ (_PANEL_FINE_PX, 3, _PANEL_TICK_COLOR if j in cols else None, None) for j in range(_PANEL_FINE_COLS) @@ -407,19 +413,20 @@ def _univariate_panel( (_PANEL_GAP_PX, row[0][1], None, None), (_PANEL_NAN_PX, row[0][1], None, nan), ] - return grid, [] + return _html_grid(grid), [] -def _axis_rows(labels: list[tuple[int, str]], lo: str, hi: str) -> list[list[_Cell]]: - """Label rows under a continuous axis, each label centered on its column. +def _axis_rows( + labels: list[tuple[int, str]], lo: str, hi: str, n: int, col_px: int +) -> list[list[_Cell]]: + """Label rows under ``n`` axis columns, each label centered on its column. Labels that would overlap move to a further row. The axis ends ``lo`` and ``hi`` are added at the edges only where they fit in an existing row. """ - n = _PANEL_FINE_COLS - def width(text: str) -> int: # columns needed by an 8pt label (~5px/char) - return min(n, -(-(5 * len(text) + 4) // _PANEL_FINE_PX)) + def width(text: str) -> int: # columns needed by an 8pt label + return min(n, -(-_label_px(text) // col_px)) placed: list[list[tuple[int, int, str]]] = [] # per row: (start, stop, text) @@ -442,19 +449,17 @@ def place(start: int, text: str, new_row: bool) -> None: out = [] for row in placed: - cells: list[_Cell] = [(_PANEL_FINE_PX, 12, None, None)] * n + cells: list[_Cell] = [(col_px, 12, None, None)] * n for start, stop, text in row: - cells[start] = (_PANEL_FINE_PX, 12, None, text) - cells[start + 1 : stop] = [(_PANEL_FINE_PX, 12, None, "")] * ( - stop - start - 1 - ) + cells[start] = (col_px, 12, None, text) + cells[start + 1 : stop] = [(col_px, 12, None, "")] * (stop - start - 1) out.append(cells) return out def _categorical_panel( node: dict, rows: np.ndarray | None, feat_names: list[str] -) -> tuple[list[list[_Cell]], list[str]]: +) -> tuple[str, list[str]]: """One column per category bin, plus a NaN column when ``rows`` has missing. A bin with no categories is the missing-value bin; it is drawn as the NaN @@ -484,10 +489,10 @@ def _categorical_panel( counts += [0, n_missing] labels += [None, "NaN"] # Widen columns to fit their 8pt labels (~5px per character). - widths = [max(w, 5 * len(t or "") + 4) for w, t in zip(widths, labels)] + widths = [max(w, _label_px(t)) for w, t in zip(widths, labels)] grid = _bar_rows(widths, parts, counts) grid.append([(w, 12, None, text) for w, text in zip(widths, labels)]) - return grid, [] + return _html_grid(grid), [] def _pair_axis_layout( @@ -528,13 +533,16 @@ def _pair_panel( rows: np.ndarray | None, group_label: Callable[[Sequence[int]], str], decimals: int, -) -> tuple[list[list[_Cell]], list[str]]: +) -> tuple[str, list[str]]: """Routing heatmap (x = first axis, y = second, high values on top). Cells holding samples are drawn in full color; with ``rows``, gray - marginal histograms sit above (x) and to the right (y). + marginal histograms sit above (x) and to the right (y). Continuous axes + get a tick at each color switch with its threshold label (x below, y to + the left); categorical axes label each category. """ x_axis, y_axis = node["pair_axes"] + n = _PANEL_PAIR_CELLS def axis_values(axis: dict) -> np.ndarray | None: if rows is None or axis["kind"] != "continuous": @@ -543,9 +551,30 @@ def axis_values(axis: dict) -> np.ndarray | None: xs, x_px, x_ext = _pair_axis_layout(node, 0, axis_values(x_axis), rows) ys, y_px, y_ext = _pair_axis_layout(node, 1, axis_values(y_axis), rows) - if y_axis["kind"] == "continuous": # high values on top, missing row last - n = _PANEL_PAIR_CELLS + if y_ext is not None: # high values on top, missing row last ys[:n], y_px[:n] = ys[:n][::-1], y_px[:n][::-1] + + def cell_label(value: object) -> str | None: + if value is None: + return "NaN" + return group_label([value]) if isinstance(value, int) else None + + if x_ext is None: # widen category columns to fit their labels + x_px = [max(w, _label_px(cell_label(v))) for w, v in zip(x_px, xs)] + + part = [ + [ + ( + None + if "gap" in (xv, yv) + else _route_pair_partition( + node, _pair_cell_row(node, x_axis, xv, y_axis, yv) + ) + ) + for xv in xs + ] + for yv in ys + ] x_counts, y_counts = [0] * len(xs), [0] * len(ys) occupied = set() for row in rows if rows is not None else []: @@ -554,28 +583,50 @@ def axis_values(axis: dict) -> np.ndarray | None: x_counts[i] += 1 y_counts[j] += 1 occupied.add((i, j)) + + def switches(lines: list[list[int | None]]) -> list[int]: + """Cells ``i`` (of the first ``n``) whose color differs from ``i - 1``.""" + return [i for i in range(1, n) if any(ln[i] != ln[i - 1] for ln in lines)] + + def boundary_label(a: int, cells: list[object], i: int) -> str: + edge = (float(cast(float, cells[i - 1])) + float(cast(float, cells[i]))) / 2 + cuts = _pair_axis_thresholds(node, a) or [edge] + return f"{min(cuts, key=lambda t: abs(t - edge)):.{decimals}f}" + + real_rows = [part[j] for j, yv in enumerate(ys) if yv != "gap"] + x_sw = switches(real_rows) if x_ext is not None else [] + real_cols = [ + [part[j][i] for j in range(len(ys))] for i, xv in enumerate(xs) if xv != "gap" + ] + y_sw = switches(real_cols) if y_ext is not None else [] + margin = rows is not None and len(rows) > 0 x_bars = _margin_lengths(x_counts) y_bars = _margin_lengths(y_counts) + + def right(h: int) -> list[_Cell]: # blank cells under/over the y margin + return ( + [(_PANEL_PAIR_PX, h, None, None)] * (_PANEL_MARGIN_CELLS + 1) + if margin + else [] + ) + + tick_w = 3 grid: list[list[_Cell]] = [] if margin: for r in range(_PANEL_MARGIN_CELLS, 0, -1): - cells: list[_Cell] = [ - (w, _PANEL_PAIR_PX, _PANEL_MARGIN_COLOR if n >= r else None, None) - for w, n in zip(x_px, x_bars) + cells: list[_Cell] = [(tick_w, _PANEL_PAIR_PX, None, None)] + cells += [ + (w, _PANEL_PAIR_PX, _PANEL_MARGIN_COLOR if b >= r else None, None) + for w, b in zip(x_px, x_bars) ] - cells += [(_PANEL_PAIR_PX, _PANEL_PAIR_PX, None, None)] * ( - _PANEL_MARGIN_CELLS + 1 - ) - grid.append(cells) - for j, (yv, h) in enumerate(zip(ys, y_px)): - cells = [] - for i, (xv, w) in enumerate(zip(xs, x_px)): + grid.append(cells + right(_PANEL_PAIR_PX)) + for j, h in enumerate(y_px): + cells = [(tick_w, h, _PANEL_TICK_COLOR if j in y_sw else None, None)] + for i, w in enumerate(x_px): + k = part[j][i] color = None - if "gap" not in (xv, yv): - k = _route_pair_partition( - node, _pair_cell_row(node, x_axis, xv, y_axis, yv) - ) + if k is not None: color = _partition_color(k) if (i, j) in occupied else _tint(k) cells.append((w, h, color, None)) if margin: @@ -591,14 +642,118 @@ def axis_values(axis: dict) -> np.ndarray | None: ] grid.append(cells) - def caption(side: str, axis: dict, ext: tuple[float, float] | None) -> str: - name = group_label(axis["columns"]) - if ext is None: - cats = ", ".join(group_label([c]) for c in axis["categories"]) - return f"{side}-axis: {name} in {{{cats}}}" - return f"{side}-axis: {name}, {ext[0]:.{decimals}f} to {ext[1]:.{decimals}f}" + # x labels under the heatmap. + tail = list(zip(x_px[n:], xs[n:])) if x_ext is not None else [] + if x_ext is not None: + tick = [(tick_w, 3, None, None)] + [ + (w, 3, _PANEL_TICK_COLOR if i in x_sw else None, None) + for i, w in enumerate(x_px) + ] + grid.append(tick + right(3)) + labels = [(i, boundary_label(0, xs, i)) for i in x_sw] + lo, hi = (f"{v:.{decimals}f}" for v in x_ext) + for r, label_row in enumerate(_axis_rows(labels, lo, hi, n, _PANEL_PAIR_PX)): + extra = [(w, 12, None, cell_label(v) if r == 0 else None) for w, v in tail] + grid.append([(tick_w, 12, None, None)] + label_row + extra + right(12)) + else: + cats = [(w, 12, None, cell_label(v)) for w, v in zip(x_px, xs)] + grid.append([(tick_w, 12, None, None)] + cats + right(12)) + + # y labels to the left, centered on their row boundary or row. + top = _PANEL_MARGIN_CELLS * _PANEL_PAIR_PX if margin else 0 + starts = [top + sum(y_px[:j]) for j in range(len(ys) + 1)] + marks: list[tuple[float, str]] = [] + ends: list[tuple[float, str]] = [] + if y_ext is not None: + marks = [(starts[j], boundary_label(1, ys, j)) for j in y_sw] + ends = [ + (starts[0] + 6, f"{y_ext[1]:.{decimals}f}"), + (starts[n] - 6, f"{y_ext[0]:.{decimals}f}"), + ] + for j, yv in enumerate(ys): + text = cell_label(yv) + if text is not None: + marks.append(((starts[j] + starts[j + 1]) / 2, text)) + height = sum(row[0][1] for row in grid) + 1 + side, side_px = _side_labels(marks, ends, height) + width = side_px + sum(w for w, *_ in grid[0]) + chart = ( + '
' + f"{side}
{_html_grid(grid)}
" + ) + names = [ + f"{s}-axis: {group_label(a['columns'])}" for s, a in zip("xy", (x_axis, y_axis)) + ] + return chart, names + + +def _pair_axis_thresholds(node: dict, a: int) -> list[float]: + return sorted( + float(s["threshold"]) + for s in node["pair_inner_tree"] + if not s["is_leaf"] and int(s["axis"]) == a and s["kind"] == "continuous" + ) + + +def _label_px(text: str | None) -> int: + """Width of an 8pt label (~5px per character).""" + return 5 * len(text or "") + 4 - return grid, [caption("x", x_axis, x_ext), caption("y", y_axis, y_ext)] + +def _side_labels( + marks: list[tuple[float, str]], ends: list[tuple[float, str]], height: int +) -> tuple[str, int]: + """Columns of right-aligned 12px labels, each centered on its pixel row. + + Returns the cells' HTML and their total width. + + Overlapping ``marks`` move to a further column (further left); ``ends`` + are added only where they fit in an existing column. + """ + columns: list[list[tuple[int, int, str]]] = [] + + def place(center: float, text: str, new_column: bool) -> None: + start = min(max(round(center - 6), 0), height - 12) + stop = start + 12 + for col in columns: + if all(stop + 1 <= a or start >= b + 1 for a, b, _ in col): + col.append((start, stop, text)) + return + if new_column: + columns.append([(start, stop, text)]) + + for center, text in marks: + place(center, text, new_column=True) + for center, text in ends: + place(center, text, new_column=False) + out, total = [], 0 + for col in reversed(columns): # the first column sits next to the heatmap + width = max(_label_px(t) for *_, t in col) + total += width + cells, cursor = [], 0 + for start, stop, text in sorted(col): + if start > cursor: + cells.append( + f'' + ) + body = f'{html.escape(text)}' + cells.append( + f'' + f"{body}" + ) + cursor = stop + if cursor < height: + cells.append( + f'' + ) + out.append( + '' + + "".join(cells) + + "
" + ) + return "".join(out), total def _margin_lengths(counts: list[int]) -> list[int]: diff --git a/tests/test_export_text.py b/tests/test_export_text.py index 41957cb..b5ce0b3 100644 --- a/tests/test_export_text.py +++ b/tests/test_export_text.py @@ -15,6 +15,7 @@ from sgtlearn._export import ( _DEFAULT_PALETTE_COLORS, _active_onehot_column, + _blend_white, _child_regions, _format_region, _is_categorical_node, @@ -441,10 +442,8 @@ def test_export_graphviz_categorical_and_pair_panels(): )[1] text = _panel_text(panels[model.tree_export()["root_index"]]) assert text[0] == "x & [red, blue]" - assert text[-2:] == [ - "x-axis: x, -1.00 to 1.00", - "y-axis: [red, blue] in {red, blue}", - ] + assert text[-2:] == ["x-axis: x", "y-axis: [red, blue]"] + assert {"red", "blue", "NaN", "-1.00", "1.00"} <= set(text) model, X = _categorical_classifier() root = _parse_dot(export_graphviz(model, X=X))[1][0] @@ -453,6 +452,46 @@ def test_export_graphviz_categorical_and_pair_panels(): assert len(labels) == 5 +def test_export_graphviz_pair_ticks_line_up(): + """Pair heatmap ticks sit on the color switches along each continuous axis, + including switches visible only in the other axis's missing margin.""" + model, X = _pair_regressor() + tree = model.tree_export() + _, panels, _ = _parse_dot(export_graphviz(model, X=X)) + child = {} + for k, color in enumerate(_PALETTE): + child[color] = child[_blend_white(color, 0.4)] = k + + def per_column(row): + return [ + (td.get("BGCOLOR") or "").lower() + for td in row.findall("TD") + for _ in range(int(td.get("COLSPAN", 1))) + ] + + n, checked = 24, 0 # continuous cells per pair axis + for node in tree["nodes"]: + if node["is_leaf"] or not _is_pair_node(node): + continue + grid = next( + t + for t in panels[node["id"]].iter("TABLE") + if t.find("TR/TD[@BGCOLOR]") is not None + ) + rows = [per_column(r) for r in grid.findall("TR")] + heat = [r for r in rows if any(c in child for c in r)] # + missing row + kids = [[child.get(c) for c in r[1:]] for r in heat] # + missing column + x_switch = {i + 1 for i in range(1, n) if any(r[i] != r[i - 1] for r in kids)} + y_switch = {j for j in range(1, n) if kids[j] != kids[j - 1]} + x_ticks = { + i for r in rows if r not in heat for i, c in enumerate(r) if c == "#555555" + } + y_ticks = {j for j, r in enumerate(heat) if r[0] == "#555555"} + assert (x_ticks, y_ticks) == (x_switch, y_switch) + checked += 1 + assert checked + + def test_export_graphviz_is_public_and_validates_inputs(): assert "export_graphviz" in sgtlearn.__all__ with pytest.raises(TypeError):