From a226a3df7c20844374847553acbf90a3d9dd4e79 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 2 Oct 2026 22:02:34 -0400 Subject: [PATCH] Fix plot_tree crash on truncated and multi-output SGTRegressor nodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit plot_tree raised TypeError for an SGTRegressor drawn with max_depth, or fitted on a multi-output y (#82). _draw_leaf_text formatted node['value'] as one float, but tree_export sets value=None on internal nodes and a per-output list on multi-output leaves, and plot_tree draws nodes cut off by max_depth as leaves. Classifiers did not crash because class_counts is filled in on internal nodes too, but multi-output classifiers silently showed only output 0's class. Truncated regressor internals now render as "…". Multi-output leaves render one entry per output as "[a, b]" for both regressors and classifiers; plot_tree resolves class names per output (pass one list per output for multi-output classifiers). The per-output and list formatting helpers move out of export_text so both renderers share them. The plot_tree docstring now states that n= is a rounded weighted count for classifiers and an unweighted count for regressors. No test caught this: the only max_depth test used a classifier, and no plot_tree test fitted a 2-D y. Add a regressor test over single and multi-output, with and without max_depth, plus a multi-output classifier test, and the test-checker blind tests promoted for _draw_leaf_text. Co-Authored-By: Claude Opus 5.5 --- .test-checker/ledger.jsonl | 13 + .../reports/r-20261002-215733-b7b8.json | 41 +++ .../reports/r-20261002-215733-b7b8.md | 38 +++ .test-checker/state.json | 49 +++ sgtlearn/_export.py | 79 +++-- ...st_sgtlearn_export_py__draw_leaf_text_1.py | 289 ++++++++++++++++++ tests/test_plot_tree.py | 58 ++++ 7 files changed, 539 insertions(+), 28 deletions(-) create mode 100644 .test-checker/reports/r-20261002-215733-b7b8.json create mode 100644 .test-checker/reports/r-20261002-215733-b7b8.md create mode 100644 tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py diff --git a/.test-checker/ledger.jsonl b/.test-checker/ledger.jsonl index f3f5ddd..1f2507f 100644 --- a/.test-checker/ledger.jsonl +++ b/.test-checker/ledger.jsonl @@ -2079,3 +2079,16 @@ {"ts":"2026-10-02T20:30:43.618Z","run":"r-20261002-162723-b8b5","type":"classified","exec":1,"counts":{"sgtlearn-ensemble-random-sgforest-regressor-py--randomsgforestregressor-check-x-y":{"accepted":47}},"queued":0} {"ts":"2026-10-02T20:30:43.673Z","run":"r-20261002-162723-b8b5","type":"run_reported","likely_bugs":0,"needs_decision":0} {"ts":"2026-10-02T20:30:43.740Z","run":"r-20261002-162723-b8b5","type":"promoted","files":["tests/regression/test_sgtlearn_ensemble_random_sgforest_regressor_py__randomsgforestregressor_check_x_y_1.py"],"verified":["sgtlearn/ensemble/random_sgforest_regressor.py"]} +{"ts":"2026-10-03T01:57:33.139Z","run":"r-20261002-215733-b7b8","type":"run_started","mode":"bugfix","revisions":{"buggy":"53a5c148d71076030fca81ed23bf8df64780e3d2","fixed":"676c746e8e2c52431adb7bb362fb3c2428833a4f"},"given":{"buggy":"HEAD","fixed":"WORKTREE"}} +{"ts":"2026-10-03T01:57:36.583Z","run":"r-20261002-215733-b7b8","type":"target_added","target":"sgtlearn-export-py--draw-leaf-text","file":"sgtlearn/_export.py","symbol":"_draw_leaf_text","lines":[685,765],"rev":"fixed","body_sha":"faf7044cbc68687fcfa0ea5b13ea538a606b5779d860a71be10b8c0c9ccb09a1"} +{"ts":"2026-10-03T01:58:04.295Z","run":"r-20261002-215733-b7b8","type":"context_set","target":"sgtlearn-export-py--draw-leaf-text","shingle_ratio":0} +{"ts":"2026-10-03T01:58:04.349Z","run":"r-20261002-215733-b7b8","type":"payload_emitted","payload":"p-e471ca8a","role":"spec","sha256":"6d27fc85bb8d000d2c86b86805d3ec975f42a1fda3dba494d8f92f355edf3aa9","target":"sgtlearn-export-py--draw-leaf-text"} +{"ts":"2026-10-03T01:59:28.987Z","run":"r-20261002-215733-b7b8","type":"spec_saved","target":"sgtlearn-export-py--draw-leaf-text","shingle_ratio":0} +{"ts":"2026-10-03T01:59:29.046Z","run":"r-20261002-215733-b7b8","type":"bundle_frozen","target":"sgtlearn-export-py--draw-leaf-text","bundle":"b-b7b8-sgtlearn-export-py--draw-leaf-text"} +{"ts":"2026-10-03T01:59:34.750Z","run":"r-20261002-215733-b7b8","type":"payload_emitted","payload":"p-c938749f","role":"writer","sha256":"2bb722170802873f243950a6b3c462a56963a47bfb1eb760125e1fc085d7860b","target":"sgtlearn-export-py--draw-leaf-text"} +{"ts":"2026-10-03T01:59:42.323Z","run":"r-20261002-215733-b7b8","type":"payload_emitted","payload":"p-4bd33228","role":"writer","sha256":"325fe67e3a0cfc9e93e785cb01f65540cddf10cd4c2a63ba227634458617b724","target":"sgtlearn-export-py--draw-leaf-text"} +{"ts":"2026-10-03T02:00:21.771Z","run":"r-20261002-215733-b7b8","type":"tests_submitted","target":"sgtlearn-export-py--draw-leaf-text","round":0,"role":"writer","payload":"p-4bd33228","files":["test_sgtlearn_export_py__draw_leaf_text_1.py"]} +{"ts":"2026-10-03T02:02:09.212Z","run":"r-20261002-215733-b7b8","type":"exec_completed","exec":1,"labels":{"buggy":{"failure":9,"pass":30,"load_errors":0},"fixed":{"pass":39,"load_errors":0}}} +{"ts":"2026-10-03T02:02:11.890Z","run":"r-20261002-215733-b7b8","type":"classified","exec":1,"counts":{"sgtlearn-export-py--draw-leaf-text":{"effective":9,"neutral":30}},"queued":0} +{"ts":"2026-10-03T02:02:15.743Z","run":"r-20261002-215733-b7b8","type":"run_reported","likely_bugs":0,"needs_decision":0} +{"ts":"2026-10-03T02:02:15.813Z","run":"r-20261002-215733-b7b8","type":"promoted","files":["tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py"],"verified":["sgtlearn/_export.py"]} diff --git a/.test-checker/reports/r-20261002-215733-b7b8.json b/.test-checker/reports/r-20261002-215733-b7b8.json new file mode 100644 index 0000000..6d22045 --- /dev/null +++ b/.test-checker/reports/r-20261002-215733-b7b8.json @@ -0,0 +1,41 @@ +{ + "run": "r-20261002-215733-b7b8", + "mode": "bugfix", + "revisions": { + "buggy": { + "given": "HEAD", + "commit": "53a5c148d71076030fca81ed23bf8df64780e3d2" + }, + "fixed": { + "given": "WORKTREE", + "commit": "676c746e8e2c52431adb7bb362fb3c2428833a4f" + } + }, + "date": "2026-10-03T02:02:15.741Z", + "harness": "claude-code", + "blind_isolation": "strong", + "summary": [ + { + "target": "sgtlearn-export-py--draw-leaf-text", + "effective": 9, + "misguided": 0, + "neutral": 30, + "broken": 0, + "flaky": 0, + "dropped": 0 + } + ], + "likely_bugs": [], + "suspicions": [], + "needs_decision": [], + "misguided": [], + "dropped": [], + "promotable": [ + { + "file": "tests/tcheck/test_sgtlearn_export_py__draw_leaf_text_1.py", + "dest": "tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py", + "tests": 39 + } + ], + "not_promotable": [] +} diff --git a/.test-checker/reports/r-20261002-215733-b7b8.md b/.test-checker/reports/r-20261002-215733-b7b8.md new file mode 100644 index 0000000..539680f --- /dev/null +++ b/.test-checker/reports/r-20261002-215733-b7b8.md @@ -0,0 +1,38 @@ +# test-checker report r-20261002-215733-b7b8 + +- Mode: bugfix +- buggy: HEAD → `53a5c148d7` +- fixed: WORKTREE → `676c746e8e` +- Date: 2026-10-03T02:02:15.741Z +- Harness: claude-code +- Blind isolation: strong + +## Summary + +| target | effective | misguided | neutral | broken | flaky | dropped | +|---|---|---|---|---|---|---| +| sgtlearn-export-py--draw-leaf-text | 9 | 0 | 30 | 0 | 0 | 0 | + +## Likely bugs + +_none_ + +## Suspicions + +_none_ + +## Needs your decision + +_none_ + +## Misguided tests and what happened to them + +_none_ + +## Dropped + +_none_ + +## Promotable tests + +- tests/tcheck/test_sgtlearn_export_py__draw_leaf_text_1.py → `tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py` (39 test(s)) diff --git a/.test-checker/state.json b/.test-checker/state.json index b39e95f..a31b627 100644 --- a/.test-checker/state.json +++ b/.test-checker/state.json @@ -32,6 +32,11 @@ "sha": "b2ad723f36c9520fda0013c7951e257e4907ba73", "run": "r-20261002-162723-b8b5", "at": "2026-10-02T20:30:43.728Z" + }, + "sgtlearn/_export.py": { + "sha": "9af0b3c1a90f76795010d09eac3d2257bab5af37", + "run": "r-20261002-215733-b7b8", + "at": "2026-10-03T02:02:15.801Z" } }, "protected": { @@ -715,6 +720,50 @@ "sgtlearn-ensemble-random-sgforest-regressor-py--randomsgforestregressor-check-x-y::tests.tcheck.test_sgtlearn_ensemble_random_sgforest_regressor_py__randomsgforestregressor_check_x_y_1::test_float32_overflow_in_sample_weight_is_rejected" ], "run": "r-20261002-162723-b8b5" + }, + "tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py": { + "test_ids": [ + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_single_output_shows_majority_class_name", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_tie_picks_lowest_class_index", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_all_equal_counts_picks_first_class", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_multi_output_shows_bracketed_comma_separated_names", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_multi_output_uses_per_output_class_names", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_empty_class_counts_shows_empty_string", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_classifier_missing_class_counts_shows_empty_string", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_none_value_shows_ellipsis", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_scalar_value_formatted_to_precision", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_scalar_value_respects_different_precision", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_scalar_value_zero_precision", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_multi_output_shows_bracketed_formatted_values", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_multi_output_tuple_sequence", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_regressor_ignores_class_names", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_is_bold", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_positioned_at_xy", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_uses_axes_transform", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_centred_horizontally_and_vertically", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_uses_given_fontsize", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_default_fontsize_when_none", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_uses_given_color", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_primary_line_is_added_to_host_axes", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_none_returns_single_text_artist", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_other_label_value_returns_single_text_artist", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_all_returns_two_text_artists", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_all_primary_artist_first", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_none_does_not_require_n_samples", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_shows_n_samples_without_impurity", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_includes_impurity_line_when_requested", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_for_classifier_node", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_is_not_bold", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_placed_below_primary", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_anchored_top_and_centred", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_uses_axes_transform", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_fontsize_one_smaller_than_integer_fontsize", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_subtitle_default_fontsize_when_fontsize_none", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_all_missing_n_samples_raises_key_error", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_all_with_impurity_missing_impurity_raises_key_error", + "sgtlearn-export-py--draw-leaf-text::tests.tcheck.test_sgtlearn_export_py__draw_leaf_text_1::test_label_all_without_impurity_does_not_need_impurity_key" + ], + "run": "r-20261002-215733-b7b8" } }, "quarantined": [], diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index acf6ccb..9af0b3c 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -82,12 +82,6 @@ def group_label(cols: Sequence[int]) -> str: prefer_logical_name=feature_names is None, ) - def per_output(items: Sequence[str]) -> str: - return items[0] if len(items) == 1 else "[" + ", ".join(items) + "]" - - def fmt_list(values: Sequence[float]) -> str: - return "[" + ", ".join(f"{v:.{decimals}f}" for v in values) + "]" - names: Any = class_names if names is None: names = getattr(estimator, "classes_", None) @@ -97,14 +91,15 @@ def leaf_text(node: dict) -> str: if not isinstance(estimator, SGTClassifier): value = node["value"] if isinstance(value, (list, tuple)): - return f"value: {fmt_list(value)}" + return f"value: {_fmt_list(value, decimals)}" return f"value: {value:.{decimals}f}" counts = node["class_counts"] - text = "class: " + per_output( + text = "class: " + _per_output( [str(names_per_output[o][int(np.argmax(c))]) for o, c in enumerate(counts)] ) if show_weights: - text = f"weights: {per_output([fmt_list(c) for c in counts])} {text}" + fmt = [_fmt_list(c, decimals) for c in counts] + text = f"weights: {_per_output(fmt)} {text}" return text lines: list[str] = [] @@ -125,6 +120,15 @@ def walk(nid: int, depth: int) -> None: return "\n".join(lines) + "\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) + "]" + + +def _fmt_list(values: Sequence[float], decimals: int) -> str: + return "[" + ", ".join(f"{v:.{decimals}f}" for v in values) + "]" + + def _check_fitted_sgt(estimator: Any, caller: str) -> None: if not isinstance(estimator, (SGTClassifier, SGTRegressor)): raise TypeError( @@ -685,7 +689,7 @@ def _draw_leaf_text( node: dict, *, is_classifier: bool, - class_names: list[str] | None, + class_names: list[list[str]] | None, criterion: str, precision: int, fontsize: int | None, @@ -695,19 +699,33 @@ def _draw_leaf_text( ) -> list: """Render a leaf as bold colored text (no box). + ``node`` may also be an internal node cut off by ``max_depth``. Regressor + internals carry no ``value``, so they render as ``"…"``; classifier + internals still carry ``class_counts`` and show their majority class. + Multi-output nodes show one entry per output, as ``[a, b]``. + Returns a list of Text artists (one bold class/value line, optional subtitle lines for ``label='all'``). """ artists: list = [] if is_classifier: - counts_by_output = list(node.get("class_counts", [])) - counts = list(counts_by_output[0]) if counts_by_output else [] - if counts: - arg = max(range(len(counts)), key=lambda i: counts[i]) - primary = class_names[arg] if class_names is not None else str(arg) - else: - primary = "" + assert class_names is not None + counts_by_output = node.get("class_counts") or [] + primary = ( + _per_output( + [ + class_names[o][int(np.argmax(counts))] + for o, counts in enumerate(counts_by_output) + ] + ) + if counts_by_output + else "" + ) + elif node["value"] is None: + primary = "…" + elif isinstance(node["value"], (list, tuple)): + primary = _fmt_list(node["value"], precision) else: primary = f"{node['value']:.{precision}f}" @@ -1384,9 +1402,13 @@ def plot_tree( feature_names : list of str, optional Display names for input columns. class_names : list of str, bool, or None, default=None - Class labels shown on classifier leaves. + Class labels shown on classifier leaves. For multi-output + classifiers, pass one list per output. label : str, default="feature" - Controls field-name labels in node text. + Controls field-name labels in node text. The ``n=`` counts are the + rounded weighted sample counts for classifiers (so they reflect + ``sample_weight`` and ``class_weight``) and unweighted training + sample counts for regressors. impurity : bool, default=False Whether to display node impurity. proportion : bool, default=False @@ -1446,19 +1468,20 @@ def plot_tree( palette = _build_palette(cmap, tree["num_partitions"]) is_classifier = isinstance(estimator, SGTClassifier) - resolved_class_names: list[str] | None + # One list of class labels per output. + resolved_class_names: list[list[str]] | None if not is_classifier: resolved_class_names = None - elif class_names is True: - clf_estimator: SGTClassifier = estimator # type: ignore[assignment] - classes: Any = ( - clf_estimator.classes_ if clf_estimator.classes_ is not None else [] - ) - resolved_class_names = [str(c) for c in classes] elif class_names in (None, False): - resolved_class_names = [str(c) for c in range(tree["num_classes"][0])] + resolved_class_names = [[str(c) for c in range(k)] for k in tree["num_classes"]] else: - resolved_class_names = list(class_names) # type: ignore[arg-type] + names: Any = ( + estimator.classes_ # type: ignore[union-attr] + if class_names is True + else class_names + ) + names_per_output = [names] if tree["num_outputs"] == 1 else names + resolved_class_names = [[str(c) for c in n] for n in names_per_output] feat_names = _resolve_feature_names(estimator, feature_names) diff --git a/tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py b/tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py new file mode 100644 index 0000000..0439380 --- /dev/null +++ b/tests/regression/test_sgtlearn_export_py__draw_leaf_text_1.py @@ -0,0 +1,289 @@ +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +import pytest +from matplotlib.text import Text + +from sgtlearn._export import _draw_leaf_text + + +@pytest.fixture +def ax(): + fig, ax = plt.subplots() + yield ax + plt.close("all") + + +def _is_bold(artist): + weight = artist.get_fontweight() + if isinstance(weight, str): + return weight == "bold" + return weight >= 700 + + +def _draw(ax, node, *, is_classifier=False, class_names=None, criterion="gini", + precision=2, fontsize=10, color="black", label="none", impurity=False, + x=0.5, y=0.5): + return _draw_leaf_text( + ax, + x, + y, + node, + is_classifier=is_classifier, + class_names=class_names, + criterion=criterion, + precision=precision, + fontsize=fontsize, + color=color, + label=label, + impurity=impurity, + ) + + +# ---------------------------------------------------------------- classifier + +def test_classifier_single_output_shows_majority_class_name(ax): + node = {"class_counts": [[1.0, 5.0, 2.0]]} + artists = _draw(ax, node, is_classifier=True, class_names=[["a", "b", "c"]]) + assert artists[0].get_text() == "b" + + +def test_classifier_tie_picks_lowest_class_index(ax): + node = {"class_counts": [[1.0, 4.0, 4.0]]} + artists = _draw(ax, node, is_classifier=True, class_names=[["a", "b", "c"]]) + assert artists[0].get_text() == "b" + + +def test_classifier_all_equal_counts_picks_first_class(ax): + node = {"class_counts": [[3.0, 3.0]]} + artists = _draw(ax, node, is_classifier=True, class_names=[["yes", "no"]]) + assert artists[0].get_text() == "yes" + + +def test_classifier_multi_output_shows_bracketed_comma_separated_names(ax): + node = {"class_counts": [[1.0, 9.0], [7.0, 2.0, 0.0]]} + class_names = [["neg", "pos"], ["x", "y", "z"]] + artists = _draw(ax, node, is_classifier=True, class_names=class_names) + assert artists[0].get_text() == "[pos, x]" + + +def test_classifier_multi_output_uses_per_output_class_names(ax): + node = {"class_counts": [[0.0, 2.0], [0.0, 2.0]]} + class_names = [["a0", "a1"], ["b0", "b1"]] + artists = _draw(ax, node, is_classifier=True, class_names=class_names) + assert artists[0].get_text() == "[a1, b1]" + + +def test_classifier_empty_class_counts_shows_empty_string(ax): + node = {"class_counts": []} + artists = _draw(ax, node, is_classifier=True, class_names=[["a", "b"]]) + assert artists[0].get_text() == "" + + +def test_classifier_missing_class_counts_shows_empty_string(ax): + node = {} + artists = _draw(ax, node, is_classifier=True, class_names=[["a", "b"]]) + assert artists[0].get_text() == "" + + +# ----------------------------------------------------------------- regressor + +def test_regressor_none_value_shows_ellipsis(ax): + artists = _draw(ax, {"value": None}) + assert artists[0].get_text() == "\u2026" + + +def test_regressor_scalar_value_formatted_to_precision(ax): + artists = _draw(ax, {"value": 1.23456}, precision=2) + assert artists[0].get_text() == "1.23" + + +def test_regressor_scalar_value_respects_different_precision(ax): + artists = _draw(ax, {"value": 3.0}, precision=4) + assert artists[0].get_text() == "3.0000" + + +def test_regressor_scalar_value_zero_precision(ax): + artists = _draw(ax, {"value": 7.0}, precision=0) + assert artists[0].get_text() == "7" + + +def test_regressor_multi_output_shows_bracketed_formatted_values(ax): + artists = _draw(ax, {"value": [1.0, 2.5, -0.333]}, precision=2) + assert artists[0].get_text() == "[1.00, 2.50, -0.33]" + + +def test_regressor_multi_output_tuple_sequence(ax): + artists = _draw(ax, {"value": (0.1, 0.2)}, precision=1) + assert artists[0].get_text() == "[0.1, 0.2]" + + +def test_regressor_ignores_class_names(ax): + artists = _draw(ax, {"value": 2.0}, class_names=None, precision=1) + assert artists[0].get_text() == "2.0" + + +# ------------------------------------------------------- primary line styling + +def test_primary_line_is_bold(ax): + artists = _draw(ax, {"value": 1.0}) + assert _is_bold(artists[0]) + + +def test_primary_line_positioned_at_xy(ax): + artists = _draw(ax, {"value": 1.0}, x=0.25, y=0.75) + assert artists[0].get_position() == pytest.approx((0.25, 0.75)) + + +def test_primary_line_uses_axes_transform(ax): + artists = _draw(ax, {"value": 1.0}) + assert artists[0].get_transform() == ax.transAxes + + +def test_primary_line_centred_horizontally_and_vertically(ax): + artists = _draw(ax, {"value": 1.0}) + assert artists[0].get_ha() == "center" + assert artists[0].get_va() == "center" + + +def test_primary_line_uses_given_fontsize(ax): + artists = _draw(ax, {"value": 1.0}, fontsize=14) + assert artists[0].get_fontsize() == pytest.approx(14) + + +def test_primary_line_default_fontsize_when_none(ax): + artists = _draw(ax, {"value": 1.0}, fontsize=None) + expected = matplotlib.font_manager.FontProperties( + size=matplotlib.rcParams["font.size"] + ).get_size_in_points() + assert artists[0].get_fontsize() == pytest.approx(expected) + + +def test_primary_line_uses_given_color(ax): + artists = _draw(ax, {"value": 1.0}, color="red") + assert mcolors.same_color(artists[0].get_color(), "red") + + +def test_primary_line_is_added_to_host_axes(ax): + artists = _draw(ax, {"value": 1.0}) + assert artists[0].axes is ax + + +# --------------------------------------------------------------- return shape + +def test_label_none_returns_single_text_artist(ax): + artists = _draw(ax, {"value": 1.0}, label="none") + assert isinstance(artists, list) + assert len(artists) == 1 + assert isinstance(artists[0], Text) + + +def test_other_label_value_returns_single_text_artist(ax): + artists = _draw(ax, {"value": 1.0}, label="feature") + assert len(artists) == 1 + + +def test_label_all_returns_two_text_artists(ax): + node = {"value": 1.0, "n_samples": 10, "impurity": 0.5} + artists = _draw(ax, node, label="all") + assert len(artists) == 2 + assert all(isinstance(a, Text) for a in artists) + + +def test_label_all_primary_artist_first(ax): + node = {"value": 1.0, "n_samples": 10, "impurity": 0.5} + artists = _draw(ax, node, label="all", precision=2) + assert artists[0].get_text() == "1.00" + assert _is_bold(artists[0]) + + +def test_label_none_does_not_require_n_samples(ax): + artists = _draw(ax, {"value": 1.0}, label="none", impurity=True) + assert len(artists) == 1 + + +# ------------------------------------------------------------------ subtitle + +def test_subtitle_shows_n_samples_without_impurity(ax): + node = {"value": 1.0, "n_samples": 42, "impurity": 0.5} + artists = _draw(ax, node, label="all", impurity=False) + assert artists[1].get_text() == "n = 42" + + +def test_subtitle_includes_impurity_line_when_requested(ax): + node = {"value": 1.0, "n_samples": 42, "impurity": 0.123456} + artists = _draw(ax, node, label="all", impurity=True, + criterion="squared_error", precision=3) + lines = artists[1].get_text().split("\n") + assert lines == ["n = 42", "squared_error = 0.123"] + + +def test_subtitle_for_classifier_node(ax): + node = {"class_counts": [[2.0, 8.0]], "n_samples": 10, "impurity": 0.32} + artists = _draw(ax, node, is_classifier=True, class_names=[["a", "b"]], + label="all", impurity=True, criterion="gini", precision=2) + assert artists[0].get_text() == "b" + assert artists[1].get_text().split("\n") == ["n = 10", "gini = 0.32"] + + +def test_subtitle_is_not_bold(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all") + assert not _is_bold(artists[1]) + + +def test_subtitle_placed_below_primary(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all", x=0.4, y=0.6) + sx, sy = artists[1].get_position() + assert sx == pytest.approx(0.4) + assert sy < 0.6 + + +def test_subtitle_anchored_top_and_centred(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all") + assert artists[1].get_va() == "top" + assert artists[1].get_ha() == "center" + + +def test_subtitle_uses_axes_transform(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all") + assert artists[1].get_transform() == ax.transAxes + + +def test_subtitle_fontsize_one_smaller_than_integer_fontsize(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all", fontsize=12) + assert artists[1].get_fontsize() == pytest.approx(11) + + +def test_subtitle_default_fontsize_when_fontsize_none(ax): + node = {"value": 1.0, "n_samples": 5, "impurity": 0.0} + artists = _draw(ax, node, label="all", fontsize=None) + expected = matplotlib.font_manager.FontProperties( + size=matplotlib.rcParams["font.size"] + ).get_size_in_points() + assert artists[1].get_fontsize() == pytest.approx(expected) + + +# -------------------------------------------------------------------- errors + +def test_label_all_missing_n_samples_raises_key_error(ax): + node = {"value": 1.0, "impurity": 0.5} + with pytest.raises(KeyError): + _draw(ax, node, label="all", impurity=False) + + +def test_label_all_with_impurity_missing_impurity_raises_key_error(ax): + node = {"value": 1.0, "n_samples": 10} + with pytest.raises(KeyError): + _draw(ax, node, label="all", impurity=True) + + +def test_label_all_without_impurity_does_not_need_impurity_key(ax): + node = {"value": 1.0, "n_samples": 10} + artists = _draw(ax, node, label="all", impurity=False) + assert artists[1].get_text() == "n = 10" diff --git a/tests/test_plot_tree.py b/tests/test_plot_tree.py index 13d7499..6e9b11b 100644 --- a/tests/test_plot_tree.py +++ b/tests/test_plot_tree.py @@ -312,3 +312,61 @@ def test_plot_tree_pair_heatmap_renders_without_training_data(): assert "NaN" in [tick.get_text() for tick in panel.get_xticklabels()] assert "NaN" in [tick.get_text() for tick in panel.get_yticklabels()] plt.close(fig) + + +def _leaf_labels(artists) -> list[str]: + return [ + a.get_text() + for a in artists + if hasattr(a, "get_fontweight") and a.get_fontweight() == "bold" + ] + + +def _is_value_label(text: str, n_outputs: int) -> bool: + parts = text[1:-1].split(", ") if n_outputs > 1 else [text] + if n_outputs > 1 and not (text.startswith("[") and text.endswith("]")): + return False + try: + return len([float(p) for p in parts]) == n_outputs + except ValueError: + return False + + +@pytest.mark.parametrize("n_outputs", [1, 2]) +@pytest.mark.parametrize("max_depth", [None, 1]) +def test_plot_tree_regressor_leaf_labels(n_outputs, max_depth): + rng = np.random.default_rng(0) + X = rng.normal(size=(300, 3)) + y = X[:, 0] + np.sin(X[:, 1]) + Y = y if n_outputs == 1 else np.column_stack([y, -y]) + est = SGTRegressor(max_depth=3, random_state=0, tao_n_runs=TEST_TAO_N_RUNS).fit( + X, Y + ) + + labels = _leaf_labels(plot_tree(est, max_depth=max_depth, label="all")) + plt.close("all") + + assert labels + values = [t for t in labels if t != "…"] + assert all(_is_value_label(t, n_outputs) for t in values), labels + if max_depth is None: + assert len(values) == len(labels) + else: + assert "…" in labels + + +def test_plot_tree_multi_output_classifier_labels_every_output(): + rng = np.random.default_rng(0) + X = rng.normal(size=(300, 3)) + Y = np.column_stack([X[:, 0] > 0, np.digitize(X[:, 1], [-1.0, 1.0])]) + est = SGTClassifier(max_depth=2, random_state=0, tao_n_runs=TEST_TAO_N_RUNS).fit( + X, Y + ) + + labels = _leaf_labels(plot_tree(est, class_names=True)) + plt.close("all") + + assert labels + for text in labels: + first, second = text.strip("[]").split(", ") + assert first in {"0", "1"} and second in {"0", "1", "2"}, labels