From aba00969df210d0129b67f42bb3a552f851eea0a Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 2 Oct 2026 22:14:47 -0400 Subject: [PATCH 1/3] Keep TAO from lowering training accuracy at lambda=0 With 3+ children, a sample can be correct in several children, so the routing problem at a node is multi-label. TAO split credit 1/numCorrect across those children, so its objective stopped tracking accuracy and it could accept rules (including the dummy) that lowered training accuracy. - Care rewards give full credit (1) to every correct child; router training rows split the sample's weight across them instead. - Each discretizer bin is relabelled to the child with the largest summed care reward, not the argmax of the fitted histogram. Empty bins and an unused NaN bin get dummyChild. - Current, dummy and candidate rules are all scored on the same full-credit rewards, so at lambda=0 an accepted change never lowers the care-set score. - Two-child nodes also try the exact best threshold on each feature, materialized as a depth-1 numeric discretizer checked against the scan. Closes #80 Co-Authored-By: Claude Opus 5.5 --- cpp/CMakeLists.txt | 1 + .../TAO/ClassificationTaoAdapter.cpp | 23 +-- .../algorithms/TAO/ClassificationTaoAdapter.h | 14 +- cpp/src/algorithms/TAO/TaoObjective.cpp | 126 +++++++++++--- cpp/src/algorithms/TAO/TaoObjective.h | 52 +++++- .../TAO/TreeAlternatingOptimization.cpp | 79 +++++++++ .../TAO/TreeAlternatingOptimization.h | 5 +- cpp/tests/test_tao_objective.cpp | 82 +++++++++ tests/test_tao.py | 158 ++++++++++++++++++ 9 files changed, 489 insertions(+), 51 deletions(-) create mode 100644 cpp/tests/test_tao_objective.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 347e58d..9804874 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -240,6 +240,7 @@ if (SGTLEARN_BUILD_TESTS) tests/test_splitters.cpp tests/test_branch_assignment.cpp tests/test_shape_assignment_search.cpp + tests/test_tao_objective.cpp tests/bench_mae_branch_assignment.cpp ) diff --git a/cpp/src/algorithms/TAO/ClassificationTaoAdapter.cpp b/cpp/src/algorithms/TAO/ClassificationTaoAdapter.cpp index 2384df6..82ed94d 100644 --- a/cpp/src/algorithms/TAO/ClassificationTaoAdapter.cpp +++ b/cpp/src/algorithms/TAO/ClassificationTaoAdapter.cpp @@ -91,23 +91,13 @@ void ClassificationTaoAdapter::childRewards( return; } - // Single-output: 1/x split uniformly across correct children. + // Single-output: full credit (1) for every correct child. Routing is + // multi-label, so a sample counts once wherever it lands correctly. const size_t label = y_(0, col); - size_t numCorrect = 0; for (size_t c = 0; c < childLeaves.size(); ++c) { const auto predictions = argMaxClass(classificationTree_.classCounts[childLeaves[c]]); - const bool correct = predictions[0] == label; - if (correct) - ++numCorrect; - reward[c] = correct ? 1.0 : 0.0; - } - if (numCorrect == 0) - return; - const double correctReward = 1.0 / static_cast(numCorrect); - for (size_t c = 0; c < reward.size(); ++c) { - if (reward[c] > 0.0) - reward[c] = correctReward; + reward[c] = predictions[0] == label ? 1.0 : 0.0; } } @@ -160,10 +150,15 @@ NodeCareSet ClassificationTaoAdapter::buildCareSet( size_t pos = 0; for (size_t i = 0; i < care.size(); ++i) { const float wi = w_(care.careCols[i]); + // Single-output: split the sample's weight across its correct children + // (soft labels), so it adds total mass wi to the router fit. childCounts + // keeps full weight: dummyChild is correct for the most care weight. + const float wRow = + nOutputs_ == 1 ? wi / static_cast(goodChildren[i].size()) : wi; for (size_t child : goodChildren[i]) { care.Xexp.col(pos) = X_.col(care.careCols[i]); care.yexp(pos) = child; - care.wexp(pos) = wi; + care.wexp(pos) = wRow; childCounts[child] += static_cast(wi); ++pos; } diff --git a/cpp/src/algorithms/TAO/ClassificationTaoAdapter.h b/cpp/src/algorithms/TAO/ClassificationTaoAdapter.h index 9443fb5..d53d60e 100644 --- a/cpp/src/algorithms/TAO/ClassificationTaoAdapter.h +++ b/cpp/src/algorithms/TAO/ClassificationTaoAdapter.h @@ -7,11 +7,14 @@ * Care-set logic (classification variant): * * - For each sample at an internal node, walk each child subtree to its leaf - * and score child ``c`` with reward ``1/x`` if the leaf's majority class equals - * the sample label, where ``x`` is the number of correct children; else ``0``. + * and score child ``c`` with reward ``1`` if the leaf's majority class equals + * the sample label, else ``0``. Routing to any correct child is fully + * correct (multi-label), so rewards are the true accuracy contribution. * - Samples where all children tie are excluded from the care set. * - Samples with multiple equally good children contribute one expanded router- - * training row per good child (multi-to-single pseudolabel expansion). + * training row per good child (multi-to-single pseudolabel expansion). In + * single-output, each row carries ``w / x`` (``x`` = number of correct + * children), so the sample adds total weight ``w`` to the router fit. * * Leaf refresh recomputes weighted class histograms at every node from * ``nodeSamples``. The routing discretizer uses the tree's own criterion @@ -59,9 +62,8 @@ class ClassificationTaoAdapter final : public ShapeGeneralizedTaoAdapter { /** * Per-child correctness rewards for one sample. * - * Single-output: ``1/x`` split uniformly across correct children (``x`` = - * count of correct), else ``0``. Multi-output: fraction of outputs correctly - * classified by each child's leaf. + * Single-output: ``1`` for each correct child, else ``0``. Multi-output: + * fraction of outputs correctly classified by each child's leaf. */ void childRewards(const std::vector &childLeaves, arma::uword col, std::vector &reward) const; diff --git a/cpp/src/algorithms/TAO/TaoObjective.cpp b/cpp/src/algorithms/TAO/TaoObjective.cpp index 550b48c..7ca5b2a 100644 --- a/cpp/src/algorithms/TAO/TaoObjective.cpp +++ b/cpp/src/algorithms/TAO/TaoObjective.cpp @@ -5,11 +5,13 @@ #include #include "algorithms/TAO/TaoObjective.h" +#include "algorithms/missing_values.h" + #include #include #include #include -#include +#include namespace tao { @@ -51,17 +53,15 @@ double TaoObjective::penalizedScore(double rewardSum, complexityScale * lambda_ * nodeSampleCount_ / totalCareWeight_; } -double TaoObjective::rewardSumForDiscretizer( - const ClassificationDiscretizer &disc, +double TaoObjective::rewardSumForBins( + const arma::Row &bins, const std::vector &binToPartition) const { - arma::Row bins; - disc.transform(X_, bins); double rewardSum = 0.0; for (size_t i = 0; i < nCare_; ++i) { const size_t bin = bins(care_.careCols[i]); if (bin >= binToPartition.size()) throw std::runtime_error( - "TaoObjective::rewardSumForDiscretizer: bin out of range"); + "TaoObjective::rewardSumForBins: bin out of range"); const size_t child = binToPartition[bin]; rewardSum += careWeight(i) * care_.careRewards[i][child]; } @@ -93,31 +93,109 @@ double TaoObjective::scoreDiscretizer( if (disc.numLeaves() < 1) return -std::numeric_limits::infinity(); - const std::vector>> &leafStats = - disc.leafStats(); - if (leafStats.empty()) + // Bin count includes a univariate discretizer's trailing NaN bin. + const size_t nBins = disc.leafStats().size(); + if (nBins == 0) throw std::runtime_error( "TaoObjective::scoreDiscretizer: discretizer has no leaf stats"); - binToPartitionOut.resize(leafStats.size()); + arma::Row bins; + disc.transform(X_, bins); - // TAO trains the inner discretizer on scalar care-set pseudo-labels (child - // partition indices 0..k-1), so each bin's stats are [1][k] not multi-output. - for (size_t b = 0; b < leafStats.size(); ++b) { - if (leafStats[b].size() != 1) - throw std::runtime_error( - "TaoObjective::scoreDiscretizer: bin " + std::to_string(b) + - " expected 1 output row of class counts, got " + - std::to_string(leafStats[b].size())); - if (leafStats[b][0].empty()) + // Relabel each bin by the true care reward, not the fitted histogram: the + // router is fit on split weights, whose argmax can miss the best child. + const size_t k = nCare_ == 0 ? 0 : care_.careRewards.front().size(); + std::vector> binReward(nBins, + std::vector(k, 0.0)); + std::vector binWeight(nBins, 0.0); + for (size_t i = 0; i < nCare_; ++i) { + const size_t bin = bins(care_.careCols[i]); + if (bin >= nBins) throw std::runtime_error( - "TaoObjective::scoreDiscretizer: bin " + std::to_string(b) + - " class histogram is empty"); - binToPartitionOut[b] = argMax(leafStats[b][0]); + "TaoObjective::scoreDiscretizer: bin out of range"); + const double w = careWeight(i); + binWeight[bin] += w; + for (size_t c = 0; c < k; ++c) + binReward[bin][c] += w * care_.careRewards[i][c]; } + binToPartitionOut.resize(nBins); + for (size_t b = 0; b < nBins; ++b) + binToPartitionOut[b] = + binWeight[b] > 0.0 ? argMax(binReward[b]) : care_.dummyChild; - return penalizedScore( - rewardSumForDiscretizer(disc, binToPartitionOut), complexityScale); + return penalizedScore(rewardSumForBins(bins, binToPartitionOut), + complexityScale); +} + +double TaoObjective::scoreBinAssignment( + const arma::Row &bins, const std::vector &binToPartition, + double complexityScale) const { + return penalizedScore(rewardSumForBins(bins, binToPartition), + complexityScale); +} + +std::optional +TaoObjective::bestThresholdCut(size_t feature, size_t minLeafSize) const { + if (nCare_ == 0) + return std::nullopt; + if (care_.careRewards.front().size() != 2) + throw std::invalid_argument( + "TaoObjective::bestThresholdCut: node must have exactly two children"); + + std::vector> finite; // (value, care index) + finite.reserve(nCare_); + double total0 = 0.0, total1 = 0.0; + double nan0 = 0.0, nan1 = 0.0, nanWeight = 0.0; + for (size_t i = 0; i < nCare_; ++i) { + const float value = + X_(static_cast(feature), care_.careCols[i]); + const double w = careWeight(i); + const double r0 = w * care_.careRewards[i][0]; + const double r1 = w * care_.careRewards[i][1]; + if (missing_values::is_finite(value)) { + finite.emplace_back(value, i); + total0 += r0; + total1 += r1; + } else { + nan0 += r0; + nan1 += r1; + nanWeight += w; + } + } + + const size_t minSide = std::max(minLeafSize, 1); + const size_t n = finite.size(); + if (n < 2 * minSide) + return std::nullopt; + std::sort(finite.begin(), finite.end()); + + const size_t nanChild = + nanWeight > 0.0 ? (nan1 > nan0 ? 1 : 0) : care_.dummyChild; + const double nanReward = nanChild == 0 ? nan0 : nan1; + + std::optional best; + double left0 = 0.0, left1 = 0.0; + for (size_t j = 1; j < n; ++j) { + const auto [prevValue, prevIdx] = finite[j - 1]; + const double w = careWeight(prevIdx); + left0 += w * care_.careRewards[prevIdx][0]; + left1 += w * care_.careRewards[prevIdx][1]; + if (j < minSide) + continue; + if (n - j < minSide) + break; + // Same tie rule as Splitter::findBestSplit, so a depth-1 discretizer + // trained on the two sides reproduces this cut. + if (finite[j].first <= prevValue + static_cast(1e-7)) + continue; + const double leftTo0 = left0 + (total1 - left1); + const double leftTo1 = left1 + (total0 - left0); + const size_t leftChild = leftTo1 > leftTo0 ? 1 : 0; + const double reward = std::max(leftTo0, leftTo1) + nanReward; + if (!best || reward > best->rewardSum) + best = ThresholdCut{prevValue, leftChild, nanChild, reward}; + } + return best; } } // namespace tao diff --git a/cpp/src/algorithms/TAO/TaoObjective.h b/cpp/src/algorithms/TAO/TaoObjective.h index 1011009..a1a94bb 100644 --- a/cpp/src/algorithms/TAO/TaoObjective.h +++ b/cpp/src/algorithms/TAO/TaoObjective.h @@ -15,10 +15,26 @@ #include "algorithms/TAO/TaoAdapter.h" #include +#include +#include #include namespace tao { +/** + * One univariate threshold routing for a two-child node: finite + * ``x <= leftMax`` goes to ``leftChild``, finite ``x > leftMax`` to + * ``1 - leftChild``, non-finite to ``nanChild``. + */ +struct ThresholdCut { + /** Largest finite care value on the left side; the cut lies above it. */ + float leftMax = 0.0f; + size_t leftChild = 0; + size_t nanChild = 0; + /** Weighted care reward of this routing (no complexity penalty). */ + double rewardSum = -std::numeric_limits::infinity(); +}; + /** * Mean care-set reward under a routing rule, with optional split penalty. * @@ -48,8 +64,11 @@ class TaoObjective { /** * Extract routing from a trained discretizer and score it on the care set. * - * Bins are mapped to child partitions by argmax over discretizer leaf stats. - * Writes the induced bin-to-partition map to the out-param. + * Each bin goes to the child with the largest summed care reward over the + * care samples in it (ties: lowest child index), not to the argmax of the + * discretizer's fitted histogram. Bins with zero care weight, including an + * unused NaN bin, go to ``dummyChild``. Writes the bin-to-partition map to + * the out-param. * * @returns Penalized mean reward, or ``-infinity`` when ``disc`` has no bins. */ @@ -57,12 +76,35 @@ class TaoObjective { std::vector &binToPartitionOut, double complexityScale) const; + /** + * Penalized mean reward when care sample ``i`` routes to + * ``binToPartition[bins(careCols[i])]``. + * + * @param bins Bin of every column of ``X`` (as from ``transform(X, bins)``). + */ + double scoreBinAssignment(const arma::Row &bins, + const std::vector &binToPartition, + double complexityScale) const; + + /** + * Exact best threshold routing on raw ``feature`` for a two-child node. + * + * Cuts lie between consecutive distinct finite care values (using the inner + * splitter's ``1e-7`` tie rule) with at least ``minLeafSize`` finite care + * samples per side; both orientations are tried. Non-finite care samples go + * to the child with the larger summed reward over them (``dummyChild`` when + * they carry no weight). + * + * @returns The best cut, or ``std::nullopt`` when no cut is feasible. + */ + std::optional bestThresholdCut(size_t feature, + size_t minLeafSize) const; + private: static size_t argMax(const std::vector &counts); - double rewardSumForDiscretizer( - const ClassificationDiscretizer &disc, - const std::vector &binToPartition) const; + double rewardSumForBins(const arma::Row &bins, + const std::vector &binToPartition) const; double meanReward(double rewardSum) const; double penalizedScore(double rewardSum, double complexityScale) const; diff --git a/cpp/src/algorithms/TAO/TreeAlternatingOptimization.cpp b/cpp/src/algorithms/TAO/TreeAlternatingOptimization.cpp index d742725..1f23251 100644 --- a/cpp/src/algorithms/TAO/TreeAlternatingOptimization.cpp +++ b/cpp/src/algorithms/TAO/TreeAlternatingOptimization.cpp @@ -12,10 +12,13 @@ #include "Discretizers/factories/DiscretizerFactories.h" #include "Discretizers/InnerDiscretizerBase.h" #include "algorithms/TAO/TaoObjective.h" +#include "algorithms/missing_values.h" #include +#include #include #include +#include #include #include @@ -54,6 +57,51 @@ computeNodeSamples(const std::vector &nodes, return nodeSamples; } +/** + * Numeric router for ``cut`` on raw ``feature``: finite ``x <= cut.leftMax`` + * -> bin 0, larger -> bin 1, non-finite -> NaN bin 2. Trained as a depth-1 + * discretizer on the finite care samples labelled by side; that split is pure, + * so the splitter recovers it with its own threshold convention. Writes the + * bin of every column of ``X`` to ``bins``. + */ +std::unique_ptr +materializeThresholdCut(LearningCriterion criterion, const NodeCareSet &care, + const arma::fmat &X, size_t feature, + const ThresholdCut &cut, arma::Row &bins) { + const auto side = [&](arma::uword col) -> size_t { + const float value = X(static_cast(feature), col); + if (!missing_values::is_finite(value)) + return 2; + return value > cut.leftMax ? 1 : 0; + }; + std::vector cols; + std::vector labels; + for (arma::uword col : care.careCols) { + const size_t s = side(col); + if (s == 2) + continue; + cols.push_back(col); + labels.push_back(s); + } + + arma::uvec featOne = {static_cast(feature)}; + auto disc = + makeClassificationDiscretizer(criterion, DiscretizerInputKind::Numeric); + disc->Train(X.cols(arma::uvec(cols)), featOne, arma::Row(labels), 2, + /*minLeafSize=*/1, /*minGainSplit=*/0.0, /*maxDepth=*/1, + /*maxLeafNodes=*/2); + disc->transform(X, bins); + + bool reproduced = disc->numLeaves() == 2; + for (size_t i = 0; reproduced && i < care.size(); ++i) + reproduced = bins(care.careCols[i]) == side(care.careCols[i]); + if (!reproduced) + throw std::runtime_error( + "tao: threshold router does not reproduce the exact cut on feature " + + std::to_string(feature)); + return disc; +} + } // namespace bool optimizeNodeInPlace( @@ -121,6 +169,37 @@ bool optimizeNodeInPlace( } } + // Two children: also try the exact best threshold on each feature, which + // the inner-tree fit above only approximates. + if (k == 2) { + std::optional bestCut; + size_t cutFeature = 0; + for (size_t f = 0; f < numFeatures; ++f) { + const auto cut = objective.bestThresholdCut(f, innerParams.minLeafSize); + if (cut && (!bestCut || cut->rewardSum > bestCut->rewardSum)) { + bestCut = cut; + cutFeature = f; + } + } + if (bestCut) { + arma::Row bins; + auto disc = materializeThresholdCut(routerCriterion, care, X, cutFeature, + *bestCut, bins); + std::vector binToPartition = { + bestCut->leftChild, 1 - bestCut->leftChild, bestCut->nanChild}; + const double score = objective.scoreBinAssignment( + bins, binToPartition, /*complexityScale=*/1.0); + if (score > bestSingleScore) { + bestSingleScore = score; + bestFeature = cutFeature; + bestBinToPartition = std::move(binToPartition); + bestDiscretizer = + std::shared_ptr(std::move(disc)); + haveSingle = true; + } + } + } + double bestPairScore = -std::numeric_limits::infinity(); bool havePair = false; std::array bestPairLogical{}; diff --git a/cpp/src/algorithms/TAO/TreeAlternatingOptimization.h b/cpp/src/algorithms/TAO/TreeAlternatingOptimization.h index b5844c8..4462289 100644 --- a/cpp/src/algorithms/TAO/TreeAlternatingOptimization.h +++ b/cpp/src/algorithms/TAO/TreeAlternatingOptimization.h @@ -14,8 +14,9 @@ * 2. Score the current rule and a constant (dummy) rule that sends all care * samples to ``dummyChild``. * 3. For each candidate feature or feature pair, train a classification - * discretizer over child-partition pseudolabels and score the induced routing - * rule. + * discretizer over child-partition pseudolabels, send each bin to the child + * with the largest care reward in it, and score the induced routing rule. + * Two-child nodes also try the exact best threshold on each feature. * 4. Accept the best non-worsening rule (current, dummy, single-feature, or pair * split). In weighted reward units, ``lambda`` penalizes single-feature splits * by ``lambda * nodeSampleCount`` and pair splits by diff --git a/cpp/tests/test_tao_objective.cpp b/cpp/tests/test_tao_objective.cpp new file mode 100644 index 0000000..b93bc3e --- /dev/null +++ b/cpp/tests/test_tao_objective.cpp @@ -0,0 +1,82 @@ +/** + * @file test_tao_objective.cpp + * @brief Catch2 tests for TAO per-node scoring (bin relabel, threshold scan). + */ + +#include +#include +#include + +#include +#include + +#include +#include + +using Catch::Matchers::WithinAbs; + +TEST_CASE("TaoObjective relabels a bin by full-credit care reward", "[tao]") { + // One bin: 10 care samples are correct in children {0, 1}, 8 only in child + // 2. Split router weights give the histogram (5, 5, 8), whose argmax is + // child 2 (8 correct); full credit gives (10, 10, 8) -> child 0 (10 correct). + tao::NodeCareSet care; + std::vector xCols; + std::vector yRows; + std::vector wRows; + for (arma::uword i = 0; i < 18; ++i) { + const bool shared = i < 10; + care.careCols.push_back(i); + care.careRewards.push_back(shared ? std::vector{1, 1, 0} + : std::vector{0, 0, 1}); + care.careWeights.push_back(1.0); + for (size_t child : shared ? std::vector{0, 1} + : std::vector{2}) { + xCols.push_back(i); + yRows.push_back(child); + wRows.push_back(shared ? 0.5f : 1.0f); + } + } + // Distinct from the bin's best child, so the unused NaN bin's map shows. + care.dummyChild = 1; + const arma::fmat X(1, 18, arma::fill::zeros); // constant feature: one bin + care.Xexp = X.cols(arma::uvec(xCols)); + care.yexp = arma::Row(yRows); + care.wexp = arma::Row(wRows); + + auto disc = makeClassificationDiscretizer(LearningCriterion::Gini, + DiscretizerInputKind::Numeric); + arma::uvec feature = {0}; + disc->Train(care.Xexp, feature, care.yexp, 3, 1, 0.0, 0, 0, care.wexp); + REQUIRE(disc->numLeaves() == 1); + REQUIRE(disc->leafStats()[0][0] == std::vector{5, 5, 8}); + + const tao::TaoObjective objective(care, X, /*lambda=*/0.0, + /*nodeSampleCount=*/18); + std::vector binToPartition; + const double score = objective.scoreDiscretizer(*disc, binToPartition, 1.0); + + REQUIRE(binToPartition == std::vector{0, 1}); + REQUIRE_THAT(score, WithinAbs(10.0 / 18.0, 1e-12)); +} + +TEST_CASE("TaoObjective finds the exact best threshold for two children", + "[tao]") { + const float nan = std::numeric_limits::quiet_NaN(); + const arma::fmat X = {{4, 1, 3, 2, nan}}; + // x = 1, 2 are correct in child 1; x = 3, 4 and NaN in child 0. + tao::NodeCareSet care; + care.careCols = {0, 1, 2, 3, 4}; + care.careRewards = {{1, 0}, {0, 1}, {1, 0}, {0, 1}, {1, 0}}; + care.dummyChild = 1; + const tao::TaoObjective objective(care, X, 0.0, 5); + + const auto cut = objective.bestThresholdCut(0, 1); + REQUIRE(cut.has_value()); + REQUIRE(cut->leftMax == 2.0f); + REQUIRE(cut->leftChild == 1); + REQUIRE(cut->nanChild == 0); + REQUIRE_THAT(cut->rewardSum, WithinAbs(5.0, 1e-12)); + + // Three finite care samples per side are needed, and only four exist. + REQUIRE_FALSE(objective.bestThresholdCut(0, 3).has_value()); +} diff --git a/tests/test_tao.py b/tests/test_tao.py index db9c195..5515f56 100644 --- a/tests/test_tao.py +++ b/tests/test_tao.py @@ -438,3 +438,161 @@ def spy(est, X, y, sw, **kw): SGTClassifier(class_weight={0: 1.0, 1: 3.0}, tao_n_runs=1, max_depth=2).fit(X, y) sw = seen[0] assert sw[y == 1][0] / sw[y == 0][0] == pytest.approx(3.0) + + +def _square_in_cross_data( + seed: int = 0, n_samples: int = 2400, n_features: int = 2, noise: float = 0.0 +) -> tuple[np.ndarray, np.ndarray]: + """Square-in-cross labels: several children of one 3-way split share a class.""" + rng = np.random.default_rng(seed) + X = rng.uniform(-2.5, 2.5, size=(n_samples, n_features)) + ix = np.abs(X[:, 0]) < 0.85 + iy = np.abs(X[:, 1]) < 0.85 + y = np.where(ix & iy, 2, np.where(ix | iy, 1, 0)) + flip = rng.random(n_samples) < noise + y[flip] = rng.integers(0, 3, flip.sum()) + return X, y + + +def test_tao_lambda0_keeps_accuracy_when_children_share_a_class() -> None: + X, y = _square_in_cross_data() + clf = SGTClassifier( + max_depth=1, + inner_max_depth=4, + inner_max_leaf_nodes=9, + num_partitions=3, + pairwise_candidates=0, + tao_n_runs=0, + random_state=0, + ).fit(X, y) + before = clf.score(X, y) + + tao.TAO_refine(clf, X, y, n_runs=10, lambda_=0.0) + + assert clf.score(X, y) >= before + + +@pytest.mark.parametrize("tao_n_runs", [0, None], ids=["no_fit_tao", "fit_tao"]) +@pytest.mark.parametrize("weighting", ["none", "sample_weight", "class_weight"]) +@pytest.mark.parametrize("num_partitions", [2, 3, 4]) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_tao_lambda0_never_lowers_weighted_training_accuracy( + seed: int, num_partitions: int, weighting: str, tao_n_runs: Optional[int] +) -> None: + X, y = _square_in_cross_data(seed, n_samples=600, n_features=3, noise=0.1) + sample_weight = None + class_weight = None + effective = np.ones(len(y)) + if weighting == "sample_weight": + sample_weight = np.random.default_rng(seed + 100).uniform(0.2, 3.0, len(y)) + effective = sample_weight + elif weighting == "class_weight": + # Same weights as sklearn's "balanced" (not accepted by SGTClassifier). + counts = np.bincount(y) + class_weight = {c: len(y) / (len(counts) * counts[c]) for c in range(len(counts))} + effective = np.array([class_weight[c] for c in y]) + # TAO sees float32 weights. + effective = effective.astype(np.float32).astype(np.float64) + + params: dict[str, Any] = dict( + max_depth=1, + num_partitions=num_partitions, + inner_max_depth=4, + inner_max_leaf_nodes=9, + class_weight=class_weight, + random_state=seed, + ) + if tao_n_runs is not None: + params["tao_n_runs"] = tao_n_runs + clf = SGTClassifier(**params).fit(X, y, sample_weight=sample_weight) + before = np.average(clf.predict(X) == y, weights=effective) + + tao.TAO_refine(clf, X, y, sample_weight=sample_weight, lambda_=0.0) + + after = np.average(clf.predict(X) == y, weights=effective) + assert after >= before - 1e-12 + + +def _k2_threshold_data(seed: int) -> tuple[np.ndarray, np.ndarray]: + """On x0, Gini's best cut isolates a pure block (600/800 correct); cutting + after the mixed middle block is more accurate (605/800). x1 is noise.""" + rng = np.random.default_rng(seed) + blocks = [(0.0, 200, 0), (1.0, 105, 0), (1.0, 100, 1), (2.0, 95, 0), (2.0, 300, 1)] + x0 = np.concatenate([lo + rng.integers(0, 10, n) / 10 for lo, n, _ in blocks]) + y = np.concatenate([np.full(n, label) for _, n, label in blocks]) + x1 = rng.integers(0, 100, len(y)) / 10 + return np.column_stack([x0, x1]), y + + +def _best_single_threshold_accuracy( + X: np.ndarray, y: np.ndarray, labels: tuple[int, int] +) -> float: + """Best accuracy of routing by one threshold to two leaves with fixed labels.""" + hits = np.column_stack([y == labels[0], y == labels[1]]).astype(float) + total = hits.sum(axis=0) + best = total.max() # constant routing + for f in range(X.shape[1]): + order = np.argsort(X[:, f], kind="stable") + xs = X[order, f] + left = np.cumsum(hits[order], axis=0) + cuts = np.flatnonzero(xs[1:] > xs[:-1]) # left side = order[: cut + 1] + keep = left[cuts, 0] + total[1] - left[cuts, 1] + swap = left[cuts, 1] + total[0] - left[cuts, 0] + best = max(best, keep.max(initial=0.0), swap.max(initial=0.0)) + return best / len(y) + + +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_tao_two_children_reach_the_best_single_threshold(seed: int) -> None: + X, y = _k2_threshold_data(seed) + clf = SGTClassifier( + max_depth=1, + num_partitions=2, + inner_max_depth=1, + inner_max_leaf_nodes=2, + pairwise_candidates=0, + tao_n_runs=0, + random_state=seed, + ).fit(X, y) + export = clf.tree_export() + root = export["nodes"][export["root_index"]] + assert len(root["children"]) == 2 + labels = tuple( + int(np.argmax(export["nodes"][cid]["class_counts"][0])) + for cid in root["children"] + ) + target = _best_single_threshold_accuracy(X, y, labels) + + tao.TAO_refine(clf, X, y, lambda_=0.0) + + assert clf.score(X, y) >= target - 1e-12 + + +@pytest.mark.parametrize("criterion", ["squared_error", "absolute_error"]) +@pytest.mark.parametrize("num_partitions", [2, 3]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_tao_lambda0_never_raises_regression_training_loss( + seed: int, num_partitions: int, criterion: str +) -> None: + X, y = make_regression( + n_samples=300, n_features=4, n_informative=3, noise=15.0, random_state=seed + ) + reg = SGTRegressor( + criterion=criterion, + max_depth=2, + num_partitions=num_partitions, + inner_max_depth=3, + inner_max_leaf_nodes=8, + tao_n_runs=0, + random_state=seed, + ).fit(X, y) + if criterion == "squared_error": + loss = lambda: np.mean((reg.predict(X) - y) ** 2) # noqa: E731 + else: + loss = lambda: np.mean(np.abs(reg.predict(X) - y)) # noqa: E731 + before = loss() + + tao.TAO_refine(reg, X, y, lambda_=0.0) + + # Native targets are float32, so allow float32-level slack. + assert loss() <= before * (1 + 1e-6) From 5cabd62854077eece71e6840cabb7c679b20b007 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 2 Oct 2026 22:32:35 -0400 Subject: [PATCH 2/3] Clarify ThresholdCut routing and the K=2 test oracle's assumptions Co-Authored-By: Claude Opus 5.5 --- cpp/src/algorithms/TAO/TaoObjective.h | 7 ++++--- tests/test_tao.py | 7 ++++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cpp/src/algorithms/TAO/TaoObjective.h b/cpp/src/algorithms/TAO/TaoObjective.h index a1a94bb..e755b9e 100644 --- a/cpp/src/algorithms/TAO/TaoObjective.h +++ b/cpp/src/algorithms/TAO/TaoObjective.h @@ -22,9 +22,10 @@ namespace tao { /** - * One univariate threshold routing for a two-child node: finite - * ``x <= leftMax`` goes to ``leftChild``, finite ``x > leftMax`` to - * ``1 - leftChild``, non-finite to ``nanChild``. + * One univariate threshold routing for a two-child node: finite care values + * ``x <= leftMax`` go to ``leftChild``, larger ones to ``1 - leftChild``, and + * non-finite ones to ``nanChild``. The installed router's threshold lies + * between ``leftMax`` and the next finite care value. */ struct ThresholdCut { /** Largest finite care value on the left side; the cut lies above it. */ diff --git a/tests/test_tao.py b/tests/test_tao.py index 5515f56..4939892 100644 --- a/tests/test_tao.py +++ b/tests/test_tao.py @@ -527,7 +527,12 @@ def _k2_threshold_data(seed: int) -> tuple[np.ndarray, np.ndarray]: def _best_single_threshold_accuracy( X: np.ndarray, y: np.ndarray, labels: tuple[int, int] ) -> float: - """Best accuracy of routing by one threshold to two leaves with fixed labels.""" + """Best accuracy of routing by one threshold to two leaves with fixed labels. + + Assumes finite values whose distinct values stay more than 1e-7 apart in + float32, as in ``_k2_threshold_data``; it does not model NaN routing or the + splitter's tie rule. + """ hits = np.column_stack([y == labels[0], y == labels[1]]).astype(float) total = hits.sum(axis=0) best = total.max() # constant routing From 5dbfde4fc463906e5fc644df8a38844e3b4c1787 Mon Sep 17 00:00:00 2001 From: Nakul Upadhya Date: Fri, 2 Oct 2026 22:32:35 -0400 Subject: [PATCH 3/3] Re-execute regression tutorial after the TAO fix Regression TAO now relabels bins by true loss and tries the exact K=2 threshold, so the printed R^2 values move slightly. The other tutorials' printed outputs and plots are unchanged by the fix. Co-Authored-By: Claude Opus 5.5 --- docs/tutorials/regression.ipynb | 36 ++++++++++++++++----------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/tutorials/regression.ipynb b/docs/tutorials/regression.ipynb index b5fe661..d76c1ee 100644 --- a/docs/tutorials/regression.ipynb +++ b/docs/tutorials/regression.ipynb @@ -16,10 +16,10 @@ "id": "f3faf212", "metadata": { "execution": { - "iopub.execute_input": "2026-09-11T17:30:50.317527Z", - "iopub.status.busy": "2026-09-11T17:30:50.317405Z", - "iopub.status.idle": "2026-09-11T17:30:51.028231Z", - "shell.execute_reply": "2026-09-11T17:30:51.027856Z" + "iopub.execute_input": "2026-10-03T02:30:25.203497Z", + "iopub.status.busy": "2026-10-03T02:30:25.203334Z", + "iopub.status.idle": "2026-10-03T02:30:27.093231Z", + "shell.execute_reply": "2026-10-03T02:30:27.092889Z" } }, "outputs": [ @@ -27,8 +27,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "predictions: [18.16 6.38 12.12 9.77 14.17]\n", - "R^2 train=0.765 test=0.591\n" + "predictions: [18.16 6.38 12.12 9.78 14.25]\n", + "R^2 train=0.767 test=0.594\n" ] } ], @@ -61,10 +61,10 @@ "id": "858de2af", "metadata": { "execution": { - "iopub.execute_input": "2026-09-11T17:30:51.029436Z", - "iopub.status.busy": "2026-09-11T17:30:51.029311Z", - "iopub.status.idle": "2026-09-11T17:30:51.137858Z", - "shell.execute_reply": "2026-09-11T17:30:51.137461Z" + "iopub.execute_input": "2026-10-03T02:30:27.094494Z", + "iopub.status.busy": "2026-10-03T02:30:27.094397Z", + "iopub.status.idle": "2026-10-03T02:30:27.199864Z", + "shell.execute_reply": "2026-10-03T02:30:27.199449Z" } }, "outputs": [ @@ -72,15 +72,15 @@ "name": "stdout", "output_type": "stream", "text": [ - "squared_error : R^2 test=0.591\n", - "absolute_error : R^2 test=0.553\n" + "squared_error : R^2 test=0.594\n", + "absolute_error : R^2 test=0.573\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "/var/folders/l5/dnbhhcmd4153b41c11_2v7p00000gn/T/ipykernel_89475/2183533904.py:2: UserWarning: Coordinate descent is disabled for the MAE objective. Set SGTLEARN_MAE_CD=1 to enable it.\n", + "/var/folders/l5/dnbhhcmd4153b41c11_2v7p00000gn/T/ipykernel_80066/2183533904.py:2: UserWarning: Coordinate descent is disabled for the MAE objective. Set SGTLEARN_MAE_CD=1 to enable it.\n", " r = SGTRegressor(max_depth=4, inner_max_depth=3, criterion=crit, random_state=0).fit(X_tr, y_tr)\n" ] } @@ -107,10 +107,10 @@ "id": "80e71be5", "metadata": { "execution": { - "iopub.execute_input": "2026-09-11T17:30:51.138914Z", - "iopub.status.busy": "2026-09-11T17:30:51.138857Z", - "iopub.status.idle": "2026-09-11T17:30:51.298894Z", - "shell.execute_reply": "2026-09-11T17:30:51.298513Z" + "iopub.execute_input": "2026-10-03T02:30:27.200848Z", + "iopub.status.busy": "2026-10-03T02:30:27.200796Z", + "iopub.status.idle": "2026-10-03T02:30:27.376623Z", + "shell.execute_reply": "2026-10-03T02:30:27.376243Z" } }, "outputs": [ @@ -118,7 +118,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "forest R^2 test=0.691\n" + "forest R^2 test=0.705\n" ] } ],