test: replace broad matrices with focused contracts

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<size_t> forward = {0, 1};
+  const std::vector<size_t> reverse = {1, 0};
+  REQUIRE((obj.assignments == forward || obj.assignments == reverse));
+  REQUIRE(obj.partitionSampleCounts() ==
+          std::vector<size_t>{samplesPerBin, samplesPerBin});
 }
 
 } // namespace
@@ -45,7 +51,7 @@ TEST_CASE("EntropyBranchAssignment coordinate descent",
   std::vector<size_t> 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<size_t> 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<size_t> 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<size_t> 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<size_t> 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",

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():

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:

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: