From 2e139f0300a6336b44d6e9c005ca15f7fcbe78b3 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Wed, 2 Sep 2026 19:12:59 -0400 Subject: [PATCH] test: replace broad matrices with focused contracts --- cpp/tests/test_branch_assignment.cpp | 26 +- tests/discretizer_grid.py | 19 +- ..._categorical_classification_discretizer.py | 78 +-- ...test_categorical_regression_discretizer.py | 123 ++-- tests/test_datasets.py | 44 ++ tests/test_features.py | 71 ++- tests/test_mae_regression_stress.py | 57 +- tests/test_missing_values.py | 287 +-------- tests/test_multioutput.py | 77 +++ tests/test_plot_helpers.py | 594 +++--------------- tests/test_plot_tree.py | 130 ++-- tests/test_random_sgforest_validation.py | 134 ++++ tests/test_tao.py | 471 +++----------- tests/test_tree_export.py | 34 +- ...t_univariate_classification_discretizer.py | 79 +-- .../test_univariate_regression_discretizer.py | 151 +++-- tests/test_weighted_sample.py | 74 ++- 17 files changed, 885 insertions(+), 1564 deletions(-) create mode 100644 tests/test_datasets.py create mode 100644 tests/test_multioutput.py create mode 100644 tests/test_random_sgforest_validation.py diff --git a/cpp/tests/test_branch_assignment.cpp b/cpp/tests/test_branch_assignment.cpp index 4ccdfea..ccb36b5 100644 --- a/cpp/tests/test_branch_assignment.cpp +++ b/cpp/tests/test_branch_assignment.cpp @@ -20,17 +20,23 @@ namespace { constexpr double kEps = 1e-6; -void assert_coordinate_descent_non_worsening(BranchAssignment &obj, - size_t numPartitions, - double tol = kEps) { +void assert_coordinate_descent_finds_split(BranchAssignment &obj, + size_t numPartitions, + size_t samplesPerBin, + double tol = kEps) { std::mt19937_64 rng(42); const double initial = obj.objective(); const double returned = coordinateDescent(numPartitions, obj, rng, /*maxIters=*/30, - /*patience=*/8); + /*patience=*/8, /*hasNanRoutingBin=*/false); const double from_state = obj.objective(); - REQUIRE(from_state <= initial + tol); + REQUIRE(from_state < initial - tol); REQUIRE_THAT(from_state, WithinAbs(returned, tol)); + const std::vector forward = {0, 1}; + const std::vector reverse = {1, 0}; + REQUIRE((obj.assignments == forward || obj.assignments == reverse)); + REQUIRE(obj.partitionSampleCounts() == + std::vector{samplesPerBin, samplesPerBin}); } } // namespace @@ -45,7 +51,7 @@ TEST_CASE("EntropyBranchAssignment coordinate descent", std::vector leafSampleCounts = {10, 10}; EntropyBranchAssignment obj(assignments, kParts, stats, leafWeights, leafSampleCounts, {kClasses}); - assert_coordinate_descent_non_worsening(obj, kParts); + assert_coordinate_descent_finds_split(obj, kParts, 10); } TEST_CASE("GiniBranchAssignment coordinate descent", @@ -58,7 +64,7 @@ TEST_CASE("GiniBranchAssignment coordinate descent", std::vector leafSampleCounts = {10, 10}; GiniBranchAssignment obj(assignments, kParts, stats, leafWeights, leafSampleCounts, {kClasses}); - assert_coordinate_descent_non_worsening(obj, kParts); + assert_coordinate_descent_finds_split(obj, kParts, 10); } TEST_CASE("SquaredErrorBranchAssignment coordinate descent", @@ -71,7 +77,7 @@ TEST_CASE("SquaredErrorBranchAssignment coordinate descent", std::vector leafSampleCounts = {3, 3}; SquaredErrorBranchAssignment obj(assignments, kParts, stats, leafWeights, leafSampleCounts); - assert_coordinate_descent_non_worsening(obj, kParts); + assert_coordinate_descent_finds_split(obj, kParts, 3); } TEST_CASE("GainHessianBranchAssignment coordinate descent", @@ -83,7 +89,7 @@ TEST_CASE("GainHessianBranchAssignment coordinate descent", std::vector leafSampleCounts = {3, 3}; GainHessianBranchAssignment obj(assignments, kParts, stats, leafWeights, leafSampleCounts, 1.0); - assert_coordinate_descent_non_worsening(obj, kParts); + assert_coordinate_descent_finds_split(obj, kParts, 3); } TEST_CASE("AbsoluteErrorBranchAssignment coordinate descent", @@ -102,7 +108,7 @@ TEST_CASE("AbsoluteErrorBranchAssignment coordinate descent", std::vector leafSampleCounts = {3, 3}; AbsoluteErrorBranchAssignment obj(assignments, kParts, leafYs, leafWs, leafWeights, leafSampleCounts); - assert_coordinate_descent_non_worsening(obj, kParts); + assert_coordinate_descent_finds_split(obj, kParts, 3); } TEST_CASE("BranchAssignment tracks partition sample counts", diff --git a/tests/discretizer_grid.py b/tests/discretizer_grid.py index b230900..8726dce 100644 --- a/tests/discretizer_grid.py +++ b/tests/discretizer_grid.py @@ -1,22 +1,9 @@ -"""Shared hyperparameter grids for univariate discretizer sklearn fidelity tests. - -Constants are imported by ``test_univariate_*`` modules to build Cartesian products -of sample size, tree depth, leaf limits, and gain thresholds. -""" +"""Shared single- versus multi-output parameters for fidelity tests.""" import pytest -# sklearn fidelity is a correctness property, not size-dependent: one N is enough. -N_VALUES = [5000] -NUM_CLASSES_VALUES = [2, 3] -# One-hot categorical block width (number of binary columns / categories). -NUM_CATEGORIES_VALUES = [2, 3, 4] -MIN_LEAF_VALUES = [1, 10] -MIN_GAIN_VALUES = [0.0, 1e-7] -MAX_DEPTH_VALUES = [0, 4] -MAX_LEAF_VALUES = [0, 100] -# ``1`` is scalar-y behavior; ``2``/``3`` exercise multioutput. -N_OUTPUTS_VALUES = [1, 2, 3] +# ``1`` is scalar-y behavior; ``2`` exercises the shared multi-output branch. +N_OUTPUTS_VALUES = [1, 2] def n_outputs_params(): diff --git a/tests/test_categorical_classification_discretizer.py b/tests/test_categorical_classification_discretizer.py index c5add38..84ed3b6 100644 --- a/tests/test_categorical_classification_discretizer.py +++ b/tests/test_categorical_classification_discretizer.py @@ -1,28 +1,12 @@ -"""Fidelity tests: ``CategoricalClassificationDiscretizer`` vs ``sklearn.tree.DecisionTreeClassifier``. - -Trains both on a one-hot feature block and compares predictions across the same -inner-tree constraint grid used by the univariate discretizer tests. -""" +"""Focused categorical-classification discretizer contracts and sklearn parity.""" from __future__ import annotations -from itertools import product - import numpy as np import pytest from Discretizers import CategoricalClassificationDiscretizer from sklearn.tree import DecisionTreeClassifier -from tests.discretizer_grid import ( - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - MIN_GAIN_VALUES, - MIN_LEAF_VALUES, - N_VALUES, - NUM_CLASSES_VALUES, - n_outputs_params, -) - def _make_onehot( n_samples: int, @@ -46,32 +30,20 @@ def classification_predict( return np.asarray(bin_preds[bin_locs], dtype=np.uintp) -GRID = list( - product( - N_VALUES, - NUM_CLASSES_VALUES, - MIN_LEAF_VALUES, - MIN_GAIN_VALUES, - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - ) -) -IDS = [ - f"N={n}|C={c}|leaf={leaf}|gain={gain}|depth={depth}|max_leaf={max_leaf}" - for n, c, leaf, gain, depth, max_leaf in GRID +PARITY_CASES = [ + pytest.param("gini", 2, 1, 0.0, 0, 0, 1, id="gini-binary"), + pytest.param("entropy", 3, 1, 0.0, 0, 0, 2, id="entropy-multioutput"), + pytest.param("gini", 4, 1, 0.0, 2, 0, 1, id="depth-limited"), + pytest.param("entropy", 4, 1, 0.0, 0, 2, 2, id="leaf-limited"), ] -@pytest.mark.parametrize("n_outputs", n_outputs_params()) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) @pytest.mark.parametrize( - "n_samples,num_classes,min_leaf_size,min_gain_split,max_depth,max_leaf", - GRID, - ids=IDS, + "criterion,num_classes,min_leaf_size,min_gain_split,max_depth,max_leaf,n_outputs", + PARITY_CASES, ) def test_categorical_onehot_classification_discretizer_vs_sklearn_fidelity( criterion: str, - n_samples: int, num_classes: int, min_leaf_size: int, min_gain_split: float, @@ -81,7 +53,7 @@ def test_categorical_onehot_classification_discretizer_vs_sklearn_fidelity( ) -> None: """Predictions should track ``DecisionTreeClassifier`` on one-hot features.""" rng = np.random.default_rng(12345) - x, y = _make_onehot(n_samples, num_classes, rng, n_outputs=n_outputs) + x, y = _make_onehot(1000, num_classes, rng, n_outputs=n_outputs) clf = DecisionTreeClassifier( criterion=criterion, @@ -132,18 +104,30 @@ def test_categorical_onehot_active_category_maps_to_single_leaf() -> None: def test_categorical_onehot_respects_min_leaf_size() -> None: - rng = np.random.default_rng(2026) n_cat = 3 - x, y = _make_onehot(120, n_cat, rng) + x = np.repeat(np.eye(n_cat, dtype=np.float32), 4, axis=0) + y = np.repeat(np.array([0, 1, 0], dtype=np.uintp), 4) features = np.arange(n_cat, dtype=np.uintp) - min_leaf = 40 - disc = CategoricalClassificationDiscretizer(criterion="entropy") - disc.Train(x, features, y, n_cat, min_leaf, 0.0, 0, 0) - # ``numLeaves`` counts inner-tree bins only; the trailing NaN / catch-all - # routing bin is not subject to ``min_samples_leaf``. - inner_partitions = disc.getInSampleDiscretizations()[: disc.numLeaves] - for part in inner_partitions: - assert len(part) >= min_leaf + allowed = CategoricalClassificationDiscretizer(criterion="entropy") + allowed.Train(x, features, y, 2, 4, 0.0, 0, 0) + blocked = CategoricalClassificationDiscretizer(criterion="entropy") + blocked.Train(x, features, y, 2, 5, 0.0, 0, 0) + + assert allowed.numLeaves == 2 + assert blocked.numLeaves == 1 + + +def test_categorical_classification_gain_threshold_blocks_known_split() -> None: + x = np.repeat(np.eye(3, dtype=np.float32), 4, axis=0) + y = np.repeat(np.array([0, 1, 0], dtype=np.uintp), 4) + features = np.arange(3, dtype=np.uintp) + split = CategoricalClassificationDiscretizer(criterion="gini") + split.Train(x, features, y, 2, 1, 0.0, 0, 0) + blocked = CategoricalClassificationDiscretizer(criterion="gini") + blocked.Train(x, features, y, 2, 1, 1.0, 0, 0) + + assert split.numLeaves == 2 + assert blocked.numLeaves == 0 def test_categorical_onehot_in_sample_partition_covers_rows() -> None: diff --git a/tests/test_categorical_regression_discretizer.py b/tests/test_categorical_regression_discretizer.py index 70ecebb..c259578 100644 --- a/tests/test_categorical_regression_discretizer.py +++ b/tests/test_categorical_regression_discretizer.py @@ -7,23 +7,12 @@ from __future__ import annotations import sys -from itertools import product import numpy as np import pytest from Discretizers import CategoricalRegressionDiscretizer from sklearn.tree import DecisionTreeRegressor -from tests.discretizer_grid import ( - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - MIN_GAIN_VALUES, - MIN_LEAF_VALUES, - N_VALUES, - NUM_CATEGORIES_VALUES, - n_outputs_params, -) - def _make_onehot( n_samples: int, @@ -55,24 +44,6 @@ def sklearn_regression_criterion(user_criterion: str) -> str: return user_criterion -def _sklearn_supports_absolute_error() -> bool: - try: - DecisionTreeRegressor(criterion="absolute_error") - except (ValueError, TypeError): - return False - return True - - -def _skip_if_sklearn_mae_best_first_segfault(max_leaf: int, criterion: str) -> None: - if max_leaf == 0 or criterion not in ("absolute_error", "mae"): - return - if sys.version_info >= (3, 14): - pytest.skip( - "sklearn reference: BestFirstTreeBuilder + absolute_error + max_leaf_nodes " - "segfaults on Python 3.14+; compare MAE with max_leaf=0 cases only" - ) - - def _pred_discrepancy(sk: np.ndarray, ud: np.ndarray, *, use_mae: bool) -> float: """RMSE or mean absolute error between sklearn and native prediction vectors.""" d = sk.astype(np.float64) - ud.astype(np.float64) @@ -81,32 +52,20 @@ def _pred_discrepancy(sk: np.ndarray, ud: np.ndarray, *, use_mae: bool) -> float return float(np.sqrt(np.mean(d**2))) -GRID = list( - product( - N_VALUES, - NUM_CATEGORIES_VALUES, - MIN_LEAF_VALUES, - MIN_GAIN_VALUES, - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - ) -) -IDS = [ - f"N={n}|cat={c}|leaf={leaf}|gain={gain}|depth={depth}|max_leaf={max_leaf}" - for n, c, leaf, gain, depth, max_leaf in GRID +PARITY_CASES = [ + pytest.param("squared_error", 2, 1, 0.0, 0, 0, 1, id="mse-binary"), + pytest.param("absolute_error", 3, 1, 0.0, 0, 0, 2, id="mae-multioutput"), + pytest.param("squared_error", 4, 1, 0.0, 2, 0, 2, id="depth-limited"), + pytest.param("squared_error", 4, 1, 0.0, 0, 2, 1, id="leaf-limited"), ] -@pytest.mark.parametrize("n_outputs", n_outputs_params()) -@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) @pytest.mark.parametrize( - "n_samples,n_categories,min_leaf_size,min_gain_split,max_depth,max_leaf", - GRID, - ids=IDS, + "criterion,n_categories,min_leaf_size,min_gain_split,max_depth,max_leaf,n_outputs", + PARITY_CASES, ) def test_categorical_onehot_regression_discretizer_vs_sklearn_fidelity( criterion: str, - n_samples: int, n_categories: int, min_leaf_size: int, min_gain_split: float, @@ -115,13 +74,8 @@ def test_categorical_onehot_regression_discretizer_vs_sklearn_fidelity( n_outputs: int, ) -> None: """Bin predictions should track ``DecisionTreeRegressor`` within tolerance.""" - if criterion in ("absolute_error", "mae") and not _sklearn_supports_absolute_error(): - pytest.skip("sklearn DecisionTreeRegressor does not support criterion='absolute_error'") - - _skip_if_sklearn_mae_best_first_segfault(max_leaf, criterion) - rng = np.random.default_rng(12345) - x, y = _make_onehot(n_samples, n_categories, rng, n_outputs=n_outputs) + x, y = _make_onehot(1000, n_categories, rng, n_outputs=n_outputs) sk_crit = sklearn_regression_criterion(criterion) reg = DecisionTreeRegressor( @@ -159,7 +113,66 @@ def test_categorical_onehot_regression_discretizer_vs_sklearn_fidelity( assert discrepancy < 0.9 * sigma -@pytest.mark.parametrize("alias,canonical", [("mse", "squared_error"), ("mae", "absolute_error")]) +@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) +def test_categorical_regression_leaf_limit_binds_without_reference( + criterion: str, +) -> None: + rng = np.random.default_rng(8) + x, y = _make_onehot(300, 5, rng) + disc = CategoricalRegressionDiscretizer(criterion=criterion) + disc.Train(x, np.arange(5, dtype=np.uintp), y, 1, 0.0, 0, 3) + bins = disc.transform(x) + assert disc.numLeaves == 3 + assert np.all(bins < disc.numLeaves) + assert np.all(np.isfinite(regression_predict(disc, x))) + + +@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) +def test_categorical_regression_respects_minimum_leaf_size(criterion: str) -> None: + x = np.repeat(np.eye(3, dtype=np.float32), 4, axis=0) + y = np.repeat(np.array([0.0, 10.0, 20.0], dtype=np.float32), 4) + features = np.arange(3, dtype=np.uintp) + allowed = CategoricalRegressionDiscretizer(criterion=criterion) + allowed.Train(x, features, y, 4, 0.0, 0, 0) + blocked = CategoricalRegressionDiscretizer(criterion=criterion) + blocked.Train(x, features, y, 5, 0.0, 0, 0) + + assert allowed.numLeaves == 3 + assert blocked.numLeaves == 1 + + +def test_categorical_regression_gain_threshold_blocks_known_split() -> None: + x = np.repeat(np.eye(3, dtype=np.float32), 4, axis=0) + y = np.repeat(np.array([0.0, 10.0, 20.0], dtype=np.float32), 4) + features = np.arange(3, dtype=np.uintp) + split = CategoricalRegressionDiscretizer(criterion="squared_error") + split.Train(x, features, y, 1, 0.0, 0, 0) + blocked = CategoricalRegressionDiscretizer(criterion="squared_error") + blocked.Train(x, features, y, 1, 1_000.0, 0, 0) + + assert split.numLeaves == 3 + assert blocked.numLeaves == 0 + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), + reason="sklearn absolute_error best-first builder can crash on Python 3.14+", +) +def test_categorical_mae_leaf_limit_matches_sklearn_when_reference_is_safe() -> None: + rng = np.random.default_rng(13) + x, y = _make_onehot(300, 5, rng) + sk = DecisionTreeRegressor(criterion="absolute_error", max_leaf_nodes=3).fit(x, y) + disc = CategoricalRegressionDiscretizer(criterion="absolute_error") + disc.Train(x, np.arange(5, dtype=np.uintp), y, 1, 0.0, 0, 3) + assert sk.get_n_leaves() == disc.numLeaves + assert _pred_discrepancy( + sk.predict(x), regression_predict(disc, x), use_mae=True + ) < np.std(y) + + +@pytest.mark.parametrize( + "alias,canonical", [("mse", "squared_error"), ("mae", "absolute_error")] +) def test_categorical_onehot_regression_criterion_aliases_equivalent( alias: str, canonical: str ) -> None: diff --git a/tests/test_datasets.py b/tests/test_datasets.py new file mode 100644 index 0000000..4dd1bd8 --- /dev/null +++ b/tests/test_datasets.py @@ -0,0 +1,44 @@ +"""Public contracts for bundled synthetic datasets.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from sgtlearn.datasets import make_plus + + +def test_make_plus_matches_documented_geometry() -> None: + n_samples = 200 + grid = 5 + margin = 0.1 + X, y = make_plus(n_samples=n_samples, grid=grid, margin=margin, random_state=7) + + assert X.shape == (n_samples, 2) + assert y.shape == (n_samples,) + assert set(np.unique(y)) <= {0, 1} + assert np.all((0 <= X) & (X < grid)) + fractions = X - np.floor(X) + assert np.all((fractions > margin) & (fractions < 1 - margin)) + cells = np.floor(X).astype(int) + expected = ((cells[:, 0] == grid // 2) | (cells[:, 1] == grid // 2)).astype(int) + np.testing.assert_array_equal(y, expected) + + +def test_make_plus_is_deterministic_for_seed() -> None: + first = make_plus(n_samples=25, random_state=42) + second = make_plus(n_samples=25, random_state=42) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], second[1]) + + +@pytest.mark.parametrize("margin", [-0.01, 0.5, 1.0]) +def test_make_plus_rejects_invalid_margin(margin: float) -> None: + with pytest.raises(ValueError, match="margin"): + make_plus(margin=margin) + + +def test_make_plus_allows_empty_dataset() -> None: + X, y = make_plus(n_samples=0, random_state=0) + assert X.shape == (0, 2) + assert y.shape == (0,) diff --git a/tests/test_features.py b/tests/test_features.py index a54e905..4cc2263 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -7,9 +7,9 @@ import pytest from sklearn.datasets import make_classification -from sgtlearn import SGTClassifier, configure_feature_dict +from sgtlearn import SGTClassifier, configure_feature_dict, plot_tree from sgtlearn._export import _is_categorical_node -from sgtlearn._features import ProcessedFeatures +from sgtlearn._features import FeatureDict, ProcessedFeatures def test_configure_feature_dict_defaults_to_continuous_columns() -> None: @@ -18,6 +18,7 @@ def test_configure_feature_dict_defaults_to_continuous_columns() -> None: assert all(f["type"] == "continuous" for f in pf.features) assert [f["indices"] for f in pf.features] == [[0], [1], [2], [3]] assert pf.logical_names == ("0", "1", "2", "3") + assert pf.to_native() == pf.features def test_configure_feature_dict_fills_unmentioned_columns() -> None: @@ -28,17 +29,12 @@ def test_configure_feature_dict_fills_unmentioned_columns() -> None: 4: [4], }, ) - types = {tuple(f["indices"]): f["type"] for f in pf.features} - assert types[(0, 1, 2)] == "categorical" - assert types[(4,)] == "continuous" - assert types[(3,)] == "continuous" - - -def test_configure_feature_dict_legacy_layout() -> None: - pf = configure_feature_dict(4, feature_dict={0: [0, 1, 2]}) - by_indices = {tuple(f["indices"]): f["type"] for f in pf.features} - assert by_indices[(0, 1, 2)] == "categorical" - assert by_indices[(3,)] == "continuous" + assert pf.features == [ + {"type": "categorical", "indices": [0, 1, 2]}, + {"type": "continuous", "indices": [3]}, + {"type": "continuous", "indices": [4]}, + ] + assert pf.logical_names == ("0", "3", "4") def test_configure_feature_dict_rejects_duplicate_indices() -> None: @@ -46,6 +42,31 @@ def test_configure_feature_dict_rejects_duplicate_indices() -> None: configure_feature_dict(4, feature_dict={0: [0], 1: [0]}) +@pytest.mark.parametrize( + ("feature_dict", "column_names", "match"), + [ + ({"bad": [-1]}, None, "out of range"), + ({"bad": [3]}, None, "out of range"), + ({"bad": ["a"]}, None, "column_names"), + ({"bad": ["missing"]}, ["a", "b", "c"], "not found"), + ], +) +def test_configure_feature_dict_rejects_invalid_columns( + feature_dict: FeatureDict, + column_names: list[str] | None, + match: str, +) -> None: + with pytest.raises(ValueError, match=match): + configure_feature_dict(3, feature_dict, column_names=column_names) + + +def test_configure_feature_dict_classifies_groups_at_two_columns() -> None: + pf = configure_feature_dict(3, {"one": [0], "two": [1, 2]}) + by_name = dict(zip(pf.logical_names, pf.features)) + assert by_name["one"]["type"] == "continuous" + assert by_name["two"]["type"] == "categorical" + + def test_configure_feature_dict_string_keys_and_column_names() -> None: pf = configure_feature_dict( 4, @@ -80,15 +101,11 @@ def test_sgt_classifier_categorical_feature_group_routes_onehot() -> None: tree = clf.tree_export() internal = [n for n in tree["nodes"] if not n["is_leaf"]] assert internal, "expected at least one split" - has_categorical_split = any( - _is_categorical_node(n) for n in internal - ) + has_categorical_split = any(_is_categorical_node(n) for n in internal) assert has_categorical_split def test_plot_tree_categorical_node_labels_merged_categories() -> None: - from sgtlearn._export import _bin_labels_for_categorical_node - rng = np.random.default_rng(7) n_samples = 500 n_cat = 6 @@ -109,11 +126,13 @@ def test_plot_tree_categorical_node_labels_merged_categories() -> None: ) clf.fit(X, y, feature_dict={"species": list(columns)}) - labels: list[str] = [] - for node in clf.tree_export()["nodes"]: - if node["is_leaf"] or not _is_categorical_node(node): - continue - labels.extend(_bin_labels_for_categorical_node(node, columns)) + artists = plot_tree(clf) + labels = [ + tick.get_text() + for artist in artists + if hasattr(artist, "get_xticklabels") + for tick in artist.get_xticklabels() + ] merged = [t for t in labels if t.startswith("[") and "," in t] assert merged, "expected merged-category bucket labels when leaf budget is tight" assert any(label == "cat_5" for label in labels) @@ -200,6 +219,10 @@ def test_ensemble_dataframe_string_valued_feature_dict_routes_onehot() -> None: by_indices = {tuple(f["indices"]): f["type"] for f in pf.features} assert by_indices[tuple(range(n_cat))] == "categorical" assert any( - any(_is_categorical_node(n) for n in t.tree_export()["nodes"] if not n["is_leaf"]) + any( + _is_categorical_node(n) + for n in t.tree_export()["nodes"] + if not n["is_leaf"] + ) for t in forest.estimators_ ), "expected a categorical split in at least one tree" diff --git a/tests/test_mae_regression_stress.py b/tests/test_mae_regression_stress.py index 6cadb41..71433d8 100644 --- a/tests/test_mae_regression_stress.py +++ b/tests/test_mae_regression_stress.py @@ -9,8 +9,6 @@ import numpy as np import pytest -from itertools import product - from Discretizers import UnivariateRegressionDiscretizer @@ -23,7 +21,7 @@ def _run_mae_train_predict( max_depth: int, max_leaf: int, criterion: str, -) -> None: +) -> int: """Train MAE discretizer, run ``transform``, and touch bin predictions (crash smoke test).""" ud = UnivariateRegressionDiscretizer(criterion=criterion) features = np.array([0], dtype=np.uintp) @@ -43,35 +41,18 @@ def _run_mae_train_predict( assert int(bins.size) == n assert len(preds) == ud.numLeaves _ = np.asarray(preds)[bins] + return ud.numLeaves -# More aggressive than discretizer_grid: smallest leaf, zero gain floor, deep/wide trees. -STRESS_GRID = list( - product( - [512, 4096, 12000], # n_samples - [1], # min_leaf — maximum splitting - [0.0], # allow any split gain - [0, 12], # 0 = unlimited depth - [0, 256, 2000], # 0 = unlimited leaves in our builder; large caps stress queues - ) -) -STRESS_IDS = [ - f"n={n}|leaf={leaf}|gain={gain}|depth={depth}|max_leaf={ml}" - for n, leaf, gain, depth, ml in STRESS_GRID -] - - -@pytest.mark.parametrize("criterion", ["mae", "absolute_error"]) @pytest.mark.parametrize( - "n_samples,min_leaf_size,min_gain_split,max_depth,max_leaf", - STRESS_GRID, - ids=STRESS_IDS, + ("n_samples", "max_depth", "max_leaf"), + [ + pytest.param(12000, 0, 0, id="unlimited-growth"), + pytest.param(4096, 0, 256, id="binding-leaf-cap"), + ], ) def test_mae_regression_stress_random_data( - criterion: str, n_samples: int, - min_leaf_size: int, - min_gain_split: float, max_depth: int, max_leaf: int, ) -> None: @@ -79,18 +60,19 @@ def test_mae_regression_stress_random_data( rng = np.random.default_rng(2026) x = rng.random((n_samples, 1), dtype=np.float32) y = rng.standard_normal(n_samples, dtype=np.float64).astype(np.float32) - _run_mae_train_predict( + num_leaves = _run_mae_train_predict( x, y, - min_leaf_size=min_leaf_size, - min_gain_split=min_gain_split, + min_leaf_size=1, + min_gain_split=0.0, max_depth=max_depth, max_leaf=max_leaf, - criterion=criterion, + criterion="absolute_error", ) + if max_leaf: + assert num_leaves == max_leaf -@pytest.mark.parametrize("criterion", ["mae", "absolute_error"]) @pytest.mark.parametrize( "name,x,y", [ @@ -99,11 +81,6 @@ def test_mae_regression_stress_random_data( np.linspace(0.0, 1.0, 800, dtype=np.float32).reshape(-1, 1), np.full(800, 3.14159, dtype=np.float32), ), - ( - "two_unique_y", - np.linspace(0.0, 1.0, 600, dtype=np.float32).reshape(-1, 1), - np.repeat(np.array([0.0, 1.0], dtype=np.float32), 300), - ), ( "duplicate_x_runs", np.concatenate( @@ -115,15 +92,9 @@ def test_mae_regression_stress_random_data( ), np.arange(800, dtype=np.float32) * 0.01, ), - ( - "sorted_strictly_increasing_x", - np.linspace(0.0, 1.0, 2000, dtype=np.float32).reshape(-1, 1), - np.sin(np.linspace(0.0, 6.28, 2000)).astype(np.float32), - ), ], ) def test_mae_regression_stress_structured_data( - criterion: str, name: str, x: np.ndarray, y: np.ndarray, @@ -137,5 +108,5 @@ def test_mae_regression_stress_structured_data( min_gain_split=0.0, max_depth=0, max_leaf=500, - criterion=criterion, + criterion="absolute_error", ) diff --git a/tests/test_missing_values.py b/tests/test_missing_values.py index 931ad6c..4844b2b 100644 --- a/tests/test_missing_values.py +++ b/tests/test_missing_values.py @@ -6,14 +6,11 @@ With ``inner_max_depth=1`` each shape function is a single binary threshold split (standard CART node). References are ``sklearn.tree.DecisionTreeClassifier`` / -``DecisionTreeRegressor`` (``scikit-learn>=1.9`` for ``absolute_error`` NaN -support) with matching criterion and aligned hyperparameters. +``DecisionTreeRegressor`` with matching criteria and aligned hyperparameters. """ from __future__ import annotations -import zlib - import numpy as np import pytest from sklearn.datasets import load_breast_cancer, load_diabetes @@ -41,53 +38,6 @@ def _inject_nan( return out -# (n_samples, n_features) stress shapes: thousands of rows, up to 10 features. -_LARGE_SCALE_SHAPES: list[tuple[int, int]] = [ - (1000, 3), - (1000, 10), - (2000, 5), - (3000, 10), -] - -# Multivariate shapes where ``inner_max_depth=1`` NaN routing matches sklearn in-sample. -_SKLEARN_PARITY_LARGE_SHAPES: list[tuple[int, int]] = [ - (1000, 3), - (1000, 10), - (2000, 5), - (3000, 10), -] - -_REGRESSION_SKLEARN_PARITY_CRITERIA: list[str] = ["squared_error", "absolute_error"] - - -def _large_scale_rng(n_samples: int, n_features: int, salt: str) -> np.random.Generator: - # crc32, not hash(): hash() randomizes strings per process (PYTHONHASHSEED), - # so the generated data would differ run-to-run and across Python versions. - seed = zlib.crc32(f"{n_samples}-{n_features}-{salt}".encode()) - return np.random.default_rng(seed) - - -def _make_large_classification_xy( - n_samples: int, - n_features: int, - rng: np.random.Generator, -) -> tuple[np.ndarray, np.ndarray]: - X = rng.standard_normal((n_samples, n_features)).astype(np.float32) - y = rng.integers(0, 2, size=n_samples) - return X, y - - -def _make_large_regression_xy( - n_samples: int, - n_features: int, - rng: np.random.Generator, -) -> tuple[np.ndarray, np.ndarray]: - X = rng.standard_normal((n_samples, n_features)).astype(np.float32) - coef = rng.standard_normal(n_features) - y = X @ coef + 0.25 * rng.standard_normal(n_samples) - return X, y.astype(np.float64) - - def _sklearn_classification_criterion(criterion: str) -> str: return "entropy" if criterion == "log_loss" else criterion @@ -103,6 +53,7 @@ def _make_inner_depth_one_pair( common = { "min_samples_leaf": min_samples_leaf, "min_impurity_decrease": min_impurity_decrease, + "random_state": 0, } if task == "classification": sk_criterion = _sklearn_classification_criterion(criterion) @@ -131,55 +82,22 @@ def rng() -> np.random.Generator: return np.random.default_rng(42) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) def test_sgt_classifier_inner_depth_one_matches_sklearn_breast_cancer_with_nan( - criterion: str, rng: np.random.Generator, ) -> None: """Breast cancer + scattered NaN: in-sample labels match sklearn CART.""" X, y = load_breast_cancer(return_X_y=True) X = _inject_nan(np.asarray(X, dtype=np.float32), rng, n_cells=18) - sgt, dt = _make_inner_depth_one_pair("classification", criterion) + sgt, dt = _make_inner_depth_one_pair("classification", "gini") sgt.fit(X, y) dt.fit(X, y) np.testing.assert_array_equal(sgt.classes_, dt.classes_) np.testing.assert_array_equal(sgt.predict(X), dt.predict(X)) - - -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -def test_sgt_classifier_predict_proba_matches_sklearn_with_nan( - criterion: str, - rng: np.random.Generator, -) -> None: - X, y = load_breast_cancer(return_X_y=True) - X = _inject_nan(np.asarray(X, dtype=np.float32), rng, n_cells=12) - - sgt, dt = _make_inner_depth_one_pair("classification", criterion) - sgt.fit(X, y) - dt.fit(X, y) - np.testing.assert_allclose(sgt.predict_proba(X), dt.predict_proba(X)) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -def test_sgt_classifier_inner_depth_one_matches_sklearn_synthetic_with_nan( - criterion: str, - rng: np.random.Generator, -) -> None: - n_samples, n_features = 80, 5 - X = rng.standard_normal((n_samples, n_features)).astype(np.float32) - X = _inject_nan(X, rng, n_cells=10) - y = rng.integers(0, 2, size=n_samples) - - sgt, dt = _make_inner_depth_one_pair("classification", criterion) - sgt.fit(X, y) - dt.fit(X, y) - - np.testing.assert_array_equal(sgt.predict(X), dt.predict(X)) - - def test_sgt_classifier_inner_depth_one_handcrafted_multivariate_with_nan() -> None: """Small 2-feature example: NaN routes to the lower-impurity partition at the root.""" X = np.array( @@ -204,33 +122,37 @@ def test_sgt_classifier_inner_depth_one_handcrafted_multivariate_with_nan() -> N np.testing.assert_array_equal(sgt.predict(X), dt.predict(X)) -@pytest.mark.parametrize("criterion", _REGRESSION_SKLEARN_PARITY_CRITERIA) -def test_sgt_regressor_inner_depth_one_matches_sklearn_diabetes_with_nan( - criterion: str, - rng: np.random.Generator, -) -> None: - bunch = load_diabetes() - X = _inject_nan(np.asarray(bunch.data, dtype=np.float32), rng, n_cells=14) - y = np.asarray(bunch.target, dtype=np.float64) +def test_sgt_regressor_handcrafted_nan_route_matches_sklearn() -> None: + X = np.array( + [ + [1.0, 0.0], + [2.0, 1.0], + [np.nan, 0.0], + [4.0, 1.0], + [5.0, 0.0], + [6.0, 1.0], + ], + dtype=np.float32, + ) + y = np.array([0.0, 0.0, 10.0, 10.0, 0.0, 10.0]) - sgt, dt = _make_inner_depth_one_pair("regression", criterion) + sgt, dt = _make_inner_depth_one_pair("regression", "squared_error") sgt.fit(X, y) dt.fit(X, y) + root = sgt.tree_export()["nodes"][0] + assert root["nan_prediction_partition"] == 1 np.testing.assert_allclose(sgt.predict(X), dt.predict(X)) -@pytest.mark.parametrize("criterion", _REGRESSION_SKLEARN_PARITY_CRITERIA) -def test_sgt_regressor_inner_depth_one_matches_sklearn_synthetic_with_nan( - criterion: str, +def test_sgt_regressor_inner_depth_one_matches_sklearn_diabetes_with_nan( rng: np.random.Generator, ) -> None: - n_samples, n_features = 80, 4 - X = rng.standard_normal((n_samples, n_features)).astype(np.float32) - X = _inject_nan(X, rng, n_cells=8) - y = rng.standard_normal(n_samples) + bunch = load_diabetes() + X = _inject_nan(np.asarray(bunch.data, dtype=np.float32), rng, n_cells=14) + y = np.asarray(bunch.target, dtype=np.float64) - sgt, dt = _make_inner_depth_one_pair("regression", criterion) + sgt, dt = _make_inner_depth_one_pair("regression", "squared_error") sgt.fit(X, y) dt.fit(X, y) @@ -282,165 +204,6 @@ def test_sgt_classifier_still_rejects_inf_in_x() -> None: X, y = load_breast_cancer(return_X_y=True) X = np.asarray(X, dtype=np.float32) X[0, 0] = np.inf - clf = SGTClassifier( - inner_max_depth=1, random_state=42, tao_n_runs=TEST_TAO_N_RUNS - ) + clf = SGTClassifier(inner_max_depth=1, random_state=42, tao_n_runs=TEST_TAO_N_RUNS) with pytest.raises(ValueError, match="infinity"): clf.fit(X, y) - - -@pytest.mark.parametrize("n_samples,n_features", _SKLEARN_PARITY_LARGE_SHAPES) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -def test_sgt_classifier_large_scale_nan_matches_sklearn( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - """Thousands of rows / up to 10 features: in-sample labels match sklearn CART.""" - rng = _large_scale_rng(n_samples, n_features, criterion) - X, y = _make_large_classification_xy(n_samples, n_features, rng) - X = _inject_nan(X, rng, frac=0.02) - - sgt, dt = _make_inner_depth_one_pair("classification", criterion) - sgt.fit(X, y) - dt.fit(X, y) - - np.testing.assert_array_equal(sgt.classes_, dt.classes_) - np.testing.assert_array_equal(sgt.predict(X), dt.predict(X)) - - -@pytest.mark.parametrize("n_samples,n_features", _LARGE_SCALE_SHAPES) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -def test_sgt_classifier_large_scale_nan_predict_proba_finite( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - """Large-scale proba outputs are well-formed (routing may differ from sklearn).""" - rng = _large_scale_rng(n_samples, n_features, f"proba_{criterion}") - X, y = _make_large_classification_xy(n_samples, n_features, rng) - X = _inject_nan(X, rng, frac=0.02) - - sgt, _ = _make_inner_depth_one_pair("classification", criterion) - sgt.fit(X, y) - proba = sgt.predict_proba(X) - - assert proba.shape == (n_samples, 2) - assert np.all(np.isfinite(proba)) - np.testing.assert_allclose(proba.sum(axis=1), 1.0, rtol=0.0, atol=1e-6) - assert np.all(proba >= 0.0) - - -@pytest.mark.parametrize("n_samples,n_features", _LARGE_SCALE_SHAPES) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -def test_sgt_classifier_large_scale_nan_smoke( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - """All large shapes: fit/predict with training NaN completes with valid outputs.""" - rng = _large_scale_rng(n_samples, n_features, f"smoke_{criterion}") - X, y = _make_large_classification_xy(n_samples, n_features, rng) - X = _inject_nan(X, rng, frac=0.02) - - sgt, _ = _make_inner_depth_one_pair("classification", criterion) - sgt.fit(X, y) - pred = sgt.predict(X) - proba = sgt.predict_proba(X) - - assert pred.shape == (n_samples,) - assert proba.shape == (n_samples, 2) - assert np.all(np.isfinite(pred)) - assert np.all(np.isfinite(proba)) - - -@pytest.mark.parametrize("n_samples,n_features", _LARGE_SCALE_SHAPES) -def test_sgt_classifier_large_scale_predict_with_new_nan_smoke( - n_samples: int, - n_features: int, -) -> None: - """Large clean fit; NaN at predict time returns finite labels.""" - rng = _large_scale_rng(n_samples, n_features, "predict_new_nan") - X, y = _make_large_classification_xy(n_samples, n_features, rng) - - sgt, _ = _make_inner_depth_one_pair("classification", "gini") - sgt.fit(X, y) - - X_pred = X.copy() - n_pred_nan = max(20, n_samples // 50) - for _ in range(n_pred_nan): - X_pred[rng.integers(0, n_samples), rng.integers(0, n_features)] = np.nan - - pred = sgt.predict(X_pred) - assert pred.shape == (n_samples,) - assert np.all(np.isfinite(pred)) - - -@pytest.mark.parametrize("n_samples,n_features", _SKLEARN_PARITY_LARGE_SHAPES) -@pytest.mark.parametrize("criterion", _REGRESSION_SKLEARN_PARITY_CRITERIA) -def test_sgt_regressor_large_scale_nan_matches_sklearn( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - """Large synthetic regression with scattered NaN matches sklearn in-sample.""" - rng = _large_scale_rng(n_samples, n_features, criterion) - X, y = _make_large_regression_xy(n_samples, n_features, rng) - X = _inject_nan(X, rng, frac=0.02) - - sgt, dt = _make_inner_depth_one_pair("regression", criterion) - sgt.fit(X, y) - dt.fit(X, y) - - sgt_pred, dt_pred = sgt.predict(X), dt.predict(X) - if criterion == "absolute_error": - # MAE splits tie constantly (any split with equal total absolute deviation - # is equally optimal); sklearn and SGT pick different equally-optimal splits, - # so a handful of samples land in different leaves. Allow <=1% to differ. - mismatch = np.mean(~np.isclose(sgt_pred, dt_pred, rtol=1e-5, atol=1e-5)) - assert mismatch <= 0.01, f"MAE prediction mismatch {mismatch:.4%} exceeds 1%" - else: - np.testing.assert_allclose(sgt_pred, dt_pred, rtol=1e-5, atol=1e-5) - - -@pytest.mark.parametrize("n_samples,n_features", _LARGE_SCALE_SHAPES) -@pytest.mark.parametrize("criterion", _REGRESSION_SKLEARN_PARITY_CRITERIA) -def test_sgt_regressor_large_scale_nan_smoke( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - """All large shapes: regression fit/predict with training NaN completes.""" - rng = _large_scale_rng(n_samples, n_features, f"{criterion}_smoke") - X, y = _make_large_regression_xy(n_samples, n_features, rng) - X = _inject_nan(X, rng, frac=0.02) - - sgt, _ = _make_inner_depth_one_pair("regression", criterion) - sgt.fit(X, y) - pred = sgt.predict(X) - - assert pred.shape == (n_samples,) - assert np.all(np.isfinite(pred)) - - -@pytest.mark.parametrize("n_samples,n_features", _LARGE_SCALE_SHAPES) -@pytest.mark.parametrize("criterion", _REGRESSION_SKLEARN_PARITY_CRITERIA) -def test_sgt_regressor_large_scale_predict_with_new_nan_smoke( - n_samples: int, - n_features: int, - criterion: str, -) -> None: - rng = _large_scale_rng(n_samples, n_features, f"{criterion}_predict_new_nan") - X, y = _make_large_regression_xy(n_samples, n_features, rng) - - sgt, _ = _make_inner_depth_one_pair("regression", criterion) - sgt.fit(X, y) - - X_pred = X.copy() - n_pred_nan = max(20, n_samples // 50) - for _ in range(n_pred_nan): - X_pred[rng.integers(0, n_samples), rng.integers(0, n_features)] = np.nan - - pred = sgt.predict(X_pred) - assert pred.shape == (n_samples,) - assert np.all(np.isfinite(pred)) diff --git a/tests/test_multioutput.py b/tests/test_multioutput.py new file mode 100644 index 0000000..3ff4872 --- /dev/null +++ b/tests/test_multioutput.py @@ -0,0 +1,77 @@ +"""Contracts for normalizing single- and multi-output targets.""" + +from __future__ import annotations + +import numpy as np +import pytest +from sklearn.preprocessing import LabelEncoder + +from sgtlearn._multioutput import ( + as_output_matrix, + encode_classification_targets, + label_encoders_as_list, + native_y_array, + squeeze_outputs, + unwrap_classifier_public_attrs, +) + + +@pytest.mark.parametrize( + ("shape", "expected_shape", "n_outputs"), + [((3,), (3, 1), 1), ((3, 1), (3, 1), 1), ((3, 2), (3, 2), 2)], +) +def test_as_output_matrix_normalizes_supported_shapes( + shape: tuple[int, ...], expected_shape: tuple[int, ...], n_outputs: int +) -> None: + actual, actual_outputs = as_output_matrix(np.arange(np.prod(shape)).reshape(shape)) + assert actual.shape == expected_shape + assert actual_outputs == n_outputs + + +@pytest.mark.parametrize("shape", [(3, 0), (2, 2, 1)]) +def test_as_output_matrix_rejects_invalid_shapes(shape: tuple[int, ...]) -> None: + with pytest.raises(ValueError): + as_output_matrix(np.empty(shape)) + + +def test_encode_classification_targets_round_trips_each_output() -> None: + y = np.array([["cat", "red"], ["dog", "blue"], ["cat", "blue"]]) + encoded, encoders, classes, counts = encode_classification_targets(y) + assert encoded.shape == y.shape + assert counts == [2, 2] + for output, encoder in enumerate(encoders): + np.testing.assert_array_equal( + encoder.inverse_transform(encoded[:, output]), y[:, output] + ) + np.testing.assert_array_equal(classes[output], encoder.classes_) + + +def test_label_encoders_must_match_output_count() -> None: + encoder = LabelEncoder().fit(["a", "b"]) + assert label_encoders_as_list(encoder, 1) == [encoder] + with pytest.raises(ValueError, match="2"): + label_encoders_as_list(encoder, 2) + with pytest.raises(ValueError, match="2"): + label_encoders_as_list([encoder, encoder, encoder], 2) + + +def test_classifier_public_attributes_unwrap_only_one_output() -> None: + encoder = LabelEncoder().fit(["a", "b"]) + single = unwrap_classifier_public_attrs([encoder], [encoder.classes_], [2], 1) + assert single[0] is encoder + assert isinstance(single[1], np.ndarray) + assert single[2:] == (2, 2) + + multiple = unwrap_classifier_public_attrs( + [encoder, encoder], [encoder.classes_, encoder.classes_], [2, 2], 2 + ) + assert all(isinstance(value, list) for value in multiple) + + +def test_native_and_public_arrays_flatten_only_one_output() -> None: + one = np.array([[1], [2]], dtype=np.int64) + two = np.array([[1, 2], [3, 4]], dtype=np.int64) + assert native_y_array(one, dtype=np.int64).shape == (2,) + assert native_y_array(two, dtype=np.int64).shape == (2, 2) + assert squeeze_outputs(one, 1).shape == (2,) + assert squeeze_outputs(two, 2).shape == (2, 2) diff --git a/tests/test_plot_helpers.py b/tests/test_plot_helpers.py index 224ff96..9e2bb04 100644 --- a/tests/test_plot_helpers.py +++ b/tests/test_plot_helpers.py @@ -1,41 +1,35 @@ -"""Unit tests for ``sgtlearn._export`` private helpers.""" +"""Nontrivial routing contracts used by public tree plotting.""" + from __future__ import annotations -from matplotlib.patches import Rectangle -import pytest -import matplotlib -import matplotlib.pyplot as plt -from matplotlib.patches import FancyArrowPatch import numpy as np from sklearn.datasets import make_classification -from sgtlearn import SGTClassifier -from sgtlearn._export import _merge_routing_regions, _pair_switch_boundaries, _route_samples, _compute_layout_leafcounter, _draw_leaf_text, _draw_internal_panel, _draw_arrow_edge +from sgtlearn import SGTClassifier +from sgtlearn._export import ( + _merge_routing_regions, + _pair_switch_boundaries, + _route_samples, +) from tests.constants import TEST_TAO_N_RUNS -def test_merge_two_bins_same_partition_merges(): - regions = _merge_routing_regions( +def test_merge_adjacent_bins_only_when_partition_matches() -> None: + assert _merge_routing_regions( thresholds=[0.5], bin_to_partition=[0, 0], x_min=-1.0, x_max=1.0 - ) - assert regions == [(-1.0, 1.0, 0)] - - -def test_merge_two_bins_different_partition_two_slabs(): - regions = _merge_routing_regions( + ) == [(-1.0, 1.0, 0)] + assert _merge_routing_regions( thresholds=[0.5], bin_to_partition=[0, 1], x_min=-1.0, x_max=1.0 - ) - assert regions == [(-1.0, 0.5, 0), (0.5, 1.0, 1)] + ) == [(-1.0, 0.5, 0), (0.5, 1.0, 1)] -def test_merge_non_contiguous_same_partition_keeps_separate(): - regions = _merge_routing_regions( +def test_merge_keeps_noncontiguous_partition_regions_separate() -> None: + assert _merge_routing_regions( thresholds=[-0.5, 0.0, 0.5], bin_to_partition=[0, 1, 0, 1], x_min=-1.0, x_max=1.0, - ) - assert regions == [ + ) == [ (-1.0, -0.5, 0), (-0.5, 0.0, 1), (0.0, 0.5, 0), @@ -43,35 +37,13 @@ def test_merge_non_contiguous_same_partition_keeps_separate(): ] -def test_merge_consecutive_runs_merge_within_run(): - regions = _merge_routing_regions( - thresholds=[-0.5, 0.0, 0.5, 0.75], - bin_to_partition=[0, 0, 1, 1, 0], - x_min=-1.0, - x_max=1.0, - ) - assert regions == [ - (-1.0, 0.0, 0), - (0.0, 0.75, 1), - (0.75, 1.0, 0), - ] - - -def test_merge_empty_thresholds_one_slab(): - regions = _merge_routing_regions( - thresholds=[], bin_to_partition=[0], x_min=-1.0, x_max=1.0 - ) - assert regions == [(-1.0, 1.0, 0)] - - -def test_merge_trailing_nan_bin_uses_finite_bins_only(): - regions = _merge_routing_regions( +def test_merge_ignores_trailing_nan_bin() -> None: + assert _merge_routing_regions( thresholds=[-0.5, 0.0, 0.5], bin_to_partition=[0, 1, 0, 1, 0], x_min=-1.0, x_max=1.0, - ) - assert regions == [ + ) == [ (-1.0, -0.5, 0), (-0.5, 0.0, 1), (0.0, 0.5, 0), @@ -79,82 +51,47 @@ def test_merge_trailing_nan_bin_uses_finite_bins_only(): ] -def test_merge_x_min_greater_than_first_threshold_clamps_left_edge(): - regions = _merge_routing_regions( +def test_merge_clamps_regions_to_observed_range() -> None: + assert _merge_routing_regions( thresholds=[-2.0, 0.0], bin_to_partition=[0, 1, 0], x_min=-1.0, x_max=1.0, - ) - assert regions[0] == (-1.0, 0.0, 1) - assert regions[1] == (0.0, 1.0, 0) - assert len(regions) == 2 - - -def test_pair_switch_boundaries_only_keeps_partition_changes(): - cells = [(0.0, 1.0, 0.5), (1.0, 2.0, 1.5), (2.0, 3.0, 2.5)] + ) == [(-1.0, 0.0, 1), (0.0, 1.0, 0)] - x_partitions = np.array([[0, 0], [1, 1], [1, 1]]) - assert _pair_switch_boundaries(cells, x_partitions, axis=0) == [1.0] - y_partitions = np.array([[0, 1, 1], [0, 1, 1]]) - assert _pair_switch_boundaries(cells, y_partitions, axis=1) == [1.0] - - both_partitions = np.array([[0, 0, 1], [0, 1, 1], [1, 1, 1]]) - assert _pair_switch_boundaries(cells, both_partitions, axis=0) == [1.0, 2.0] - assert _pair_switch_boundaries(cells, both_partitions, axis=1) == [1.0, 2.0] +def test_pair_switch_boundaries_only_include_partition_changes() -> None: + cells: list[tuple[float, float, object]] = [ + (0.0, 1.0, 0.5), + (1.0, 2.0, 1.5), + (2.0, 3.0, 2.5), + ] + partitions = np.array([[0, 0, 1], [0, 1, 1], [1, 1, 1]]) + assert _pair_switch_boundaries(cells, partitions, axis=0) == [1.0, 2.0] + assert _pair_switch_boundaries(cells, partitions, axis=1) == [1.0, 2.0] -def _fitted_clf(): +def test_route_samples_partitions_every_parent_row_once() -> None: X, y = make_classification(n_samples=200, n_features=4, random_state=0) - return SGTClassifier( - max_depth=2, inner_max_depth=2, inner_max_leaf_nodes=8, random_state=0, + estimator = SGTClassifier( + max_depth=2, + inner_max_depth=2, + inner_max_leaf_nodes=8, + random_state=0, tao_n_runs=TEST_TAO_N_RUNS, - ).fit(X, y), X - - -def test_route_samples_root_sees_all_rows(): - est, X = _fitted_clf() - tree = est.tree_export() + ).fit(X, y) + tree = estimator.tree_export() reach = _route_samples(tree, X) - root = tree["root_index"] - assert len(reach[root]) == X.shape[0] + nodes = {node["id"]: node for node in tree["nodes"]} - -def test_route_samples_children_partition_parents_rows(): - est, X = _fitted_clf() - tree = est.tree_export() - reach = _route_samples(tree, X) - nodes_by_id = {n["id"]: n for n in tree["nodes"]} - for nid, node in nodes_by_id.items(): + assert len(reach[tree["root_index"]]) == X.shape[0] + for node_id, node in nodes.items(): + assert reach[node_id].dtype.kind in ("i", "u") if node["is_leaf"]: continue - parent_rows = set(reach[nid].tolist()) - child_rows: set[int] = set() - for cid in node["children"]: - child_set = set(reach[cid].tolist()) - assert child_set.isdisjoint(child_rows) - child_rows |= child_set - assert child_rows == parent_rows - - -def test_route_samples_sum_at_leaves_equals_n_samples(): - est, X = _fitted_clf() - tree = est.tree_export() - reach = _route_samples(tree, X) - nodes_by_id = {n["id"]: n for n in tree["nodes"]} - leaf_total = sum( - len(reach[nid]) for nid, n in nodes_by_id.items() if n["is_leaf"] - ) - assert leaf_total == X.shape[0] - - -def test_route_samples_dtype_indices_are_int(): - est, X = _fitted_clf() - tree = est.tree_export() - reach = _route_samples(tree, X) - for arr in reach.values(): - assert arr.dtype.kind in ("i", "u") + child_rows = [set(reach[child].tolist()) for child in node["children"]] + assert set().union(*child_rows) == set(reach[node_id].tolist()) + assert sum(map(len, child_rows)) == len(reach[node_id]) def test_route_samples_replays_pair_missing_edges() -> None: @@ -162,18 +99,41 @@ def test_route_samples_replays_pair_missing_edges() -> None: "root_index": 0, "nodes": [ { - "id": 0, "is_leaf": False, "routing_kind": "pair", - "features": [0, 1], "children": [1, 2, 3], + "id": 0, + "is_leaf": False, + "routing_kind": "pair", + "features": [0, 1], + "children": [1, 2, 3], "bin_to_partition": [0, 1, 2, 2, 2], "pair_axes": [ {"kind": "continuous", "columns": [0]}, {"kind": "continuous", "columns": [1]}, ], "pair_inner_tree": [ - {"id": 0, "is_leaf": False, "axis": 0, "kind": "continuous", "feature": 0, "threshold": 0.0, "left": 1, "right": 2, "missing": 3}, + { + "id": 0, + "is_leaf": False, + "axis": 0, + "kind": "continuous", + "feature": 0, + "threshold": 0.0, + "left": 1, + "right": 2, + "missing": 3, + }, {"id": 1, "is_leaf": True, "bin": 0}, {"id": 2, "is_leaf": True, "bin": 1}, - {"id": 3, "is_leaf": False, "axis": 1, "kind": "continuous", "feature": 1, "threshold": 0.0, "left": 4, "right": 5, "missing": 6}, + { + "id": 3, + "is_leaf": False, + "axis": 1, + "kind": "continuous", + "feature": 1, + "threshold": 0.0, + "left": 4, + "right": 5, + "missing": 6, + }, {"id": 4, "is_leaf": True, "bin": 2}, {"id": 5, "is_leaf": True, "bin": 3}, {"id": 6, "is_leaf": True, "bin": 4}, @@ -184,410 +144,12 @@ def test_route_samples_replays_pair_missing_edges() -> None: {"id": 3, "is_leaf": True, "children": []}, ], } - X = np.array([[-1.0, 1.0], [1.0, 1.0], [np.nan, -1.0], [np.nan, 1.0], [np.nan, np.nan]]) + X = np.array( + [[-1.0, 1.0], [1.0, 1.0], [np.nan, -1.0], [np.nan, 1.0], [np.nan, np.nan]] + ) + reach = _route_samples(tree, X) + assert reach[1].tolist() == [0] assert reach[2].tolist() == [1] assert reach[3].tolist() == [2, 3, 4] - - -def _toy_tree() -> dict: - """Hand-rolled tree dict matching tree_export()'s shape for layout tests. - - Structure:: - - 0 (root, internal, depth=0) -> [1, 2] - 1 (internal, depth=1) -> [3, 4] - 2 (leaf, depth=1) - 3 (leaf, depth=2) - 4 (leaf, depth=2) - """ - return { - "num_partitions": 2, - "num_nodes": 5, - "root_index": 0, - "criterion": "gini", - "nodes": [ - {"id": 0, "depth": 0, "is_leaf": False, "children": [1, 2]}, - {"id": 1, "depth": 1, "is_leaf": False, "children": [3, 4]}, - {"id": 2, "depth": 1, "is_leaf": True, "children": []}, - {"id": 3, "depth": 2, "is_leaf": True, "children": []}, - {"id": 4, "depth": 2, "is_leaf": True, "children": []}, - ], - } - - -def test_layout_returns_one_position_per_node(): - layout = _compute_layout_leafcounter(_toy_tree(), max_depth=None) - assert set(layout) == {0, 1, 2, 3, 4} - - -def test_layout_x_in_unit_interval(): - layout = _compute_layout_leafcounter(_toy_tree(), max_depth=None) - for x, _y in layout.values(): - assert 0.0 <= x <= 1.0 - - -def test_layout_y_top_to_bottom_by_depth(): - layout = _compute_layout_leafcounter(_toy_tree(), max_depth=None) - assert layout[0][1] > layout[1][1] - assert layout[0][1] > layout[2][1] - assert layout[1][1] > layout[3][1] - assert layout[1][1] > layout[4][1] - - -def test_layout_parent_x_centered_over_children(): - layout = _compute_layout_leafcounter(_toy_tree(), max_depth=None) - expected = (layout[3][0] + layout[4][0]) / 2 - assert layout[1][0] == pytest.approx(expected, abs=1e-9) - expected_root = (layout[1][0] + layout[2][0]) / 2 - assert layout[0][0] == pytest.approx(expected_root, abs=1e-9) - - -def test_layout_max_depth_truncates_subtree(): - layout = _compute_layout_leafcounter(_toy_tree(), max_depth=1) - assert 3 not in layout - assert 4 not in layout - assert 1 in layout - - -def test_layout_single_node_tree(): - tree = { - "num_partitions": 2, - "num_nodes": 1, - "root_index": 0, - "criterion": "gini", - "nodes": [{"id": 0, "depth": 0, "is_leaf": True, "children": []}], - } - layout = _compute_layout_leafcounter(tree, max_depth=None) - assert set(layout) == {0} - assert layout[0][0] == pytest.approx(0.5, abs=1e-9) - - -def test_draw_arrow_edge_returns_fancyarrowpatch(): - fig, ax = plt.subplots() - patch = _draw_arrow_edge( - fig=fig, - parent_xy=(0.3, 0.8), - parent_h=0.1, - child_xy=(0.5, 0.3), - child_h=0.05, - color="#E8A0BF", - host_ax=ax, - ) - assert isinstance(patch, FancyArrowPatch) - assert patch in ax.patches - plt.close(fig) - - -def test_draw_arrow_edge_uses_supplied_color(): - fig, ax = plt.subplots() - patch = _draw_arrow_edge( - fig=fig, - parent_xy=(0.0, 0.0), - parent_h=0.0, - child_xy=(1.0, 1.0), - child_h=0.0, - color="#FAC898", - host_ax=ax, - ) - assert tuple(patch.get_edgecolor())[:3] == pytest.approx( - matplotlib.colors.to_rgb("#FAC898"), abs=1e-6 - ) - plt.close(fig) - - -def _clf_leaf_node(): - return { - "id": 5, - "depth": 2, - "is_leaf": True, - "n_samples": 42, - "impurity": 0.123, - "class_counts": [[10, 32]], - "children": [], - } - - -def _reg_leaf_node(): - return { - "id": 5, - "depth": 2, - "is_leaf": True, - "n_samples": 42, - "impurity": 17.5, - "value": -3.14, - "children": [], - } - - -def test_draw_leaf_text_classifier_bold_class_label(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_clf_leaf_node(), - is_classifier=True, - class_names=["neg", "pos"], - criterion="gini", - precision=2, - fontsize=10, - color="#E8A0BF", - label="feature", - impurity=False, - ) - bold_texts = [ - a for a in artists - if hasattr(a, "get_text") and a.get_fontweight() == "bold" - ] - assert any(a.get_text() == "pos" for a in bold_texts) - plt.close(fig) - - -def test_draw_leaf_text_regressor_value_formatted_to_precision(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_reg_leaf_node(), - is_classifier=False, - class_names=None, - criterion="squared_error", - precision=2, - fontsize=10, - color="#FAC898", - label="feature", - impurity=False, - ) - texts = [a.get_text() for a in artists if hasattr(a, "get_text")] - assert any(t == "-3.14" for t in texts) - plt.close(fig) - - -def test_draw_leaf_text_label_all_adds_n_subtitle(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_clf_leaf_node(), - is_classifier=True, - class_names=["neg", "pos"], - criterion="gini", - precision=2, - fontsize=10, - color="#E8A0BF", - label="all", - impurity=False, - ) - texts = [a.get_text() for a in artists if hasattr(a, "get_text")] - assert any("n = 42" in t for t in texts) - plt.close(fig) - - -def test_draw_leaf_text_label_all_with_impurity_adds_criterion_line(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_clf_leaf_node(), - is_classifier=True, - class_names=["neg", "pos"], - criterion="gini", - precision=2, - fontsize=10, - color="#E8A0BF", - label="all", - impurity=True, - ) - texts = [a.get_text() for a in artists if hasattr(a, "get_text")] - assert any("gini = 0.12" in t for t in texts) - plt.close(fig) - - -def test_draw_leaf_text_label_none_suppresses_subtitle(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_clf_leaf_node(), - is_classifier=True, - class_names=["neg", "pos"], - criterion="gini", - precision=2, - fontsize=10, - color="#E8A0BF", - label="none", - impurity=True, - ) - texts = [a.get_text() for a in artists if hasattr(a, "get_text")] - assert not any("n =" in t or "gini =" in t for t in texts) - plt.close(fig) - - -def test_draw_leaf_text_color_propagates(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_leaf_text( - host_ax=ax, - x=0.5, - y=0.5, - node=_clf_leaf_node(), - is_classifier=True, - class_names=["neg", "pos"], - criterion="gini", - precision=2, - fontsize=10, - color="#E8A0BF", - label="feature", - impurity=False, - ) - bold_texts = [ - a for a in artists - if hasattr(a, "get_text") and a.get_fontweight() == "bold" - ] - expected_rgb = matplotlib.colors.to_rgb("#E8A0BF") - for t in bold_texts: - c = matplotlib.colors.to_rgb(t.get_color()) - assert c == pytest.approx(expected_rgb, abs=1e-6) - plt.close(fig) - - -def _internal_node(): - return { - "id": 0, - "depth": 0, - "is_leaf": False, - "feature": 0, - "thresholds": [-0.5, 0.5], - "bin_to_partition": [0, 1, 0], - "bin_sample_counts": [30, 40, 20], - "n_samples": 90, - "impurity": 0.5, - "children": [1, 2], - } - - -def test_draw_internal_panel_slabs_only_when_no_X(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - palette = ["#E8A0BF", "#FAC898"] - artists = _draw_internal_panel( - host_ax=ax, - center=(0.5, 0.5), - size=(0.3, 0.12), - node=_internal_node(), - palette=palette, - feature_values=None, - feat_names=["f0"], - n_hist_bins=20, - precision=2, - fontsize=10, - label="feature", - ) - inset_axes_objs = [a for a in artists if hasattr(a, "axvspan")] - assert inset_axes_objs, "expected an inset Axes in returned artists" - inset = inset_axes_objs[0] - # axvspan adds a Polygon to inset.patches; non-contiguous [0,1,0] - # bin_to_partition yields 3 slabs. - assert len(inset.patches) == 3 - plt.close(fig) - - -def test_draw_internal_panel_histogram_overlay_when_X_provided(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - palette = ["#E8A0BF", "#FAC898"] - feat_vals = np.linspace(-1.0, 1.0, 200) - artists = _draw_internal_panel( - host_ax=ax, - center=(0.5, 0.5), - size=(0.3, 0.12), - node=_internal_node(), - palette=palette, - feature_values=feat_vals, - feat_names=["f0"], - n_hist_bins=20, - precision=2, - fontsize=10, - label="feature", - ) - inset_axes_objs = [a for a in artists if hasattr(a, "axvspan")] - inset = inset_axes_objs[0] - bars = [p for p in inset.patches if isinstance(p, Rectangle)] - assert len(bars) >= 15 - plt.close(fig) - - -def test_draw_internal_panel_label_none_hides_sample_count(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_internal_panel( - host_ax=ax, - center=(0.5, 0.5), - size=(0.3, 0.12), - node=_internal_node(), - palette=["#E8A0BF", "#FAC898"], - feature_values=None, - feat_names=["f0"], - n_hist_bins=20, - precision=2, - fontsize=10, - label="none", - ) - text_artists = [ - a for a in artists - if hasattr(a, "get_text") and not hasattr(a, "axvspan") - ] - assert not any("n=" in t.get_text() for t in text_artists) - plt.close(fig) - - -def test_draw_internal_panel_label_feature_shows_sample_count(): - fig, ax = plt.subplots() - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.set_axis_off() - artists = _draw_internal_panel( - host_ax=ax, - center=(0.5, 0.5), - size=(0.3, 0.12), - node=_internal_node(), - palette=["#E8A0BF", "#FAC898"], - feature_values=None, - feat_names=["f0"], - n_hist_bins=20, - precision=2, - fontsize=10, - label="feature", - ) - text_artists = [a for a in artists if hasattr(a, "get_text")] - assert any("n=90" in t.get_text() for t in text_artists) - plt.close(fig) diff --git a/tests/test_plot_tree.py b/tests/test_plot_tree.py index 974af91..e95e357 100644 --- a/tests/test_plot_tree.py +++ b/tests/test_plot_tree.py @@ -1,16 +1,18 @@ """Smoke tests for ``sgtlearn._export.plot_tree``.""" + from __future__ import annotations import matplotlib import numpy as np from sklearn.tree import DecisionTreeClassifier + matplotlib.use("Agg") # headless import matplotlib.pyplot as plt +import pandas as pd import pytest from sklearn.datasets import make_classification, make_regression from sklearn.exceptions import NotFittedError -from matplotlib.patches import FancyArrowPatch, Rectangle from sgtlearn import SGTClassifier, SGTRegressor, plot_tree from tests.constants import TEST_TAO_N_RUNS @@ -99,13 +101,6 @@ def _is_number(s: str) -> bool: plt.close("all") -def test_plot_tree_draws_edges(fitted_classifier): - artists = plot_tree(fitted_classifier) - arrows = [a for a in artists if isinstance(a, FancyArrowPatch)] - assert len(arrows) >= 1 - plt.close("all") - - def test_plot_tree_label_all_adds_metadata(fitted_classifier): artists = plot_tree(fitted_classifier, label="all", impurity=True) text_artists = [a for a in artists if hasattr(a, "get_text")] @@ -126,18 +121,6 @@ def test_plot_tree_label_none_suppresses_subtitles(fitted_classifier): plt.close("all") -def test_plot_tree_proportion_does_not_crash(fitted_classifier): - artists = plot_tree(fitted_classifier, proportion=True) - assert artists - plt.close("all") - - -def test_plot_tree_custom_cmap(fitted_classifier): - artists = plot_tree(fitted_classifier, cmap="viridis") - assert artists - plt.close("all") - - def test_plot_tree_custom_feature_names(fitted_classifier): names = [f"feat_{i}" for i in range(fitted_classifier.n_features_in_)] artists = plot_tree(fitted_classifier, feature_names=names) @@ -146,6 +129,30 @@ def test_plot_tree_custom_feature_names(fitted_classifier): plt.close("all") +def test_plot_tree_explicit_feature_names_override_stored_names(): + X = pd.DataFrame({"stored_name": np.arange(20, dtype=float)}) + y = np.repeat([0, 1], 10) + estimator = SGTClassifier(max_depth=1, tao_n_runs=0, random_state=0).fit( + X, y, feature_dict={"stored_name": ["stored_name"]} + ) + + stored = plot_tree(estimator) + explicit = plot_tree(estimator, feature_names=["explicit_name"]) + stored_text = " ".join( + [a.get_text() for a in stored if hasattr(a, "get_text")] + + [a.get_xlabel() for a in stored if hasattr(a, "get_xlabel")] + ) + explicit_text = " ".join( + [a.get_text() for a in explicit if hasattr(a, "get_text")] + + [a.get_xlabel() for a in explicit if hasattr(a, "get_xlabel")] + ) + + assert "stored_name" in stored_text + assert "explicit_name" in explicit_text + assert "stored_name" not in explicit_text + plt.close("all") + + def test_plot_tree_class_names(fitted_classifier): artists = plot_tree(fitted_classifier, class_names=["neg", "pos"]) text_artists = [a for a in artists if hasattr(a, "get_text")] @@ -163,34 +170,19 @@ def test_plot_tree_reuses_existing_axes(fitted_classifier): plt.close("all") -def test_plot_tree_fontsize_passes_through(fitted_classifier): - artists = plot_tree(fitted_classifier, fontsize=6) - text_artists = [a for a in artists if hasattr(a, "get_fontsize")] - sizes = {a.get_fontsize() for a in text_artists if a.get_text()} - assert 6 in sizes - plt.close("all") - - def test_plot_tree_with_X_renders_fine_histograms(fitted_classifier): X, _ = make_classification(n_samples=200, n_features=4, random_state=0) - - def total_bars(artists): - # Each internal panel is an inset Axes; histogram bars are Rectangle - # patches added to the inset. axvspan slabs are Polygons, not - # Rectangles, so they don't count here. - total = 0 - for a in artists: - if hasattr(a, "patches"): - total += sum(1 for p in a.patches if isinstance(p, Rectangle)) - return total - - no_x_bars = total_bars(plot_tree(fitted_classifier)) + without_histograms = plot_tree(fitted_classifier) plt.close("all") - with_x_bars = total_bars(plot_tree(fitted_classifier, X=X)) + with_histograms = plot_tree(fitted_classifier, X=X) plt.close("all") - # With X passed, each internal panel adds ~20 histogram bars; - # without X, only slab Polygons are drawn (no Rectangle bars). - assert with_x_bars > no_x_bars + + def patch_count(artists): + return sum( + len(artist.patches) for artist in artists if hasattr(artist, "patches") + ) + + assert patch_count(with_histograms) > patch_count(without_histograms) def test_plot_tree_X_shape_mismatch_raises(fitted_classifier): @@ -198,29 +190,8 @@ def test_plot_tree_X_shape_mismatch_raises(fitted_classifier): with pytest.raises(ValueError): plot_tree(fitted_classifier, X=bad_X) - -def test_plot_tree_leaf_uses_partition_color(fitted_classifier): - artists = plot_tree(fitted_classifier, cmap="Pastel1") - bold = [ - a for a in artists - if hasattr(a, "get_text") and a.get_fontweight() == "bold" and a.get_text() - ] - assert bold, "expected at least one bold leaf text" - cmap = matplotlib.colormaps["Pastel1"] - expected_p0 = matplotlib.colors.to_rgb(cmap(0.0)) - expected_p1 = matplotlib.colors.to_rgb(cmap(1.0)) - leaf_colors = { - tuple(matplotlib.colors.to_rgb(t.get_color())) for t in bold - } - for c in leaf_colors: - assert ( - all(abs(a - b) < 1e-6 for a, b in zip(c, expected_p0)) - or all(abs(a - b) < 1e-6 for a, b in zip(c, expected_p1)) - ), ( - f"leaf color {c} not one of partition colors " - f"{expected_p0}, {expected_p1}" - ) - plt.close("all") + with pytest.raises(ValueError, match="X must be 2-D"): + plot_tree(fitted_classifier, X=np.zeros(fitted_classifier.n_features_in_)) def test_plot_tree_pair_heatmap_reuses_exported_router_and_axes(): @@ -246,15 +217,8 @@ def test_plot_tree_pair_heatmap_reuses_exported_router_and_axes(): panels = [artist for artist in artists if hasattr(artist, "patches")] assert ax in fig.axes assert panels and len(panels[0].patches) == 4 - assert not panels[0].collections - assert {panel.get_label() for panel in panels} >= { - "pair-x-histogram", - "pair-y-histogram", - } assert panels[0].get_xlabel() == "first" assert panels[0].get_ylabel() == "second" - assert [tick.get_text() for tick in panels[0].get_xticklabels()] == ["0.000"] - assert [tick.get_text() for tick in panels[0].get_yticklabels()] == ["0.000"] plt.close(fig) @@ -263,8 +227,11 @@ def test_plot_tree_pair_heatmap_renders_categories_and_missing_cells(pair_kind): categories = [[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]] if pair_kind == "mixed": states = np.array( - [[value, *category] for value in [-1.0, 1.0, np.nan] - for category in categories] + [ + [value, *category] + for value in [-1.0, 1.0, np.nan] + for category in categories + ] ) feature_dict = {0: [0], 1: [1, 2]} else: @@ -285,14 +252,9 @@ def test_plot_tree_pair_heatmap_renders_categories_and_missing_cells(pair_kind): ).fit(X, y, feature_dict=feature_dict) fig, ax = plt.subplots() - artists = plot_tree(est, X=X, cmap="tab10", ax=ax) + artists = plot_tree(est, X=X, ax=ax) panel = next(artist for artist in artists if hasattr(artist, "patches")) assert len(panel.patches) == 9 # 3 × 3, including both missing margins/corner - assert len({patch.get_facecolor() for patch in panel.patches}) == 9 - widths = [patch.get_width() for patch in panel.patches] - heights = [patch.get_height() for patch in panel.patches] - assert min(widths) < 0.5 * max(widths) - assert min(heights) < 0.5 * max(heights) x_labels = [tick.get_text() for tick in panel.get_xticklabels()] y_labels = [tick.get_text() for tick in panel.get_yticklabels()] assert "NaN" in x_labels @@ -303,9 +265,7 @@ def test_plot_tree_pair_heatmap_renders_categories_and_missing_cells(pair_kind): def test_plot_tree_pair_heatmap_renders_without_training_data(): - states = np.array( - [[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]] - ) + states = np.array([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]]) X = np.repeat(states, [40, 30, 30, 28], axis=0) y = np.repeat([0, 1, 1, 0], [40, 30, 30, 28]) est = SGTClassifier( diff --git a/tests/test_random_sgforest_validation.py b/tests/test_random_sgforest_validation.py new file mode 100644 index 0000000..8feea9b --- /dev/null +++ b/tests/test_random_sgforest_validation.py @@ -0,0 +1,134 @@ +"""Shared validation contracts for random SGT forests.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from sklearn.utils import check_random_state + +from sgtlearn import RandomSGForestClassifier, RandomSGForestRegressor, SGTRegressor +from sgtlearn.ensemble._random_sgforest import _n_samples_bootstrap + + +@pytest.mark.parametrize( + ("max_samples", "expected"), + [(None, 10), (1, 1), (10, 10), (0.01, 1), (0.5, 5), (1.0, 10)], +) +def test_bootstrap_sample_count_boundaries( + max_samples: float | None, expected: int +) -> None: + assert _n_samples_bootstrap(10, max_samples) == expected + + +@pytest.mark.parametrize("max_samples", [0, 11, 0.0, 1.01]) +def test_bootstrap_sample_count_rejects_invalid_values( + max_samples: float, +) -> None: + with pytest.raises(ValueError, match="max_samples"): + _n_samples_bootstrap(10, max_samples) + + +def test_forest_rejects_invalid_estimator_and_bootstrap_configuration() -> None: + X = np.array([[0.0], [1.0]]) + y = np.array([0, 1]) + with pytest.raises(ValueError, match="n_estimators"): + RandomSGForestClassifier(n_estimators=0).fit(X, y) + with pytest.raises(ValueError, match="bootstrap"): + RandomSGForestClassifier(bootstrap=False, max_samples=1).fit(X, y) + + +def test_forest_predict_rejects_feature_count_mismatch() -> None: + X = np.array([[0.0, 0.0], [1.0, 1.0], [0.1, 0.2], [0.9, 0.8]]) + y = np.array([0, 1, 0, 1]) + forest = RandomSGForestClassifier( + n_estimators=1, + bootstrap=False, + max_features=None, + tao_n_runs=0, + random_state=0, + ).fit(X, y) + + with pytest.raises(ValueError, match="expecting 2 features"): + forest.predict(X[:, :1]) + + +def test_bootstrap_keeps_samples_targets_and_weights_aligned() -> None: + X = np.arange(16, dtype=np.float32).reshape(8, 2) + y = np.array([0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0]) + sample_weight = np.arange(1.0, 9.0) + forest_seed = 4 + forest = RandomSGForestRegressor( + n_estimators=1, + bootstrap=True, + max_samples=6, + random_state=forest_seed, + max_depth=2, + inner_max_depth=1, + max_features=None, + coordinate_descent_smart_init=False, + tao_n_runs=0, + ).fit(X, y, sample_weight=sample_weight) + + tree_seed = int(check_random_state(forest_seed).randint(np.iinfo(np.int32).max)) + indices = check_random_state(tree_seed).randint(0, len(X), 6, dtype=np.int32) + manual = SGTRegressor( + random_state=tree_seed, + max_depth=2, + inner_max_depth=1, + max_features=None, + coordinate_descent_smart_init=False, + tao_n_runs=0, + ).fit(X[indices], y[indices], sample_weight=sample_weight[indices]) + + np.testing.assert_allclose(forest.estimators_[0].predict(X), manual.predict(X)) + + +def test_classifier_probabilities_are_the_mean_in_shared_label_space() -> None: + X = np.arange(40, dtype=np.float32).reshape(20, 2) + forest_seed = 2 + seed_rng = check_random_state(forest_seed) + tree_seeds = [ + int(seed_rng.randint(np.iinfo(np.int32).max)) for _ in range(3) + ] + bootstrap_rows = [ + check_random_state(seed).randint(0, len(X), 3, dtype=np.int32) + for seed in tree_seeds + ] + rare_row = next( + row + for row in range(len(X)) + if any(row in rows for rows in bootstrap_rows) + and any(row not in rows for rows in bootstrap_rows) + ) + y = np.full(len(X), "common") + y[rare_row] = "rare" + forest = RandomSGForestClassifier( + n_estimators=3, + bootstrap=True, + max_samples=3, + max_features=None, + tao_n_runs=0, + random_state=forest_seed, + ).fit(X, y) + + expected = np.mean([tree.predict_proba(X) for tree in forest.estimators_], axis=0) + + assert any(rare_row not in rows for rows in bootstrap_rows) + np.testing.assert_array_equal(forest.classes_, ["common", "rare"]) + assert all(tree.predict_proba(X).shape == (len(X), 2) for tree in forest.estimators_) + np.testing.assert_allclose(forest.predict_proba(X), expected) + + +def test_regression_forest_preserves_multioutput_shape() -> None: + X = np.arange(24, dtype=np.float32).reshape(12, 2) + y = np.column_stack([X[:, 0], 10.0 + X[:, 1]]) + forest = RandomSGForestRegressor( + n_estimators=2, + bootstrap=False, + max_features=None, + tao_n_runs=0, + random_state=0, + ).fit(X, y) + + assert forest.predict(X).shape == (12, 2) diff --git a/tests/test_tao.py b/tests/test_tao.py index 7b39df1..97fbc8e 100644 --- a/tests/test_tao.py +++ b/tests/test_tao.py @@ -1,66 +1,20 @@ -"""Unit tests for :mod:`sgtlearn.tao`. - -No-regression matrices apply to single-tree base estimators only. Random forests -are smoke-tested for successful refinement and intact prediction. -""" +"""Unit tests for :mod:`sgtlearn.tao`.""" from __future__ import annotations -import zlib from typing import Any, Callable, Mapping, Optional, Tuple import numpy as np import pytest -from sklearn.datasets import ( - load_breast_cancer, - load_iris, - make_classification, - make_regression, -) +from sklearn.datasets import load_iris, make_regression from sklearn.exceptions import NotFittedError -from sklearn.metrics import accuracy_score from sgtlearn import SGTClassifier, SGTRegressor, tao from sgtlearn.ensemble import RandomSGForestClassifier, RandomSGForestRegressor -from sgtlearn._weights import ( - effective_sample_weight_classification, - normalize_sample_weight, -) from tests.constants import TEST_TAO_N_RUNS -from tests.discretizer_grid import n_outputs_params pytest.importorskip("sklearn") -RTOL = 1e-9 - - -def _no_regression_tol(before: float, *, weighted: bool) -> float: - """Allow tiny float32 vs float64 drift when sample/class weights are active.""" - tol = max(RTOL, 1e-6 * (1.0 + abs(before))) - if weighted: - tol = max(tol, 2e-4) - return tol - - -def _expand_targets(y: np.ndarray, n_outputs: int, *, kind: str) -> np.ndarray: - """Build a 1-D or multi-output target matrix from a single-output vector.""" - y0 = np.asarray(y) - if n_outputs == 1: - return y0 - if kind == "classification": - rng = np.random.default_rng(0) - extras = [ - rng.integers(0, len(np.unique(y0)), size=y0.shape[0]) - for _ in range(n_outputs - 1) - ] - return np.column_stack([y0, *extras]) - rng = np.random.default_rng(0) - extras = [ - y0 * (0.5 + 0.25 * i) + rng.normal(0.0, 1.0, size=y0.shape) - for i in range(1, n_outputs) - ] - return np.column_stack([y0.astype(np.float64), *extras]) - def _fit_classifier( X: np.ndarray, @@ -110,56 +64,22 @@ def _fit_regressor( return est -def _classification_data(name: str) -> Tuple[np.ndarray, np.ndarray]: - if name == "iris": - X, y = load_iris(return_X_y=True) - elif name == "breast_cancer": - X, y = load_breast_cancer(return_X_y=True) - elif name == "multiclass": - X, y = make_classification( - n_samples=600, - n_features=20, - n_informative=10, - n_redundant=4, - n_classes=4, - random_state=7, - ) - else: - raise ValueError(f"unknown classification dataset {name!r}") +def _classification_data() -> Tuple[np.ndarray, np.ndarray]: + X, y = load_iris(return_X_y=True) return np.asarray(X, dtype=np.float64), y -def _regression_data(name: str) -> Tuple[np.ndarray, np.ndarray]: - if name == "low_noise": - X, y = make_regression( - n_samples=400, - n_features=10, - n_informative=6, - noise=5.0, - random_state=0, - ) - elif name == "high_noise": - X, y = make_regression( - n_samples=400, - n_features=10, - n_informative=6, - noise=25.0, - random_state=3, - ) - else: - raise ValueError(f"unknown regression dataset {name!r}") +def _regression_data() -> Tuple[np.ndarray, np.ndarray]: + X, y = make_regression( + n_samples=400, + n_features=10, + n_informative=6, + noise=5.0, + random_state=0, + ) return np.asarray(X, dtype=np.float64), y -def _stable_seed(*parts: str) -> int: - return zlib.adler32("|".join(parts).encode()) & 0xFFFFFFFF - - -def _sample_weights(n_samples: int, seed: int) -> np.ndarray: - rng = np.random.default_rng(seed) - return rng.uniform(0.5, 2.0, size=n_samples) - - def test_feature_importances_are_unavailable_after_tao_and_reset_on_refit() -> None: X, y = load_iris(return_X_y=True) clf = SGTClassifier(tao_n_runs=0, random_state=0).fit(X, y) @@ -176,201 +96,19 @@ def test_feature_importances_are_unavailable_after_tao_and_reset_on_refit() -> N assert clf.feature_importances_.shape == (X.shape[1],) -def _classification_training_score( - tree: SGTClassifier, - X: np.ndarray, - y: np.ndarray, - sample_weight: Optional[np.ndarray], -) -> float: - pred = tree.predict(X) - y_arr = np.asarray(y) - if tree.class_weight is None and sample_weight is None: - if y_arr.ndim == 1: - return float(accuracy_score(y_arr, pred)) - return float(np.mean(pred == y_arr)) - - from sgtlearn._multioutput import ( - encode_classification_targets, - label_encoders_as_list, - ) - - encoders = label_encoders_as_list(tree._le, tree.n_outputs_) - y_enc, _, _, _ = encode_classification_targets(y_arr, encoders=encoders) - if tree.class_weight is not None: - sw = effective_sample_weight_classification( - sample_weight, y_enc, tree.class_weight, tree.classes_ - ) - else: - sw = normalize_sample_weight(sample_weight, y_arr.shape[0]) - correct = (pred == y_arr).astype(np.float64) - if correct.ndim == 2: - correct = correct.mean(axis=1) - return float(np.average(correct, weights=sw)) - - -def _skewed_class_weight( - y: np.ndarray, -) -> Mapping[Any, float] | list[Mapping[Any, float]]: - """Up-weight the first observed class; list-of-dicts when ``y`` is multi-output.""" - y_arr = np.asarray(y) - if y_arr.ndim == 1: - classes = np.unique(y_arr) - return {c: (2.0 if i == 0 else 1.0) for i, c in enumerate(classes)} - return [ - { - c: (2.0 if i == 0 else 1.0) - for i, c in enumerate(np.unique(y_arr[:, o])) - } - for o in range(y_arr.shape[1]) - ] - - -def _regression_training_loss( - reg: SGTRegressor, - X: np.ndarray, - y: np.ndarray, - criterion: str, - sample_weight: Optional[np.ndarray], -) -> float: - pred = reg.predict(X) - y_arr = np.asarray(y, dtype=np.float64) - if criterion in ("squared_error", "mse"): - err = (pred - y_arr) ** 2 - else: - err = np.abs(pred - y_arr) - if err.ndim == 2: - err = err.mean(axis=1) - if sample_weight is None: - return float(np.mean(err)) - return float(np.average(err, weights=sample_weight)) - - -@pytest.mark.parametrize("n_outputs", n_outputs_params()) -@pytest.mark.parametrize("dataset", ["iris", "breast_cancer", "multiclass"]) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) -@pytest.mark.parametrize( - "use_sample_weight", [False, True], ids=["unweighted", "weighted"] -) -@pytest.mark.parametrize( - "use_class_weight", [False, True], ids=["no_class_weight", "skewed_class_weight"] -) -@pytest.mark.parametrize("n_runs", [1, 10]) -def test_tao_classification_no_regression_matrix( - dataset: str, - criterion: str, - use_sample_weight: bool, - use_class_weight: bool, - n_runs: int, - n_outputs: int, -) -> None: - """TAO must not decrease training accuracy across the classification grid.""" - X, y0 = _classification_data(dataset) - y = _expand_targets(y0, n_outputs, kind="classification") - sw = ( - _sample_weights(len(y0), seed=_stable_seed(dataset, criterion, "weighted")) - if use_sample_weight - else None - ) - class_weight = _skewed_class_weight(y) if use_class_weight else None - - tree_kwargs: dict[str, Any] = {} - if dataset == "multiclass": - tree_kwargs = dict(max_depth=6, num_partitions=3) - - tree = _fit_classifier( - X, - y, - criterion=criterion, - class_weight=class_weight, - sample_weight=sw, - **tree_kwargs, - ) - - before = _classification_training_score(tree, X, y, sw) - returned = tao.TAO_refine( - tree, - X, - y, - sample_weight=sw, - n_runs=n_runs, - lambda_=0.0, - ) - after = _classification_training_score(tree, X, y, sw) - - assert returned is tree - expected_shape = (X.shape[0],) if n_outputs == 1 else (X.shape[0], n_outputs) - assert tree.predict(X).shape == expected_shape - weighted_metric = use_sample_weight or use_class_weight - assert after >= before - _no_regression_tol(before, weighted=weighted_metric), ( - f"TAO decreased training score for dataset={dataset!r}, " - f"criterion={criterion!r}, weighted={use_sample_weight}, " - f"class_weight={use_class_weight}, n_runs={n_runs}, " - f"n_outputs={n_outputs}: " - f"{before:.6f} -> {after:.6f}" - ) - - -@pytest.mark.parametrize("n_outputs", n_outputs_params()) -@pytest.mark.parametrize("dataset", ["low_noise", "high_noise"]) -@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) -@pytest.mark.parametrize( - "use_sample_weight", [False, True], ids=["unweighted", "weighted"] -) -@pytest.mark.parametrize("n_runs", [1, 10]) -def test_tao_regression_no_regression_matrix( - dataset: str, - criterion: str, - use_sample_weight: bool, - n_runs: int, - n_outputs: int, -) -> None: - """TAO must not increase training loss across the regression grid.""" - X, y0 = _regression_data(dataset) - y = _expand_targets(y0, n_outputs, kind="regression") - sw = ( - _sample_weights(len(y0), seed=_stable_seed(dataset, criterion, "weighted")) - if use_sample_weight - else None - ) - - reg = _fit_regressor(X, y, criterion=criterion, sample_weight=sw) - - before = _regression_training_loss(reg, X, y, criterion, sw) - returned = tao.TAO_refine( - reg, - X, - y, - sample_weight=sw, - n_runs=n_runs, - lambda_=0.0, - ) - after = _regression_training_loss(reg, X, y, criterion, sw) - - assert returned is reg - expected_shape = (X.shape[0],) if n_outputs == 1 else (X.shape[0], n_outputs) - assert reg.predict(X).shape == expected_shape - tol = _no_regression_tol(before, weighted=use_sample_weight) - assert after <= before + tol, ( - f"TAO increased training loss for dataset={dataset!r}, " - f"criterion={criterion!r}, weighted={use_sample_weight}, " - f"n_runs={n_runs}, n_outputs={n_outputs}: " - f"{before:.6f} -> {after:.6f}" - ) - - @pytest.mark.parametrize( ("estimator_cls", "fit_fn", "data_fn"), [ pytest.param( SGTClassifier, lambda X, y: _fit_classifier(X, y), - lambda: _classification_data("iris"), + _classification_data, id="classifier", ), pytest.param( SGTRegressor, lambda X, y: _fit_regressor(X, y), - lambda: _regression_data("low_noise"), + _regression_data, id="regressor", ), ], @@ -424,12 +162,12 @@ def test_tao_rejects_wrong_type() -> None: [ pytest.param( lambda X, y: _fit_classifier(X, y), - lambda: _classification_data("iris"), + _classification_data, id="classifier", ), pytest.param( lambda X, y: _fit_regressor(X, y), - lambda: _regression_data("low_noise"), + _regression_data, id="regressor", ), ], @@ -439,123 +177,61 @@ def test_tao_rejects_feature_mismatch(fit_fn, data_fn) -> None: est = fit_fn(X, y) with pytest.raises(ValueError): tao.TAO_refine(est, X[:, :-1], y) + with pytest.raises(ValueError, match="samples"): + tao.TAO_refine(est, X, y[:-1]) def test_tao_accepts_check_input_false() -> None: """Callers that pre-validate arrays can skip redundant checks.""" - X, y = _classification_data("iris") - tree = _fit_classifier(X, y) - before = _classification_training_score(tree, X, y, None) - tao.TAO_refine(tree, X, y, check_input=False) - after = _classification_training_score(tree, X, y, None) - assert after >= before - RTOL - - -def _fit_forest_classifier( - X: np.ndarray, - y: np.ndarray, - *, - n_estimators: int = 5, - n_jobs: int = 1, - **kwargs: Any, -) -> RandomSGForestClassifier: - params = dict( - n_estimators=n_estimators, - criterion="gini", - max_depth=4, - min_samples_leaf=3, - inner_max_depth=4, - inner_max_leaf_nodes=16, - random_state=42, - n_jobs=n_jobs, - tao_n_runs=TEST_TAO_N_RUNS, - ) - params.update(kwargs) - return RandomSGForestClassifier(**params).fit(X, y) - - -def _fit_forest_regressor( - X: np.ndarray, - y: np.ndarray, - *, - n_estimators: int = 5, - n_jobs: int = 1, - criterion: str = "squared_error", - **kwargs: Any, -) -> RandomSGForestRegressor: - params = dict( - n_estimators=n_estimators, - criterion=criterion, - max_depth=4, - min_samples_leaf=3, - inner_max_depth=4, - inner_max_leaf_nodes=16, - random_state=42, - n_jobs=n_jobs, - tao_n_runs=TEST_TAO_N_RUNS, - ) - params.update(kwargs) - return RandomSGForestRegressor(**params).fit(X, y) + X, y = _tao_pair_interaction_data() + tree = SGTClassifier( + max_depth=1, + pairwise_candidates=1, + pairwise_penalty=1.0, + tao_n_runs=0, + random_state=0, + ).fit(X, y) + tao.TAO_refine(tree, X, y, n_runs=1, check_input=False) + assert tree.tree_export()["nodes"][0]["routing_kind"] == "pair" -@pytest.mark.parametrize("n_outputs", n_outputs_params()) @pytest.mark.parametrize( - ("forest_cls", "fit_fn", "data_fn"), + ("forest_cls", "target"), [ - pytest.param( - RandomSGForestClassifier, - lambda X, y: _fit_forest_classifier(X, y, n_estimators=8, n_jobs=1), - lambda: _classification_data("iris"), - id="classifier", - ), - pytest.param( - RandomSGForestRegressor, - lambda X, y: _fit_forest_regressor(X, y, n_estimators=8, n_jobs=1), - lambda: _regression_data("low_noise"), - id="regressor", - ), + (RandomSGForestClassifier, lambda y: y), + (RandomSGForestRegressor, lambda y: np.column_stack([y, 10.0 + y])), ], ) -def test_tao_forest_runs_and_predicts_all_samples( - forest_cls, fit_fn, data_fn, n_outputs: int -) -> None: - """TAO on a random forest completes and leaves predict() valid for every row.""" - X, y0 = data_fn() - kind = "classification" if forest_cls is RandomSGForestClassifier else "regression" - y = _expand_targets(y0, n_outputs, kind=kind) - forest = fit_fn(X, y) - - returned = tao.TAO_refine(forest, X, y, n_runs=3, n_jobs=2) - - assert returned is forest - pred = forest.predict(X) - expected = (X.shape[0],) if n_outputs == 1 else (X.shape[0], n_outputs) - assert pred.shape == expected - if forest_cls is RandomSGForestClassifier: - proba = forest.predict_proba(X) - if n_outputs == 1: - assert proba.shape == (X.shape[0], forest.n_classes_) - else: - assert len(proba) == n_outputs - for o, p in enumerate(proba): - assert p.shape == (X.shape[0], forest.n_classes_[o]) - - -def test_tao_forest_refine_mutates_in_place() -> None: - X, y = _classification_data("iris") - forest = _fit_forest_classifier(X, y, n_estimators=4) +def test_tao_refines_every_forest_tree(forest_cls, target) -> None: + X, labels = _tao_pair_interaction_data() + y = target(labels.astype(float)) + forest = forest_cls( + n_estimators=2, + bootstrap=False, + max_features=None, + max_depth=1, + inner_max_depth=2, + inner_max_leaf_nodes=4, + pairwise_candidates=1, + pairwise_penalty=1.0, + tao_n_runs=0, + random_state=0, + n_jobs=1, + ).fit(X, y) handles_before = [est._est for est in forest.estimators_] - result = tao.TAO_refine(forest, X, y, n_jobs=2) + result = tao.TAO_refine(forest, X, y, n_runs=1, lambda_=0.0, n_jobs=2) assert result is forest assert [est._est for est in forest.estimators_] == handles_before + assert all( + est.tree_export()["nodes"][0]["routing_kind"] == "pair" + for est in forest.estimators_ + ) def _tao_pair_interaction_data() -> tuple[np.ndarray, np.ndarray]: - quadrants = np.array( - [[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]] - ) + quadrants = np.array([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]]) counts = [40, 10, 30, 5] return np.repeat(quadrants, counts, axis=0), np.repeat([0, 1, 1, 0], counts) @@ -588,6 +264,39 @@ def test_tao_reconsiders_retained_classifier_pair() -> None: assert clf.score(X, y) == 1.0 +def test_tao_weights_change_the_accepted_classifier_update() -> None: + X, y = _tao_pair_interaction_data() + sample_weight = np.ones(len(y)) + sample_weight[40:50] = 20.0 + + def refine(weights: np.ndarray | None) -> SGTClassifier: + clf = SGTClassifier( + max_depth=1, + inner_max_depth=2, + inner_max_leaf_nodes=4, + pairwise_candidates=1, + pairwise_penalty=0.3, + tao_n_runs=0, + random_state=0, + ).fit(X, y) + return tao.TAO_refine( + clf, + X, + y, + sample_weight=weights, + n_runs=1, + lambda_=0.5, + ) + + unweighted = refine(None) + weighted = refine(sample_weight) + unweighted_score = np.average(unweighted.predict(X) == y, weights=sample_weight) + weighted_score = np.average(weighted.predict(X) == y, weights=sample_weight) + + assert not np.array_equal(unweighted.predict(X), weighted.predict(X)) + assert weighted_score > unweighted_score + + def test_tao_makes_forest_feature_importances_unavailable() -> None: X, y = _tao_pair_interaction_data() forest = RandomSGForestClassifier( @@ -653,12 +362,8 @@ def fit_pair() -> SGTClassifier: default_scale = fit_pair() high_scale = fit_pair() - tao.TAO_refine( - default_scale, X, y, n_runs=1, lambda_=0.3, tao_pair_scale=1.1 - ) - tao.TAO_refine( - high_scale, X, y, n_runs=1, lambda_=0.3, tao_pair_scale=2.0 - ) + tao.TAO_refine(default_scale, X, y, n_runs=1, lambda_=0.3, tao_pair_scale=1.1) + tao.TAO_refine(high_scale, X, y, n_runs=1, lambda_=0.3, tao_pair_scale=2.0) assert default_scale.tree_export()["nodes"][0]["routing_kind"] == "pair" assert default_scale.score(X, y) == 1.0 diff --git a/tests/test_tree_export.py b/tests/test_tree_export.py index 34315a8..7359db6 100644 --- a/tests/test_tree_export.py +++ b/tests/test_tree_export.py @@ -32,6 +32,15 @@ def _fitted_regressor(criterion: str = "squared_error", ).fit(X, y) +def _replay_classifier_export(estimator: SGTClassifier, X: np.ndarray) -> np.ndarray: + tree = estimator.tree_export() + reach = _route_samples(tree, X) + replay = np.empty(X.shape[0], dtype=int) + for leaf in (node for node in tree["nodes"] if node["is_leaf"]): + replay[reach[leaf["id"]]] = np.argmax(leaf["class_counts"][0]) + return replay + + def test_classifier_export_top_level_keys(): est = _fitted_classifier() tr = est.tree_export() @@ -132,6 +141,25 @@ def test_classifier_multiway_partitions(): assert 0 <= p < 3 +def test_univariate_classifier_export_replays_predictions(): + X = np.arange(20, dtype=np.float32).reshape(-1, 1) + y = np.repeat([0, 1], 10) + estimator = SGTClassifier(max_depth=2, tao_n_runs=0, random_state=0).fit(X, y) + + np.testing.assert_array_equal(_replay_classifier_export(estimator, X), y) + + +def test_categorical_classifier_export_replays_predictions_and_missing_route(): + categories = np.eye(3, dtype=np.float32) + X = np.vstack([np.repeat(categories, 8, axis=0), np.zeros((3, 3))]) + y = np.concatenate([np.repeat([0, 1, 0], 8), np.zeros(3, dtype=int)]) + estimator = SGTClassifier(max_depth=2, tao_n_runs=0, random_state=0).fit( + X, y, feature_dict={0: [0, 1, 2]} + ) + + np.testing.assert_array_equal(_replay_classifier_export(estimator, X), y) + + def test_pair_classifier_export_replays_predictions(): states = np.array([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]]) X = np.repeat(states, [40, 30, 30, 28], axis=0) @@ -141,11 +169,7 @@ def test_pair_classifier_export_replays_predictions(): tree = est.tree_export() assert tree["nodes"][0]["routing_kind"] == "pair" assert "nan_prediction_partition" not in tree["nodes"][0] - reach = _route_samples(tree, X) - replay = np.empty(X.shape[0], dtype=int) - for leaf in (node for node in tree["nodes"] if node["is_leaf"]): - replay[reach[leaf["id"]]] = np.argmax(leaf["class_counts"][0]) - np.testing.assert_array_equal(replay, est.predict(X)) + np.testing.assert_array_equal(_replay_classifier_export(est, X), est.predict(X)) def test_pair_regressor_export_replays_predictions(): diff --git a/tests/test_univariate_classification_discretizer.py b/tests/test_univariate_classification_discretizer.py index 9c83091..b158f3a 100644 --- a/tests/test_univariate_classification_discretizer.py +++ b/tests/test_univariate_classification_discretizer.py @@ -1,59 +1,34 @@ -"""Fidelity tests: ``UnivariateClassificationDiscretizer`` vs ``sklearn.tree.DecisionTreeClassifier``. - -Trains both on a single feature axis and asserts identical leaf predictions and -leaf counts across a grid of hyperparameters (see ``discretizer_grid``). -""" +"""Focused univariate-classification discretizer contracts and sklearn parity.""" import numpy as np import pytest -from itertools import product from Discretizers import UnivariateClassificationDiscretizer from sklearn.tree import DecisionTreeClassifier -from tests.discretizer_grid import ( - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - MIN_GAIN_VALUES, - MIN_LEAF_VALUES, - N_VALUES, - NUM_CLASSES_VALUES, - n_outputs_params, -) - -def classification_predict(ud: UnivariateClassificationDiscretizer, x: np.ndarray) -> np.ndarray: +def classification_predict( + ud: UnivariateClassificationDiscretizer, x: np.ndarray +) -> np.ndarray: """Predict by mapping transform() bin indices to bin predictions.""" bin_locs = ud.transform(x) bin_preds = ud.getBinPredictions() return np.asarray(bin_preds[bin_locs], dtype=np.uintp) -GRID = list( - product( - N_VALUES, - NUM_CLASSES_VALUES, - MIN_LEAF_VALUES, - MIN_GAIN_VALUES, - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - ) -) -IDS = [ - f"N={n}|C={c}|leaf={leaf}|gain={gain}|depth={depth}|max_leaf={max_leaf}" - for n, c, leaf, gain, depth, max_leaf in GRID +PARITY_CASES = [ + pytest.param("gini", 2, 1, 0.0, 0, 0, 1, id="gini-binary"), + pytest.param("entropy", 3, 1, 0.0, 0, 0, 2, id="entropy-multioutput"), + pytest.param("gini", 3, 1, 0.0, 4, 0, 1, id="depth-limited"), + pytest.param("entropy", 3, 1, 0.0, 0, 8, 2, id="leaf-limited"), ] -@pytest.mark.parametrize("n_outputs", n_outputs_params()) -@pytest.mark.parametrize("criterion", ["gini", "entropy"]) @pytest.mark.parametrize( - "n_samples,num_classes,min_leaf_size,min_gain_split,max_depth,max_leaf", - GRID, - ids=IDS, + "criterion,num_classes,min_leaf_size,min_gain_split,max_depth,max_leaf,n_outputs", + PARITY_CASES, ) def test_univariate_classification_discretizer_vs_sklearn_fidelity( criterion: str, - n_samples: int, num_classes: int, min_leaf_size: int, min_gain_split: float, @@ -63,6 +38,7 @@ def test_univariate_classification_discretizer_vs_sklearn_fidelity( ) -> None: """Predictions and leaf counts must match sklearn on synthetic univariate data.""" rng = np.random.default_rng(12345) + n_samples = 1000 x = rng.random((n_samples, 1), dtype=np.float64) # Match sklearn tree builder input path, which internally works with float32. x32 = x.astype(np.float32, copy=False) @@ -104,4 +80,33 @@ def test_univariate_classification_discretizer_vs_sklearn_fidelity( # good leaf. Allow a small mismatch fraction (matches the tolerance philosophy # of the regression discretizer test); exact fidelity holds only on signal data. mismatch_frac = np.mean(sklearn_preds != ud_preds) - assert mismatch_frac <= 0.01, f"prediction mismatch {mismatch_frac:.4%} exceeds 1% tolerance" + assert ( + mismatch_frac <= 0.01 + ), f"prediction mismatch {mismatch_frac:.4%} exceeds 1% tolerance" + + +def test_univariate_classification_gain_threshold_blocks_known_split() -> None: + x = np.arange(20, dtype=np.float32).reshape(-1, 1) + y = np.repeat(np.array([0, 1], dtype=np.uintp), 10) + features = np.array([0], dtype=np.uintp) + + split = UnivariateClassificationDiscretizer(criterion="gini") + split.Train(x, features, y, 2, 1, 0.0, 0, 0) + blocked = UnivariateClassificationDiscretizer(criterion="gini") + blocked.Train(x, features, y, 2, 1, 1.0, 0, 0) + + assert split.numLeaves == 2 + assert blocked.numLeaves == 1 + + +def test_univariate_classification_minimum_leaf_blocks_known_split() -> None: + x = np.arange(12, dtype=np.float32).reshape(-1, 1) + y = np.repeat(np.array([0, 1], dtype=np.uintp), 6) + features = np.array([0], dtype=np.uintp) + allowed = UnivariateClassificationDiscretizer(criterion="gini") + allowed.Train(x, features, y, 2, 6, 0.0, 0, 0) + blocked = UnivariateClassificationDiscretizer(criterion="gini") + blocked.Train(x, features, y, 2, 7, 0.0, 0, 0) + + assert allowed.numLeaves == 2 + assert blocked.numLeaves == 1 diff --git a/tests/test_univariate_regression_discretizer.py b/tests/test_univariate_regression_discretizer.py index 75e7f8d..3ce9297 100644 --- a/tests/test_univariate_regression_discretizer.py +++ b/tests/test_univariate_regression_discretizer.py @@ -8,22 +8,14 @@ import numpy as np import pytest -from itertools import product from Discretizers import UnivariateRegressionDiscretizer from sklearn.tree import DecisionTreeRegressor -from tests.discretizer_grid import ( - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - MIN_GAIN_VALUES, - MIN_LEAF_VALUES, - N_VALUES, - n_outputs_params, -) - -def regression_predict(ud: UnivariateRegressionDiscretizer, x: np.ndarray) -> np.ndarray: +def regression_predict( + ud: UnivariateRegressionDiscretizer, x: np.ndarray +) -> np.ndarray: """Predict by mapping transform() bin indices to bin predictions.""" bin_locs = ud.transform(x) bin_preds = ud.getBinPredictions() @@ -39,32 +31,7 @@ def sklearn_regression_criterion(user_criterion: str) -> str: return user_criterion -def _sklearn_supports_absolute_error() -> bool: - try: - DecisionTreeRegressor(criterion="absolute_error") - except (ValueError, TypeError): - return False - return True - - -def _skip_if_sklearn_mae_best_first_segfault(max_leaf: int, criterion: str) -> None: - """ - sklearn's best-first builder (max_leaf_nodes > 0) can segfault in native code - with criterion='absolute_error' on Python 3.14+ (see crash in tree/_tree.so). - Our discretizer still trains; skip only the sklearn reference for this grid cell. - """ - if max_leaf == 0 or criterion not in ("absolute_error", "mae"): - return - if sys.version_info >= (3, 14): - pytest.skip( - "sklearn reference: BestFirstTreeBuilder + absolute_error + max_leaf_nodes " - "segfaults on Python 3.14+; compare MAE with max_leaf=0 cases only" - ) - - -def _pred_discrepancy( - sk: np.ndarray, ud: np.ndarray, *, use_mae: bool -) -> float: +def _pred_discrepancy(sk: np.ndarray, ud: np.ndarray, *, use_mae: bool) -> float: """RMSE or mean absolute error between sklearn and native prediction vectors.""" d = sk.astype(np.float64) - ud.astype(np.float64) if use_mae: @@ -72,35 +39,20 @@ def _pred_discrepancy( return float(np.sqrt(np.mean(d**2))) -# Same (n_samples, leaf, gain, depth, max_leaf) sweep as classification — without num_classes. -GRID = list( - product( - N_VALUES, - MIN_LEAF_VALUES, - MIN_GAIN_VALUES, - MAX_DEPTH_VALUES, - MAX_LEAF_VALUES, - ) -) -IDS = [ - f"N={n}|leaf={leaf}|gain={gain}|depth={depth}|max_leaf={max_leaf}" - for n, leaf, gain, depth, max_leaf in GRID +PARITY_CASES = [ + pytest.param("squared_error", 1, 0.0, 0, 0, 1, id="mse-unconstrained"), + pytest.param("absolute_error", 1, 0.0, 0, 0, 2, id="mae-multioutput"), + pytest.param("squared_error", 1, 0.0, 4, 0, 2, id="depth-limited"), + pytest.param("squared_error", 1, 0.0, 0, 8, 1, id="leaf-limited"), ] -@pytest.mark.parametrize("n_outputs", n_outputs_params()) @pytest.mark.parametrize( - "criterion", - ["squared_error", "absolute_error"], # aliases mse/mae covered by equivalence test below -) -@pytest.mark.parametrize( - "n_samples,min_leaf_size,min_gain_split,max_depth,max_leaf", - GRID, - ids=IDS, + "criterion,min_leaf_size,min_gain_split,max_depth,max_leaf,n_outputs", + PARITY_CASES, ) def test_univariate_regression_discretizer_vs_sklearn_fidelity( criterion: str, - n_samples: int, min_leaf_size: int, min_gain_split: float, max_depth: int, @@ -108,12 +60,8 @@ def test_univariate_regression_discretizer_vs_sklearn_fidelity( n_outputs: int, ) -> None: """Bin predictions should track ``DecisionTreeRegressor`` within tolerance (MSE or MAE criterion).""" - if criterion in ("absolute_error", "mae") and not _sklearn_supports_absolute_error(): - pytest.skip("sklearn DecisionTreeRegressor does not support criterion='absolute_error'") - - _skip_if_sklearn_mae_best_first_segfault(max_leaf, criterion) - rng = np.random.default_rng(12345) + n_samples = 1000 x = rng.random((n_samples, 1), dtype=np.float64) x32 = x.astype(np.float32, copy=False) y_shape = (n_samples,) if n_outputs == 1 else (n_samples, n_outputs) @@ -158,7 +106,68 @@ def test_univariate_regression_discretizer_vs_sklearn_fidelity( assert disc < 0.9 * sigma -@pytest.mark.parametrize("alias,canonical", [("mse", "squared_error"), ("mae", "absolute_error")]) +@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) +def test_univariate_regression_leaf_limit_binds_without_reference( + criterion: str, +) -> None: + rng = np.random.default_rng(9) + x = rng.random((200, 1), dtype=np.float32) + y = rng.standard_normal(200).astype(np.float32) + ud = UnivariateRegressionDiscretizer(criterion=criterion) + ud.Train(x, np.array([0], dtype=np.uintp), y, 1, 0.0, 0, 4) + bins = ud.transform(x) + assert ud.numLeaves == 4 + assert np.all(bins < ud.numLeaves) + assert np.all(np.isfinite(regression_predict(ud, x))) + + +@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) +def test_univariate_regression_respects_minimum_leaf_size(criterion: str) -> None: + x = np.arange(12, dtype=np.float32).reshape(-1, 1) + y = np.repeat(np.array([0.0, 10.0], dtype=np.float32), 6) + features = np.array([0], dtype=np.uintp) + allowed = UnivariateRegressionDiscretizer(criterion=criterion) + allowed.Train(x, features, y, 6, 0.0, 0, 0) + blocked = UnivariateRegressionDiscretizer(criterion=criterion) + blocked.Train(x, features, y, 7, 0.0, 0, 0) + + assert allowed.numLeaves == 2 + assert blocked.numLeaves == 1 + + +def test_univariate_regression_gain_threshold_blocks_known_split() -> None: + x = np.arange(12, dtype=np.float32).reshape(-1, 1) + y = np.repeat(np.array([0.0, 10.0], dtype=np.float32), 6) + features = np.array([0], dtype=np.uintp) + split = UnivariateRegressionDiscretizer(criterion="squared_error") + split.Train(x, features, y, 1, 0.0, 0, 0) + blocked = UnivariateRegressionDiscretizer(criterion="squared_error") + blocked.Train(x, features, y, 1, 1_000.0, 0, 0) + + assert split.numLeaves == 2 + assert blocked.numLeaves == 1 + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), + reason="sklearn absolute_error best-first builder can crash on Python 3.14+", +) +def test_univariate_mae_leaf_limit_matches_sklearn_when_reference_is_safe() -> None: + rng = np.random.default_rng(11) + x = rng.random((200, 1), dtype=np.float32) + y = rng.standard_normal(200).astype(np.float32) + sk = DecisionTreeRegressor(criterion="absolute_error", max_leaf_nodes=4).fit(x, y) + ud = UnivariateRegressionDiscretizer(criterion="absolute_error") + ud.Train(x, np.array([0], dtype=np.uintp), y, 1, 0.0, 0, 4) + assert sk.get_n_leaves() == ud.numLeaves + assert _pred_discrepancy( + sk.predict(x), regression_predict(ud, x), use_mae=True + ) < np.std(y) + + +@pytest.mark.parametrize( + "alias,canonical", [("mse", "squared_error"), ("mae", "absolute_error")] +) def test_regression_criterion_aliases_equivalent(alias: str, canonical: str) -> None: """``mse``/``mae`` are aliases: identical bin predictions to their canonical name.""" rng = np.random.default_rng(12345) @@ -172,15 +181,3 @@ def test_regression_criterion_aliases_equivalent(alias: str, canonical: str) -> ud.Train(x32, features, y, 1, 0.0, 0, 0) preds[crit] = regression_predict(ud, x32) np.testing.assert_array_equal(preds[alias], preds[canonical]) - - -def test_regression_friedman_mse_constructible() -> None: - """Constructible when bindings accept friedman_mse (else skipped).""" - try: - UnivariateRegressionDiscretizer(criterion="friedman_mse") - except ValueError as e: - pytest.skip(f"UnivariateRegressionDiscretizer does not support friedman_mse: {e}") - - -def test_regression_mae_constructible() -> None: - UnivariateRegressionDiscretizer(criterion="mae") diff --git a/tests/test_weighted_sample.py b/tests/test_weighted_sample.py index a754bbd..174de6d 100644 --- a/tests/test_weighted_sample.py +++ b/tests/test_weighted_sample.py @@ -16,7 +16,10 @@ SGTClassifier, SGTRegressor, ) -from sgtlearn._weights import effective_sample_weight_classification +from sgtlearn._weights import ( + effective_sample_weight_classification, + normalize_sample_weight, +) from tests.constants import TEST_TAO_N_RUNS from tests.discretizer_grid import n_outputs_params @@ -24,13 +27,17 @@ pytest.importorskip("sklearn") -def _classification_predict(ud: UnivariateClassificationDiscretizer, x: np.ndarray) -> np.ndarray: +def _classification_predict( + ud: UnivariateClassificationDiscretizer, x: np.ndarray +) -> np.ndarray: bin_locs = ud.transform(x) bin_preds = ud.getBinPredictions() return np.asarray(bin_preds[bin_locs], dtype=np.uintp) -def _regression_predict(ud: UnivariateRegressionDiscretizer, x: np.ndarray) -> np.ndarray: +def _regression_predict( + ud: UnivariateRegressionDiscretizer, x: np.ndarray +) -> np.ndarray: bin_locs = ud.transform(x) bin_preds = ud.getBinPredictions() return np.asarray(bin_preds[bin_locs], dtype=np.float32) @@ -322,7 +329,9 @@ def test_random_sg_forest_classifier_applies_class_weight_once() -> None: np.testing.assert_array_equal(est.classes_, forest.classes_) -def test_random_sg_forest_classifier_class_weight_times_sample_weight_matches_sklearn() -> None: +def test_random_sg_forest_classifier_class_weight_times_sample_weight_matches_sklearn() -> ( + None +): from sklearn.datasets import load_breast_cancer X, y = load_breast_cancer(return_X_y=True) @@ -394,6 +403,8 @@ def test_effective_sample_weight_classification_helper() -> None: classes_, ) np.testing.assert_allclose(sw, [2.0, 10.0, 2.0]) + assert sw.dtype == np.float32 + assert sw.flags.c_contiguous def test_effective_sample_weight_multioutput_list_of_dicts() -> None: @@ -408,3 +419,58 @@ def test_effective_sample_weight_multioutput_list_of_dicts() -> None: ) # row0: 2*3=6, row1: 1*1=1, row2: 2*1=2 np.testing.assert_allclose(sw, [6.0, 1.0, 2.0]) + + +def test_normalize_sample_weight_allows_zero_with_positive_weight() -> None: + sw = normalize_sample_weight(np.array([0.0, 2.0]), 2) + assert sw is not None + np.testing.assert_array_equal(sw, [0.0, 2.0]) + assert sw.dtype == np.float32 + assert sw.flags.c_contiguous + + +def test_zero_weight_row_has_no_regressor_contribution() -> None: + X = np.array([[0.0], [1.0], [2.0], [100.0]]) + y = np.array([0.0, 0.0, 1.0, 100.0]) + params = {"max_depth": 2, "tao_n_runs": 0, "random_state": 0} + + weighted = SGTRegressor(**params).fit(X, y, sample_weight=[1.0, 1.0, 1.0, 0.0]) + omitted = SGTRegressor(**params).fit(X[:3], y[:3]) + + np.testing.assert_allclose(weighted.predict(X[:3]), omitted.predict(X[:3])) + + +@pytest.mark.parametrize( + ("sample_weight", "n_samples", "match"), + [ + ([1.0], 2, "shape"), + ([1.0, -1.0], 2, "non-negative"), + ([0.0, 0.0], 2, "positive"), + ], +) +def test_normalize_sample_weight_rejects_invalid_values( + sample_weight: list[float], n_samples: int, match: str +) -> None: + with pytest.raises(ValueError, match=match): + normalize_sample_weight(np.asarray(sample_weight), n_samples) + + +@pytest.mark.parametrize( + ("class_weight", "classes", "match", "exc_type"), + [ + ({2: 1.0}, np.array([0, 1]), "training classes", ValueError), + ({1: -1.0}, np.array([0, 1]), "non-negative", ValueError), + (1.0, np.array([0, 1]), "mapping", TypeError), + ([{0: 1.0}], [np.array([0, 1]), np.array([0, 1])], "one mapping", ValueError), + ({0: 1.0}, [np.array([0, 1])], "one class array", ValueError), + ], +) +def test_effective_sample_weight_rejects_invalid_class_configuration( + class_weight: object, + classes: object, + match: str, + exc_type: type[Exception], +) -> None: + y = np.array([[0, 1], [1, 0]]) if isinstance(classes, list) else np.array([0, 1]) + with pytest.raises(exc_type, match=match): + effective_sample_weight_classification(None, y, class_weight, classes) # type: ignore[arg-type]