diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 41f7d43..bb50f40 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -84,6 +84,8 @@ add_library(sgtlearn_core STATIC src/Discretizers/univariate/GainHessianUnivariateDiscretizer.h src/algorithms/WaveletTreeMAE.h src/algorithms/WaveletTreeMAE.cpp + src/algorithms/WeightedMAETree.h + src/algorithms/WeightedMAETree.cpp src/Discretizers/univariate/GainHessianUnivariateDiscretizer.cpp src/Splitters/univariate/AbsoluteErrorSplitter.h src/Splitters/univariate/AbsoluteErrorSplitter.cpp @@ -119,8 +121,14 @@ add_library(sgtlearn_core STATIC src/Estimators/RegressionShapeGeneralizedTree.h src/Estimators/RegressionShapeGeneralizedTree.cpp src/BranchAssignmentObjectives/BranchAssignment.h + src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentCommon.h src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp + src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.h + src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.cpp + src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.h + src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.cpp + src/BranchAssignmentObjectives/MaeBranchConfig.h src/BranchAssignmentObjectives/LeafAggregateProcessor.h src/BranchAssignmentObjectives/LeafAggregationBranchAssignment.h src/BranchAssignmentObjectives/LeafAggregationBranchAssignment.cpp @@ -223,8 +231,10 @@ if (SGTLEARN_BUILD_TESTS) add_executable(cpp_tests tests/test_wavelet_tree_mae.cpp + tests/test_weighted_mae_tree.cpp tests/test_splitters.cpp tests/test_branch_assignment.cpp + tests/bench_mae_branch_assignment.cpp ) target_link_libraries(cpp_tests PRIVATE diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp index bf4d31b..5235161 100644 --- a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp @@ -1,14 +1,15 @@ /** * @file AbsoluteErrorBranchAssignment.cpp - * @brief MAE objective with partition medians over raw per-leaf target samples. + * @brief Default MAE branch assignment (sorted merge / filter). */ -#include #include "AbsoluteErrorBranchAssignment.h" +#include "AbsoluteErrorBranchAssignmentCommon.h" #include "Criterion.h" #include +#include #include AbsoluteErrorBranchAssignment::AbsoluteErrorBranchAssignment( @@ -19,80 +20,171 @@ AbsoluteErrorBranchAssignment::AbsoluteErrorBranchAssignment( : BranchAssignment(assignments, numPartitions, leafSampleCounts), leafYs_(leafYs), leafWs_(leafWs), leafWeights_(leafWeights) { - if (assignments.size() != leafYs.size() || leafYs.size() != leafWs.size() || - leafYs.size() != leafWeights.size()) - throw std::runtime_error( - "assignments, leafYs, leafWs, and leafWeights must have the same length"); - if (leafSampleCounts.size() != leafYs.size()) - throw std::runtime_error( - "leafSampleCounts must have the same length as bin statistics"); - - for (const auto &binOutputs : leafYs) { - if (!binOutputs.empty()) { - nOutputs_ = binOutputs.size(); - break; - } - } + absolute_error_branch::validateInputs(assignments, numPartitions, leafYs, + leafWs, leafWeights, leafSampleCounts, + nOutputs_); - for (size_t i = 0; i < leafYs.size(); ++i) { - if (assignments[i] >= numPartitions) - throw std::runtime_error("assignments[i] must be a valid partition index"); - for (const auto &outputYs : leafYs[i]) { - if (outputYs.size() != leafWs[i].size()) - throw std::runtime_error( - "leafYs[i][o] and leafWs[i] must have the same length"); - } - } + const size_t numLeaves = assignments.size(); + binYsSorted_.assign(numLeaves, {}); + binWsSorted_.assign(numLeaves, {}); + for (size_t b = 0; b < numLeaves; ++b) + buildSortedBin(b); partitionWeight_.assign(numPartitions, 0.0); partitionLoss_.assign(numPartitions, 0.0); + partYs_.assign(numPartitions, std::vector>(nOutputs_)); + partWs_.assign(numPartitions, std::vector>(nOutputs_)); + partSrcBin_.assign(numPartitions, + std::vector>(nOutputs_)); - const size_t numLeaves = assignments.size(); - for (size_t b = 0; b < numLeaves; b++) { - if (assignments[b] < numPartitions) { - partitionWeight_[assignments[b]] += leafWeights_[b]; - partitionSampleCount_[assignments[b]] += leafSampleCounts_[b]; - } + for (size_t b = 0; b < numLeaves; ++b) { + if (assignments[b] >= numPartitions) + continue; + partitionWeight_[assignments[b]] += leafWeights_[b]; + partitionSampleCount_[assignments[b]] += leafSampleCounts_[b]; + mergeLeafIntoPartition(b, assignments[b]); } - for (size_t p = 0; p < numPartitions; p++) { + for (size_t p = 0; p < numPartitions; ++p) { sumNumberOfSamples_ += partitionWeight_[p]; partitionLoss_[p] = computePartitionMae(p); weightedSumLoss_ += partitionWeight_[p] * partitionLoss_[p]; } } +void AbsoluteErrorBranchAssignment::buildSortedBin(size_t leaf) { + binYsSorted_[leaf].assign(nOutputs_, {}); + binWsSorted_[leaf].assign(nOutputs_, {}); + const auto &ws = leafWs_[leaf]; + for (size_t o = 0; o < nOutputs_; ++o) { + if (o >= leafYs_[leaf].size()) + continue; + const auto &ys = leafYs_[leaf][o]; + const size_t n = ys.size(); + std::vector order(n); + for (size_t i = 0; i < n; ++i) + order[i] = i; + std::sort(order.begin(), order.end(), + [&ys](size_t a, size_t b) { return ys[a] < ys[b]; }); + auto &ysOut = binYsSorted_[leaf][o]; + auto &wsOut = binWsSorted_[leaf][o]; + ysOut.resize(n); + wsOut.resize(n); + for (size_t i = 0; i < n; ++i) { + ysOut[i] = ys[order[i]]; + wsOut[i] = ws[order[i]]; + } + } +} + +void AbsoluteErrorBranchAssignment::mergeLeafIntoPartition(size_t leaf, + size_t partition) { + for (size_t o = 0; o < nOutputs_; ++o) { + const auto &binY = binYsSorted_[leaf][o]; + const auto &binW = binWsSorted_[leaf][o]; + auto &partY = partYs_[partition][o]; + auto &partW = partWs_[partition][o]; + auto &partSrc = partSrcBin_[partition][o]; + + if (binY.empty()) + continue; + if (partY.empty()) { + partY = binY; + partW = binW; + partSrc.assign(binY.size(), leaf); + continue; + } + + std::vector outY; + std::vector outW; + std::vector outSrc; + outY.reserve(partY.size() + binY.size()); + outW.reserve(partW.size() + binW.size()); + outSrc.reserve(partSrc.size() + binY.size()); + + size_t i = 0; + size_t j = 0; + while (i < partY.size() && j < binY.size()) { + if (partY[i] <= binY[j]) { + outY.push_back(partY[i]); + outW.push_back(partW[i]); + outSrc.push_back(partSrc[i]); + ++i; + } else { + outY.push_back(binY[j]); + outW.push_back(binW[j]); + outSrc.push_back(leaf); + ++j; + } + } + while (i < partY.size()) { + outY.push_back(partY[i]); + outW.push_back(partW[i]); + outSrc.push_back(partSrc[i]); + ++i; + } + while (j < binY.size()) { + outY.push_back(binY[j]); + outW.push_back(binW[j]); + outSrc.push_back(leaf); + ++j; + } + partY.swap(outY); + partW.swap(outW); + partSrc.swap(outSrc); + } +} + +void AbsoluteErrorBranchAssignment::filterLeafFromPartition(size_t leaf, + size_t partition) { + for (size_t o = 0; o < nOutputs_; ++o) { + auto &partY = partYs_[partition][o]; + auto &partW = partWs_[partition][o]; + auto &partSrc = partSrcBin_[partition][o]; + if (partY.empty()) + continue; + + std::vector outY; + std::vector outW; + std::vector outSrc; + outY.reserve(partY.size()); + outW.reserve(partW.size()); + outSrc.reserve(partSrc.size()); + for (size_t i = 0; i < partY.size(); ++i) { + if (partSrc[i] == leaf) + continue; + outY.push_back(partY[i]); + outW.push_back(partW[i]); + outSrc.push_back(partSrc[i]); + } + partY.swap(outY); + partW.swap(outW); + partSrc.swap(outSrc); + } +} + double AbsoluteErrorBranchAssignment::objective() { - // if (!allLeavesAssigned_) - // throw std::runtime_error( - // "Cannot compute objective if any leaves have been unassigned"); return sumNumberOfSamples_ > 0.0 ? weightedSumLoss_ / static_cast(sumNumberOfSamples_) : 0.0; } void AbsoluteErrorBranchAssignment::addLeaf(size_t leaf, size_t partition) { - // if (allLeavesAssigned_) - // throw std::runtime_error("Cannot assign a leaf if none ever left"); - weightedSumLoss_ -= partitionWeight_[partition] * partitionLoss_[partition]; partitionWeight_[partition] += leafWeights_[leaf]; partitionSampleCount_[partition] += leafSampleCounts_[leaf]; sumNumberOfSamples_ += leafWeights_[leaf]; + mergeLeafIntoPartition(leaf, partition); + partitionLoss_[partition] = computePartitionMae(partition); weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; assignments[leaf] = partition; - // allLeavesAssigned_ = true; } void AbsoluteErrorBranchAssignment::removeLeaf(size_t leaf) { - // if (!allLeavesAssigned_) - // throw std::runtime_error( - // "More than one leaf cannot be removed from the objective"); - const size_t partition = assignments[leaf]; if (partition >= numPartitions) throw std::runtime_error( @@ -103,52 +195,20 @@ void AbsoluteErrorBranchAssignment::removeLeaf(size_t leaf) { partitionWeight_[partition] -= leafWeights_[leaf]; partitionSampleCount_[partition] -= leafSampleCounts_[leaf]; + filterLeafFromPartition(leaf, partition); + partitionLoss_[partition] = computePartitionMae(partition); weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; - assignments[leaf] = kUnassignedPartition(numPartitions); - // allLeavesAssigned_ = false; -} - -void AbsoluteErrorBranchAssignment::collectPartitionSamples( - size_t partition, size_t output, std::vector &ys, - std::vector &ws) const { - ys.clear(); - ws.clear(); - for (size_t b = 0; b < assignments.size(); ++b) { - if (assignments[b] != partition) - continue; - if (output < leafYs_[b].size()) - ys.insert(ys.end(), leafYs_[b][output].begin(), leafYs_[b][output].end()); - ws.insert(ws.end(), leafWs_[b].begin(), leafWs_[b].end()); - } + assignments[leaf] = absolute_error_branch::unassignedPartition(numPartitions); } double AbsoluteErrorBranchAssignment::computePartitionMae(size_t partition) const { double total = 0.0; - std::vector ys; - std::vector ws; - for (size_t o = 0; o < nOutputs_; ++o) { - collectPartitionSamples(partition, o, ys, ws); - if (ys.size() <= 1) { - total += Criterion::absoluteError(ys, ws).mae; - continue; - } - std::vector order(ys.size()); - for (size_t i = 0; i < order.size(); ++i) - order[i] = i; - std::sort(order.begin(), order.end(), - [&ys](size_t a, size_t b) { return ys[a] < ys[b]; }); - std::vector ysSorted; - std::vector wsSorted; - ysSorted.reserve(ys.size()); - wsSorted.reserve(ws.size()); - for (size_t idx : order) { - ysSorted.push_back(ys[idx]); - wsSorted.push_back(ws[idx]); - } - total += Criterion::absoluteError(ysSorted, wsSorted).mae; - } + for (size_t o = 0; o < nOutputs_; ++o) + total += Criterion::absoluteErrorPresorted(partYs_[partition][o], + partWs_[partition][o]) + .mae; return total; } diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h index c53954b..8ca0fca 100644 --- a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h @@ -2,7 +2,11 @@ /** * @file AbsoluteErrorBranchAssignment.h - * @brief MAE branch assignment using per-leaf raw ``y`` samples and partition medians. + * @brief Default MAE branch assignment: sorted bins + merge/filter partitions. + * + * Production AbsoluteError backend. Deprecated alternatives: + * ``AbsoluteErrorBranchAssignmentBst``, ``AbsoluteErrorBranchAssignmentSort``. + * Hot-swap via ``SGTLEARN_MAE_BACKEND`` (default ``merge``). */ #include @@ -10,26 +14,20 @@ #include /** - * Multi-output MAE branch-assignment objective: per-partition loss is the SUM - * over outputs of the MAE about that output's median over the y values in the - * partition. Holds raw per-leaf, per-output y samples (``leafYs[bin][output]``) - * with per-sample weights shared across outputs (``leafWs[bin]``); add/remove - * recomputes the summed MAE for the affected partition(s). Single-output - * matches the old scalar path. + * Multi-output MAE using pre-sorted per-bin arrays and sorted partitions. + * Join: mergesort-style merge (``O(n + k)``). Leave: filter by source-bin id + * (``O(n)``). MAE uses ``Criterion::absoluteErrorPresorted``. */ class AbsoluteErrorBranchAssignment : public BranchAssignment { public: AbsoluteErrorBranchAssignment( std::vector &assignments, size_t numPartitions, std::vector>> &leafYs, - std::vector> &leafWs, - std::vector &leafWeights, + std::vector> &leafWs, std::vector &leafWeights, const std::vector &leafSampleCounts); double objective() override; - void addLeaf(size_t leaf, size_t partition) override; - void removeLeaf(size_t leaf) override; private: @@ -40,18 +38,24 @@ class AbsoluteErrorBranchAssignment : public BranchAssignment { double weightedSumLoss_ = 0; double sumNumberOfSamples_ = 0; - bool allLeavesAssigned_ = true; std::vector partitionWeight_; std::vector partitionLoss_; - void collectPartitionSamples(size_t partition, size_t output, - std::vector &ys, - std::vector &ws) const; - double computePartitionMae(size_t partition) const; + std::vector>> binYsSorted_; + std::vector>> binWsSorted_; - /** Valid partitions are [0, numPartitions); this marks a leaf not in any partition. */ - static constexpr size_t kUnassignedPartition(size_t numPartitions) { - return numPartitions; - } + std::vector>> partYs_; + std::vector>> partWs_; + std::vector>> partSrcBin_; + + void buildSortedBin(size_t leaf); + void mergeLeafIntoPartition(size_t leaf, size_t partition); + void filterLeafFromPartition(size_t leaf, size_t partition); + double computePartitionMae(size_t partition) const; }; + +/** @deprecated Prefer ``AbsoluteErrorBranchAssignment`` (merge is default). */ +using AbsoluteErrorBranchAssignmentMerge [[deprecated( + "Use AbsoluteErrorBranchAssignment; merge is the default backend")]] = + AbsoluteErrorBranchAssignment; diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.cpp b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.cpp new file mode 100644 index 0000000..79b43c6 --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.cpp @@ -0,0 +1,125 @@ +/** + * @file AbsoluteErrorBranchAssignmentBst.cpp + * @brief Deprecated AVL/order-statistic MAE branch assignment. + */ + +#include "AbsoluteErrorBranchAssignmentBst.h" + +#include "AbsoluteErrorBranchAssignmentCommon.h" + +#include +#include + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + +AbsoluteErrorBranchAssignmentBst::AbsoluteErrorBranchAssignmentBst( + std::vector &assignments, size_t numPartitions, + std::vector>> &leafYs, + std::vector> &leafWs, std::vector &leafWeights, + const std::vector &leafSampleCounts) + : BranchAssignment(assignments, numPartitions, leafSampleCounts), + leafYs_(leafYs), leafWs_(leafWs), leafWeights_(leafWeights) { + + absolute_error_branch::validateInputs(assignments, numPartitions, leafYs, + leafWs, leafWeights, leafSampleCounts, + nOutputs_); + + partitionWeight_.assign(numPartitions, 0.0); + partitionLoss_.assign(numPartitions, 0.0); + trees_.resize(numPartitions); + for (size_t p = 0; p < numPartitions; ++p) + trees_[p].resize(nOutputs_); + + const size_t numLeaves = assignments.size(); + for (size_t b = 0; b < numLeaves; ++b) { + if (assignments[b] >= numPartitions) + continue; + partitionWeight_[assignments[b]] += leafWeights_[b]; + partitionSampleCount_[assignments[b]] += leafSampleCounts_[b]; + insertLeafIntoPartition(b, assignments[b]); + } + + for (size_t p = 0; p < numPartitions; ++p) { + sumNumberOfSamples_ += partitionWeight_[p]; + partitionLoss_[p] = computePartitionMae(p); + weightedSumLoss_ += partitionWeight_[p] * partitionLoss_[p]; + } +} + +double AbsoluteErrorBranchAssignmentBst::objective() { + return sumNumberOfSamples_ > 0.0 + ? weightedSumLoss_ / static_cast(sumNumberOfSamples_) + : 0.0; +} + +void AbsoluteErrorBranchAssignmentBst::insertLeafIntoPartition( + size_t leaf, size_t partition) { + if (nOutputs_ == 0) + return; + const auto &ws = leafWs_[leaf]; + for (size_t o = 0; o < nOutputs_; ++o) { + if (o >= leafYs_[leaf].size()) + continue; + trees_[partition][o].insert_batch(leafYs_[leaf][o], ws); + } +} + +void AbsoluteErrorBranchAssignmentBst::eraseLeafFromPartition( + size_t leaf, size_t partition) { + if (nOutputs_ == 0) + return; + const auto &ws = leafWs_[leaf]; + for (size_t o = 0; o < nOutputs_; ++o) { + if (o >= leafYs_[leaf].size()) + continue; + trees_[partition][o].remove_batch(leafYs_[leaf][o], ws); + } +} + +void AbsoluteErrorBranchAssignmentBst::addLeaf(size_t leaf, size_t partition) { + weightedSumLoss_ -= partitionWeight_[partition] * partitionLoss_[partition]; + + partitionWeight_[partition] += leafWeights_[leaf]; + partitionSampleCount_[partition] += leafSampleCounts_[leaf]; + sumNumberOfSamples_ += leafWeights_[leaf]; + + insertLeafIntoPartition(leaf, partition); + + partitionLoss_[partition] = computePartitionMae(partition); + weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; + + assignments[leaf] = partition; +} + +void AbsoluteErrorBranchAssignmentBst::removeLeaf(size_t leaf) { + const size_t partition = assignments[leaf]; + if (partition >= numPartitions) + throw std::runtime_error( + "removeLeaf: leaf is not assigned to a valid partition"); + + weightedSumLoss_ -= partitionWeight_[partition] * partitionLoss_[partition]; + sumNumberOfSamples_ -= leafWeights_[leaf]; + partitionWeight_[partition] -= leafWeights_[leaf]; + partitionSampleCount_[partition] -= leafSampleCounts_[leaf]; + + eraseLeafFromPartition(leaf, partition); + + partitionLoss_[partition] = computePartitionMae(partition); + weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; + + assignments[leaf] = absolute_error_branch::unassignedPartition(numPartitions); +} + +double +AbsoluteErrorBranchAssignmentBst::computePartitionMae(size_t partition) const { + double total = 0.0; + for (size_t o = 0; o < nOutputs_; ++o) + total += trees_[partition][o].mae(); + return total; +} + +#pragma GCC diagnostic pop +#pragma clang diagnostic pop diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.h b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.h new file mode 100644 index 0000000..7f71fa1 --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.h @@ -0,0 +1,50 @@ +#pragma once + +/** + * @file AbsoluteErrorBranchAssignmentBst.h + * @brief Deprecated MAE branch assignment using per-partition WeightedMAETree. + */ + +#include +#include "BranchAssignment.h" +#include "algorithms/WeightedMAETree.h" +#include + +/** + * Multi-output MAE with per-partition ``WeightedMAETree`` multisets. + * + * @deprecated Prefer ``AbsoluteErrorBranchAssignment`` (merge/filter). Kept for + * benchmarks and A/B via ``SGTLEARN_MAE_BACKEND=bst``. + */ +class [[deprecated( + "Use AbsoluteErrorBranchAssignment (merge/filter); set " + "SGTLEARN_MAE_BACKEND=bst only for benchmarks")]] AbsoluteErrorBranchAssignmentBst + : public BranchAssignment { +public: + AbsoluteErrorBranchAssignmentBst( + std::vector &assignments, size_t numPartitions, + std::vector>> &leafYs, + std::vector> &leafWs, std::vector &leafWeights, + const std::vector &leafSampleCounts); + + double objective() override; + void addLeaf(size_t leaf, size_t partition) override; + void removeLeaf(size_t leaf) override; + +private: + std::vector>> &leafYs_; + std::vector> &leafWs_; + std::vector &leafWeights_; + size_t nOutputs_ = 0; + + double weightedSumLoss_ = 0; + double sumNumberOfSamples_ = 0; + + std::vector partitionWeight_; + std::vector partitionLoss_; + std::vector> trees_; + + double computePartitionMae(size_t partition) const; + void insertLeafIntoPartition(size_t leaf, size_t partition); + void eraseLeafFromPartition(size_t leaf, size_t partition); +}; diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentCommon.h b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentCommon.h new file mode 100644 index 0000000..8a3bc8a --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentCommon.h @@ -0,0 +1,51 @@ +#pragma once + +/** + * @file AbsoluteErrorBranchAssignmentCommon.h + * @brief Shared validation helpers for AbsoluteError branch-assignment backends. + */ + +#include +#include +#include + +namespace absolute_error_branch { + +inline constexpr size_t unassignedPartition(size_t numPartitions) { + return numPartitions; +} + +inline void validateInputs( + const std::vector &assignments, size_t numPartitions, + const std::vector>> &leafYs, + const std::vector> &leafWs, + const std::vector &leafWeights, + const std::vector &leafSampleCounts, size_t &nOutputs) { + if (assignments.size() != leafYs.size() || leafYs.size() != leafWs.size() || + leafYs.size() != leafWeights.size()) + throw std::runtime_error( + "assignments, leafYs, leafWs, and leafWeights must have the same length"); + if (leafSampleCounts.size() != leafYs.size()) + throw std::runtime_error( + "leafSampleCounts must have the same length as bin statistics"); + + nOutputs = 0; + for (const auto &binOutputs : leafYs) { + if (!binOutputs.empty()) { + nOutputs = binOutputs.size(); + break; + } + } + + for (size_t i = 0; i < leafYs.size(); ++i) { + if (assignments[i] >= numPartitions) + throw std::runtime_error("assignments[i] must be a valid partition index"); + for (const auto &outputYs : leafYs[i]) { + if (outputYs.size() != leafWs[i].size()) + throw std::runtime_error( + "leafYs[i][o] and leafWs[i] must have the same length"); + } + } +} + +} // namespace absolute_error_branch diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.cpp b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.cpp new file mode 100644 index 0000000..5d9d6a5 --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.cpp @@ -0,0 +1,128 @@ +/** + * @file AbsoluteErrorBranchAssignmentSort.cpp + * @brief Deprecated full re-sort MAE branch assignment. + */ + +#include "AbsoluteErrorBranchAssignmentSort.h" + +#include "AbsoluteErrorBranchAssignmentCommon.h" +#include "Criterion.h" + +#include +#include +#include + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + +AbsoluteErrorBranchAssignmentSort::AbsoluteErrorBranchAssignmentSort( + std::vector &assignments, size_t numPartitions, + std::vector>> &leafYs, + std::vector> &leafWs, std::vector &leafWeights, + const std::vector &leafSampleCounts) + : BranchAssignment(assignments, numPartitions, leafSampleCounts), + leafYs_(leafYs), leafWs_(leafWs), leafWeights_(leafWeights) { + + absolute_error_branch::validateInputs(assignments, numPartitions, leafYs, + leafWs, leafWeights, leafSampleCounts, + nOutputs_); + + partitionWeight_.assign(numPartitions, 0.0); + partitionLoss_.assign(numPartitions, 0.0); + + const size_t numLeaves = assignments.size(); + for (size_t b = 0; b < numLeaves; ++b) { + if (assignments[b] < numPartitions) { + partitionWeight_[assignments[b]] += leafWeights_[b]; + partitionSampleCount_[assignments[b]] += leafSampleCounts_[b]; + } + } + + for (size_t p = 0; p < numPartitions; ++p) { + sumNumberOfSamples_ += partitionWeight_[p]; + partitionLoss_[p] = computePartitionMae(p); + weightedSumLoss_ += partitionWeight_[p] * partitionLoss_[p]; + } +} + +double AbsoluteErrorBranchAssignmentSort::objective() { + return sumNumberOfSamples_ > 0.0 + ? weightedSumLoss_ / static_cast(sumNumberOfSamples_) + : 0.0; +} + +void AbsoluteErrorBranchAssignmentSort::addLeaf(size_t leaf, size_t partition) { + weightedSumLoss_ -= partitionWeight_[partition] * partitionLoss_[partition]; + + partitionWeight_[partition] += leafWeights_[leaf]; + partitionSampleCount_[partition] += leafSampleCounts_[leaf]; + sumNumberOfSamples_ += leafWeights_[leaf]; + assignments[leaf] = partition; + + partitionLoss_[partition] = computePartitionMae(partition); + weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; +} + +void AbsoluteErrorBranchAssignmentSort::removeLeaf(size_t leaf) { + const size_t partition = assignments[leaf]; + if (partition >= numPartitions) + throw std::runtime_error( + "removeLeaf: leaf is not assigned to a valid partition"); + + weightedSumLoss_ -= partitionWeight_[partition] * partitionLoss_[partition]; + sumNumberOfSamples_ -= leafWeights_[leaf]; + partitionWeight_[partition] -= leafWeights_[leaf]; + partitionSampleCount_[partition] -= leafSampleCounts_[leaf]; + assignments[leaf] = absolute_error_branch::unassignedPartition(numPartitions); + + partitionLoss_[partition] = computePartitionMae(partition); + weightedSumLoss_ += partitionWeight_[partition] * partitionLoss_[partition]; +} + +void AbsoluteErrorBranchAssignmentSort::collectPartitionSamples( + size_t partition, size_t output, std::vector &ys, + std::vector &ws) const { + ys.clear(); + ws.clear(); + for (size_t b = 0; b < assignments.size(); ++b) { + if (assignments[b] != partition) + continue; + if (output < leafYs_[b].size()) + ys.insert(ys.end(), leafYs_[b][output].begin(), leafYs_[b][output].end()); + ws.insert(ws.end(), leafWs_[b].begin(), leafWs_[b].end()); + } +} + +double +AbsoluteErrorBranchAssignmentSort::computePartitionMae(size_t partition) const { + double total = 0.0; + std::vector ys; + std::vector ws; + for (size_t o = 0; o < nOutputs_; ++o) { + collectPartitionSamples(partition, o, ys, ws); + if (ys.size() <= 1) { + total += Criterion::absoluteError(ys, ws).mae; + continue; + } + std::vector order(ys.size()); + for (size_t i = 0; i < order.size(); ++i) + order[i] = i; + std::sort(order.begin(), order.end(), + [&ys](size_t a, size_t b) { return ys[a] < ys[b]; }); + std::vector ysSorted; + std::vector wsSorted; + ysSorted.reserve(ys.size()); + wsSorted.reserve(ws.size()); + for (size_t idx : order) { + ysSorted.push_back(ys[idx]); + wsSorted.push_back(ws[idx]); + } + total += Criterion::absoluteErrorPresorted(ysSorted, wsSorted).mae; + } + return total; +} + +#pragma GCC diagnostic pop +#pragma clang diagnostic pop diff --git a/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.h b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.h new file mode 100644 index 0000000..065b643 --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.h @@ -0,0 +1,49 @@ +#pragma once + +/** + * @file AbsoluteErrorBranchAssignmentSort.h + * @brief Deprecated MAE branch assignment that re-sorts partitions on each move. + */ + +#include +#include "BranchAssignment.h" +#include + +/** + * Reference MAE path: collect partition samples and re-sort on every add/remove. + * + * @deprecated Prefer ``AbsoluteErrorBranchAssignment`` (merge/filter). Kept for + * benchmarks and A/B via ``SGTLEARN_MAE_BACKEND=sort``. + */ +class [[deprecated( + "Use AbsoluteErrorBranchAssignment (merge/filter); set " + "SGTLEARN_MAE_BACKEND=sort only for benchmarks")]] AbsoluteErrorBranchAssignmentSort + : public BranchAssignment { +public: + AbsoluteErrorBranchAssignmentSort( + std::vector &assignments, size_t numPartitions, + std::vector>> &leafYs, + std::vector> &leafWs, std::vector &leafWeights, + const std::vector &leafSampleCounts); + + double objective() override; + void addLeaf(size_t leaf, size_t partition) override; + void removeLeaf(size_t leaf) override; + +private: + std::vector>> &leafYs_; + std::vector> &leafWs_; + std::vector &leafWeights_; + size_t nOutputs_ = 0; + + double weightedSumLoss_ = 0; + double sumNumberOfSamples_ = 0; + + std::vector partitionWeight_; + std::vector partitionLoss_; + + void collectPartitionSamples(size_t partition, size_t output, + std::vector &ys, + std::vector &ws) const; + double computePartitionMae(size_t partition) const; +}; diff --git a/cpp/src/BranchAssignmentObjectives/BranchAssignmentFactory.cpp b/cpp/src/BranchAssignmentObjectives/BranchAssignmentFactory.cpp index 8e02861..cb374bc 100644 --- a/cpp/src/BranchAssignmentObjectives/BranchAssignmentFactory.cpp +++ b/cpp/src/BranchAssignmentObjectives/BranchAssignmentFactory.cpp @@ -3,15 +3,23 @@ * @brief Factory implementation for ``BranchAssignment`` objects. */ -#include #include +#include #include "BranchAssignmentFactory.h" #include "AbsoluteErrorBranchAssignment.h" +#include "AbsoluteErrorBranchAssignmentBst.h" +#include "AbsoluteErrorBranchAssignmentSort.h" #include "BranchAssignmentVariants.h" +#include "MaeBranchConfig.h" #include +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + std::unique_ptr makeBranchAssignment( LearningCriterion criterion, std::vector &assignments, size_t numPartitions, std::vector> &leafStats, @@ -26,9 +34,21 @@ std::unique_ptr makeBranchAssignment( throw std::invalid_argument( "makeBranchAssignment(AbsoluteError): maeLeafYs and maeLeafWs " "required"); - return std::make_unique( - assignments, numPartitions, *maeLeafYs, *maeLeafWs, leafWeights, - leafSampleCounts); + switch (mae_branch_config::backend()) { + case mae_branch_config::Backend::Sort: + return std::make_unique( + assignments, numPartitions, *maeLeafYs, *maeLeafWs, leafWeights, + leafSampleCounts); + case mae_branch_config::Backend::Bst: + return std::make_unique( + assignments, numPartitions, *maeLeafYs, *maeLeafWs, leafWeights, + leafSampleCounts); + case mae_branch_config::Backend::Merge: + default: + return std::make_unique( + assignments, numPartitions, *maeLeafYs, *maeLeafWs, leafWeights, + leafSampleCounts); + } case LearningCriterion::SquaredError: case LearningCriterion::Entropy: case LearningCriterion::Gini: @@ -41,6 +61,9 @@ std::unique_ptr makeBranchAssignment( } } +#pragma GCC diagnostic pop +#pragma clang diagnostic pop + std::unique_ptr makeBranchAssignment( LearningCriterion criterion, std::vector &assignments, size_t numPartitions, diff --git a/cpp/src/BranchAssignmentObjectives/BranchAssignmentVariants.h b/cpp/src/BranchAssignmentObjectives/BranchAssignmentVariants.h index 339d86c..c85cf04 100644 --- a/cpp/src/BranchAssignmentObjectives/BranchAssignmentVariants.h +++ b/cpp/src/BranchAssignmentObjectives/BranchAssignmentVariants.h @@ -7,6 +7,8 @@ #include #include "AbsoluteErrorBranchAssignment.h" +#include "AbsoluteErrorBranchAssignmentBst.h" +#include "AbsoluteErrorBranchAssignmentSort.h" #include "BranchAssignmentFactory.h" #include "LeafAggregationBranchAssignment.h" #include diff --git a/cpp/src/BranchAssignmentObjectives/MaeBranchConfig.h b/cpp/src/BranchAssignmentObjectives/MaeBranchConfig.h new file mode 100644 index 0000000..8061647 --- /dev/null +++ b/cpp/src/BranchAssignmentObjectives/MaeBranchConfig.h @@ -0,0 +1,39 @@ +#pragma once + +/** + * @file MaeBranchConfig.h + * @brief Runtime toggles for AbsoluteError branch assignment (benchmarking / experiments). + * + * Environment variables (read on each call): + * - ``SGTLEARN_MAE_BACKEND``: ``merge`` (default), ``bst``, or ``sort`` + * - ``SGTLEARN_MAE_CD``: ``1`` / ``true`` enables coordinate descent for + * ``absolute_error`` (default off for sklearn CART parity) + */ + +#include +#include + +namespace mae_branch_config { + +enum class Backend { Merge, Bst, Sort }; + +inline Backend backend() { + const char *v = std::getenv("SGTLEARN_MAE_BACKEND"); + if (v == nullptr) + return Backend::Merge; + if (std::strcmp(v, "sort") == 0 || std::strcmp(v, "Sort") == 0) + return Backend::Sort; + if (std::strcmp(v, "bst") == 0 || std::strcmp(v, "Bst") == 0) + return Backend::Bst; + // Explicit merge, unknown values, or empty → default merge. + return Backend::Merge; +} + +inline bool coordinateDescentEnabled() { + const char *v = std::getenv("SGTLEARN_MAE_CD"); + return v != nullptr && + (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || + std::strcmp(v, "TRUE") == 0 || std::strcmp(v, "yes") == 0); +} + +} // namespace mae_branch_config diff --git a/cpp/src/Criterion.cpp b/cpp/src/Criterion.cpp index 273b908..71e3a2a 100644 --- a/cpp/src/Criterion.cpp +++ b/cpp/src/Criterion.cpp @@ -69,40 +69,34 @@ double Criterion::squaredError( } Criterion::AbsoluteErrorStats -Criterion::absoluteError(const std::vector &ys, - const std::vector &weights) { +Criterion::absoluteErrorPresorted(const std::vector &ys, + const std::vector &weights) { AbsoluteErrorStats out; const size_t n = ys.size(); if (n == 0 || n != weights.size()) return out; - std::vector> pairs; - pairs.reserve(n); for (size_t i = 0; i < n; ++i) { const double w = static_cast(weights[i]); if (w < 0.0) - return out; - pairs.emplace_back(static_cast(ys[i]), w); + return AbsoluteErrorStats{}; out.totalWeight += w; } if (out.totalWeight <= 0.0) return out; - std::sort(pairs.begin(), pairs.end(), - [](const auto &a, const auto &b) { return a.first < b.first; }); - const double half = 0.5 * out.totalWeight; double wLeft = 0.0; double wyLeft = 0.0; double totalWy = 0.0; - for (const auto &[y, w] : pairs) - totalWy += w * y; + for (size_t i = 0; i < n; ++i) + totalWy += static_cast(weights[i]) * static_cast(ys[i]); - int medianRank = static_cast(pairs.size()) - 1; + int medianRank = static_cast(n) - 1; int medianPrevRank = medianRank > 0 ? medianRank - 1 : -1; bool found = false; - for (size_t rank = 0; rank < pairs.size(); ++rank) { - const double w = pairs[rank].second; + for (size_t rank = 0; rank < n; ++rank) { + const double w = static_cast(weights[rank]); if (wLeft + w > half) { medianRank = static_cast(rank); medianPrevRank = rank > 0 ? static_cast(rank - 1) : -1; @@ -110,18 +104,19 @@ Criterion::absoluteError(const std::vector &ys, break; } wLeft += w; - wyLeft += w * pairs[rank].first; + wyLeft += w * static_cast(ys[rank]); } if (!found) { - wLeft = out.totalWeight - pairs.back().second; - wyLeft = totalWy - pairs.back().second * pairs.back().first; + const double wLast = static_cast(weights.back()); + wLeft = out.totalWeight - wLast; + wyLeft = totalWy - wLast * static_cast(ys.back()); } if (medianPrevRank >= 0 && std::fabs(wLeft - half) <= 1e-12) { - out.median = 0.5 * (pairs[static_cast(medianPrevRank)].first + - pairs[static_cast(medianRank)].first); + out.median = 0.5 * (static_cast(ys[static_cast(medianPrevRank)]) + + static_cast(ys[static_cast(medianRank)])); } else { - out.median = pairs[static_cast(medianRank)].first; + out.median = static_cast(ys[static_cast(medianRank)]); } const double wRight = out.totalWeight - wLeft; @@ -132,6 +127,30 @@ Criterion::absoluteError(const std::vector &ys, return out; } +Criterion::AbsoluteErrorStats +Criterion::absoluteError(const std::vector &ys, + const std::vector &weights) { + AbsoluteErrorStats out; + const size_t n = ys.size(); + if (n == 0 || n != weights.size()) + return out; + + std::vector ysSorted; + std::vector wsSorted; + ysSorted.reserve(n); + wsSorted.reserve(n); + std::vector order(n); + for (size_t i = 0; i < n; ++i) + order[i] = i; + std::sort(order.begin(), order.end(), + [&ys](size_t a, size_t b) { return ys[a] < ys[b]; }); + for (size_t idx : order) { + ysSorted.push_back(ys[idx]); + wsSorted.push_back(weights[idx]); + } + return absoluteErrorPresorted(ysSorted, wsSorted); +} + double Criterion::gainAndHessian(const std::vector &derivatives, double lambda) { const double g = static_cast(derivatives[0]); diff --git a/cpp/src/Criterion.h b/cpp/src/Criterion.h index 54d27ee..50c483c 100644 --- a/cpp/src/Criterion.h +++ b/cpp/src/Criterion.h @@ -41,5 +41,12 @@ struct AbsoluteErrorStats { AbsoluteErrorStats absoluteError(const std::vector &ys, const std::vector &weights); +/** + * Same as ``absoluteError`` but assumes ``ys`` are already sorted ascending + * (weights aligned). Skips the internal sort — used by merge/filter MAE CD. + */ +AbsoluteErrorStats absoluteErrorPresorted(const std::vector &ys, + const std::vector &weights); + double gainAndHessian(const std::vector &derivatives, double lambda); } // namespace Criterion diff --git a/cpp/src/Estimators/ShapeFunctions/ShapeFunctionSplitSearch.cpp b/cpp/src/Estimators/ShapeFunctions/ShapeFunctionSplitSearch.cpp index 429a3b4..d17dc37 100644 --- a/cpp/src/Estimators/ShapeFunctions/ShapeFunctionSplitSearch.cpp +++ b/cpp/src/Estimators/ShapeFunctions/ShapeFunctionSplitSearch.cpp @@ -11,6 +11,7 @@ #include "BranchAssignmentObjectives/BranchAssignment.h" #include "BranchAssignmentObjectives/BranchAssignmentFactory.h" #include "BranchAssignmentObjectives/LeafAggregationBranchAssignment.h" +#include "BranchAssignmentObjectives/MaeBranchConfig.h" #include "algorithms/BinPartitionAssignments.h" #include "algorithms/CoordinateDescent.h" @@ -45,6 +46,32 @@ void refineShapeBranchAssignmentNested( leafSampleCounts, classesPerOutput, nOutputs); } +void refineShapeBranchAssignmentAbsoluteError( + std::unique_ptr &branchObj, size_t k, + size_t numRoutingBins, const CoordinateDescentParams &cdParams, + std::mt19937_64 &rng, + std::vector>> &maeLeafYs, + std::vector> &maeLeafWs, + std::vector &leafWeights, + const std::vector &leafSampleCounts) { + if (k >= numRoutingBins || !mae_branch_config::coordinateDescentEnabled()) + return; + + const std::vector snapshot = branchObj->assignments; + const double objBeforeCd = branchObj->objective(); + coordinateDescent(k, *branchObj, rng, cdParams.maxIters, cdParams.patience); + const double objAfterCd = branchObj->objective(); + if (std::isfinite(objAfterCd) && + objAfterCd <= objBeforeCd + kShapeFunctionCdImprovementEps) + return; + + std::vector rollback = snapshot; + std::vector> dummyLeafStats(maeLeafYs.size()); + branchObj = makeBranchAssignment( + LearningCriterion::AbsoluteError, rollback, k, dummyLeafStats, leafWeights, + leafSampleCounts, &maeLeafYs, &maeLeafWs); +} + void seedTrialBinAssignments(size_t k, size_t numRoutingBins, const std::vector> &stats, const std::vector &sizes, @@ -147,6 +174,9 @@ ShapeBranchAssignmentSearchResult searchShapeBranchAssignmentFromDiscretizer( branchObj = makeBranchAssignment(criterion, trialAssignments, k, dummyLeafStats, leafWeights, sizes, maeLeafYs, maeLeafWs); + refineShapeBranchAssignmentAbsoluteError( + branchObj, k, numRoutingBins, cdParams, rng, maeLeafYsStorage, + maeLeafWsStorage, leafWeights, sizes); } else { branchObj = makeBranchAssignment(criterion, trialAssignments, k, stats, leafWeights, sizes, classesPerOutput, diff --git a/cpp/src/algorithms/WeightedMAETree.cpp b/cpp/src/algorithms/WeightedMAETree.cpp new file mode 100644 index 0000000..2517d2f --- /dev/null +++ b/cpp/src/algorithms/WeightedMAETree.cpp @@ -0,0 +1,284 @@ +/** + * @file WeightedMAETree.cpp + * @brief Augmented AVL implementation for dynamic weighted median / MAE. + */ + +#include "algorithms/WeightedMAETree.h" + +#include +#include +#include + +namespace { + +constexpr double kWeightEps = 1e-15; +constexpr double kHalfTieEps = 1e-12; + +} // namespace + +void WeightedMAETree::pull(Node *n) { + if (!n) + return; + n->height = 1 + std::max(heightOf(n->left), heightOf(n->right)); + n->subWeight = n->weight + weightOf(n->left) + weightOf(n->right); + n->subWy = n->key * n->weight + wyOf(n->left) + wyOf(n->right); +} + +WeightedMAETree::Node *WeightedMAETree::rotateLeft(Node *x) { + Node *y = x->right; + x->right = y->left; + y->left = x; + pull(x); + pull(y); + return y; +} + +WeightedMAETree::Node *WeightedMAETree::rotateRight(Node *y) { + Node *x = y->left; + y->left = x->right; + x->right = y; + pull(y); + pull(x); + return x; +} + +WeightedMAETree::Node *WeightedMAETree::balance(Node *n) { + pull(n); + const int bf = heightOf(n->left) - heightOf(n->right); + if (bf > 1) { + if (heightOf(n->left->right) > heightOf(n->left->left)) + n->left = rotateLeft(n->left); + return rotateRight(n); + } + if (bf < -1) { + if (heightOf(n->right->left) > heightOf(n->right->right)) + n->right = rotateRight(n->right); + return rotateLeft(n); + } + return n; +} + +void WeightedMAETree::destroy(Node *n) { + if (!n) + return; + destroy(n->left); + destroy(n->right); + delete n; +} + +void WeightedMAETree::clear() { + destroy(root_); + root_ = nullptr; + totalWeight_ = 0.0; + totalWy_ = 0.0; +} + +WeightedMAETree::Node *WeightedMAETree::insertNode(Node *n, double key, + double w) { + if (!n) { + Node *created = new Node(); + created->key = key; + created->weight = w; + pull(created); + return created; + } + if (key < n->key) + n->left = insertNode(n->left, key, w); + else if (key > n->key) + n->right = insertNode(n->right, key, w); + else + n->weight += w; + return balance(n); +} + +WeightedMAETree::Node *WeightedMAETree::minNode(Node *n) { + while (n && n->left) + n = n->left; + return n; +} + +WeightedMAETree::Node *WeightedMAETree::eraseMin(Node *n) { + if (!n->left) + return n->right; + n->left = eraseMin(n->left); + return balance(n); +} + +WeightedMAETree::Node *WeightedMAETree::eraseNode(Node *n, double key, + double w) { + if (!n) + throw std::runtime_error("WeightedMAETree::erase: key not found"); + + if (key < n->key) + n->left = eraseNode(n->left, key, w); + else if (key > n->key) + n->right = eraseNode(n->right, key, w); + else { + n->weight -= w; + if (n->weight < -kWeightEps) + throw std::runtime_error("WeightedMAETree::erase: weight underflow"); + if (n->weight <= kWeightEps) { + Node *left = n->left; + Node *right = n->right; + delete n; + if (!right) + return left; + if (!left) + return right; + Node *m = minNode(right); + m->right = eraseMin(right); + m->left = left; + return balance(m); + } + } + return balance(n); +} + +void WeightedMAETree::insert(double y, double w) { + if (w < 0.0) + throw std::invalid_argument("WeightedMAETree::insert: negative weight"); + if (w <= kWeightEps) + return; + root_ = insertNode(root_, y, w); + totalWeight_ += w; + totalWy_ += y * w; +} + +void WeightedMAETree::erase(double y, double w) { + if (w < 0.0) + throw std::invalid_argument("WeightedMAETree::erase: negative weight"); + if (w <= kWeightEps) + return; + root_ = eraseNode(root_, y, w); + totalWeight_ -= w; + totalWy_ -= y * w; + if (totalWeight_ < 0.0 && totalWeight_ > -kWeightEps) + totalWeight_ = 0.0; + if (std::fabs(totalWy_) < kWeightEps) + totalWy_ = 0.0; +} + +void WeightedMAETree::insert_batch(const std::vector &ys, + const std::vector &ws) { + if (ys.size() != ws.size()) + throw std::invalid_argument( + "WeightedMAETree::insert_batch: ys/ws size mismatch"); + for (size_t i = 0; i < ys.size(); ++i) + insert(static_cast(ys[i]), static_cast(ws[i])); +} + +void WeightedMAETree::remove_batch(const std::vector &ys, + const std::vector &ws) { + if (ys.size() != ws.size()) + throw std::invalid_argument( + "WeightedMAETree::remove_batch: ys/ws size mismatch"); + for (size_t i = 0; i < ys.size(); ++i) + erase(static_cast(ys[i]), static_cast(ws[i])); +} + +void WeightedMAETree::aggregatesLessThan(const Node *n, double key, double &wOut, + double &wyOut) { + while (n) { + if (key <= n->key) { + n = n->left; + } else { + wOut += weightOf(n->left) + n->weight; + wyOut += wyOf(n->left) + n->key * n->weight; + n = n->right; + } + } +} + +bool WeightedMAETree::predecessor(const Node *n, double key, double &out) { + bool found = false; + while (n) { + if (n->key < key) { + out = n->key; + found = true; + n = n->right; + } else { + n = n->left; + } + } + return found; +} + +double WeightedMAETree::median() const { return medianAndMae().first; } + +double WeightedMAETree::mae() const { return medianAndMae().second; } + +std::pair WeightedMAETree::medianAndMae() const { + if (!root_ || totalWeight_ <= 0.0) + return {0.0, 0.0}; + + const double half = 0.5 * totalWeight_; + double wLeft = 0.0; + double wyLeft = 0.0; + const Node *n = root_; + const Node *medianNode = nullptr; + + while (n) { + const double leftW = weightOf(n->left); + if (wLeft + leftW > half) { + n = n->left; + continue; + } + if (wLeft + leftW + n->weight > half) { + wLeft += leftW; + wyLeft += wyOf(n->left); + medianNode = n; + break; + } + wLeft += leftW + n->weight; + wyLeft += wyOf(n->left) + n->key * n->weight; + n = n->right; + } + + if (!medianNode) { + // Degenerate: land on rightmost key (mirrors Criterion fallback). + n = root_; + while (n->right) + n = n->right; + medianNode = n; + wLeft = totalWeight_ - n->weight; + wyLeft = totalWy_ - n->key * n->weight; + } + + double m = medianNode->key; + if (std::fabs(wLeft - half) <= kHalfTieEps) { + double pred = 0.0; + if (predecessor(root_, medianNode->key, pred)) + m = 0.5 * (pred + medianNode->key); + } + + // Pinball about m with value split (y < m vs y > m); ties at m contribute 0. + double wLt = 0.0; + double wyLt = 0.0; + aggregatesLessThan(root_, m, wLt, wyLt); + + double wEq = 0.0; + double wyEq = 0.0; + if (std::fabs(m - medianNode->key) <= kHalfTieEps) { + wEq = medianNode->weight; + wyEq = medianNode->key * medianNode->weight; + } else { + // Half-tie average is not an inserted key; scan for an exact match anyway. + const Node *eq = root_; + while (eq) { + if (m < eq->key) + eq = eq->left; + else if (m > eq->key) + eq = eq->right; + else { + wEq = eq->weight; + wyEq = eq->key * eq->weight; + break; + } + } + } + + const double wGt = totalWeight_ - wLt - wEq; + const double wyGt = totalWy_ - wyLt - wyEq; + const double pinball = (m * wLt - wyLt) + (wyGt - m * wGt); + return {m, pinball / totalWeight_}; +} diff --git a/cpp/src/algorithms/WeightedMAETree.h b/cpp/src/algorithms/WeightedMAETree.h new file mode 100644 index 0000000..e3f87f6 --- /dev/null +++ b/cpp/src/algorithms/WeightedMAETree.h @@ -0,0 +1,110 @@ +#pragma once + +/** + * @file WeightedMAETree.h + * @brief Augmented AVL multiset for dynamic weighted median and pinball MAE. + * + * Keys are merged by value: each distinct ``y`` stores total weight at that + * value. Inserts/erases of a batch of size ``K`` are ``O(K log N)``. Median and + * MAE queries are ``O(log N)`` via subtree weight / ``Σw·y`` aggregates. + */ + +#include +#include +#include + +/** + * Self-balancing BST ordered by ``y``, maintaining per-subtree + * ``(Σw, Σw·y)`` so weighted median and MAE match ``Criterion::absoluteError`` + * under batch membership updates. + */ +class WeightedMAETree { +public: + WeightedMAETree() = default; + ~WeightedMAETree() { clear(); } + + WeightedMAETree(const WeightedMAETree &) = delete; + WeightedMAETree &operator=(const WeightedMAETree &) = delete; + + WeightedMAETree(WeightedMAETree &&other) noexcept + : root_(other.root_), totalWeight_(other.totalWeight_), + totalWy_(other.totalWy_) { + other.root_ = nullptr; + other.totalWeight_ = 0.0; + other.totalWy_ = 0.0; + } + + WeightedMAETree &operator=(WeightedMAETree &&other) noexcept { + if (this != &other) { + clear(); + root_ = other.root_; + totalWeight_ = other.totalWeight_; + totalWy_ = other.totalWy_; + other.root_ = nullptr; + other.totalWeight_ = 0.0; + other.totalWy_ = 0.0; + } + return *this; + } + + void clear(); + + /** Insert ``(y[i], w[i])`` for all ``i``; ``O(K log N)``. */ + void insert_batch(const std::vector &ys, + const std::vector &ws); + + /** Erase the same multiset of pairs previously inserted; ``O(K log N)``. */ + void remove_batch(const std::vector &ys, + const std::vector &ws); + + void insert(double y, double w); + void erase(double y, double w); + + double totalWeight() const { return totalWeight_; } + + /** Weighted median (sklearn half-tie average). ``O(log N)``. */ + double median() const; + + /** Mean absolute error about ``median()``. ``O(log N)``. */ + double mae() const; + + /** ``(median, mae)`` in one walk + aggregate query. */ + std::pair medianAndMae() const; + +private: + struct Node { + double key = 0.0; + double weight = 0.0; + double subWeight = 0.0; + double subWy = 0.0; + int height = 1; + Node *left = nullptr; + Node *right = nullptr; + }; + + Node *root_ = nullptr; + double totalWeight_ = 0.0; + double totalWy_ = 0.0; + + static int heightOf(const Node *n) { return n ? n->height : 0; } + static double weightOf(const Node *n) { return n ? n->subWeight : 0.0; } + static double wyOf(const Node *n) { return n ? n->subWy : 0.0; } + + static void pull(Node *n); + static Node *rotateLeft(Node *x); + static Node *rotateRight(Node *y); + static Node *balance(Node *n); + + Node *insertNode(Node *n, double key, double w); + Node *eraseNode(Node *n, double key, double w); + static Node *minNode(Node *n); + static Node *eraseMin(Node *n); + static void destroy(Node *n); + + /** Weight / ``Σw·y`` over keys strictly ``< key``. */ + static void aggregatesLessThan(const Node *n, double key, double &wOut, + double &wyOut); + + /** Largest key strictly less than ``key``, or false if none. */ + static bool predecessor(const Node *n, double key, double &out); +}; diff --git a/cpp/tests/bench_mae_branch_assignment.cpp b/cpp/tests/bench_mae_branch_assignment.cpp new file mode 100644 index 0000000..d9a759b --- /dev/null +++ b/cpp/tests/bench_mae_branch_assignment.cpp @@ -0,0 +1,239 @@ +/** + * @file bench_mae_branch_assignment.cpp + * @brief CD wall-time: sort vs BST vs merge AbsoluteError backends. + * + * Writes a concise colleague-facing CSV under ``benchmarks/results/``. + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using Catch::Matchers::WithinAbs; +using clock_type = std::chrono::steady_clock; + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + +namespace { + +struct MaeScenario { + std::vector assignments; + std::vector>> leafYs; + std::vector> leafWs; + std::vector leafWeights; + std::vector leafSampleCounts; + size_t numPartitions = 0; +}; + +MaeScenario makeScenario(size_t numBins, size_t samplesPerBin, + size_t numPartitions, size_t nOutputs, + uint64_t seed) { + std::mt19937_64 rng(seed); + std::uniform_real_distribution yDist(-100.0F, 100.0F); + std::uniform_real_distribution wDist(0.5F, 2.0F); + + MaeScenario s; + s.numPartitions = numPartitions; + s.assignments.resize(numBins); + s.leafYs.resize(numBins); + s.leafWs.resize(numBins); + s.leafWeights.resize(numBins); + s.leafSampleCounts.assign(numBins, samplesPerBin); + + for (size_t b = 0; b < numBins; ++b) { + s.assignments[b] = b % numPartitions; + s.leafYs[b].resize(nOutputs); + s.leafWs[b].resize(samplesPerBin); + double wSum = 0.0; + for (size_t i = 0; i < samplesPerBin; ++i) { + s.leafWs[b][i] = wDist(rng); + wSum += static_cast(s.leafWs[b][i]); + } + s.leafWeights[b] = wSum; + for (size_t o = 0; o < nOutputs; ++o) { + s.leafYs[b][o].resize(samplesPerBin); + for (size_t i = 0; i < samplesPerBin; ++i) + s.leafYs[b][o][i] = yDist(rng); + } + } + return s; +} + +struct BenchResult { + double ms = 0.0; + double objective = 0.0; +}; + +template +BenchResult timeCd(MaeScenario &base, uint64_t seed, int repeats, + MakeObj &&makeObj) { + double totalMs = 0.0; + double lastObj = 0.0; + for (int r = 0; r < repeats; ++r) { + auto asg = base.assignments; + auto obj = makeObj(asg); + std::mt19937_64 rng(seed + static_cast(r)); + const auto t0 = clock_type::now(); + lastObj = coordinateDescent(base.numPartitions, obj, rng, 8, 3); + const auto t1 = clock_type::now(); + totalMs += std::chrono::duration(t1 - t0).count(); + } + return {totalMs / static_cast(repeats), lastObj}; +} + +std::filesystem::path resultsDir() { + namespace fs = std::filesystem; + const fs::path candidates[] = { + fs::path("benchmarks") / "results", + fs::path("..") / "benchmarks" / "results", + fs::path("..") / ".." / "benchmarks" / "results", + }; + for (const auto &p : candidates) { + std::error_code ec; + if (fs::exists(p.parent_path(), ec)) + return p; + } + return fs::path("benchmarks") / "results"; +} + +} // namespace + +TEST_CASE("AbsoluteError BST/sort/merge objectives match under CD", + "[branch_assignment][absolute_error][correctness]") { + auto scenario = makeScenario(/*numBins=*/32, /*samplesPerBin=*/40, + /*numPartitions=*/4, /*nOutputs=*/1, /*seed=*/99); + + auto asgBst = scenario.assignments; + auto asgSort = scenario.assignments; + auto asgMerge = scenario.assignments; + + AbsoluteErrorBranchAssignmentBst bst(asgBst, scenario.numPartitions, + scenario.leafYs, scenario.leafWs, + scenario.leafWeights, + scenario.leafSampleCounts); + AbsoluteErrorBranchAssignmentSort sortObj(asgSort, scenario.numPartitions, + scenario.leafYs, scenario.leafWs, + scenario.leafWeights, + scenario.leafSampleCounts); + AbsoluteErrorBranchAssignment mergeObj(asgMerge, scenario.numPartitions, + scenario.leafYs, scenario.leafWs, + scenario.leafWeights, + scenario.leafSampleCounts); + + REQUIRE_THAT(bst.objective(), WithinAbs(sortObj.objective(), 1e-6)); + REQUIRE_THAT(mergeObj.objective(), WithinAbs(sortObj.objective(), 1e-6)); + + std::mt19937_64 rngBst(123); + std::mt19937_64 rngSort(123); + std::mt19937_64 rngMerge(123); + const double bstFinal = + coordinateDescent(scenario.numPartitions, bst, rngBst, 8, 3); + const double sortFinal = + coordinateDescent(scenario.numPartitions, sortObj, rngSort, 8, 3); + const double mergeFinal = + coordinateDescent(scenario.numPartitions, mergeObj, rngMerge, 8, 3); + REQUIRE_THAT(bstFinal, WithinAbs(sortFinal, 1e-5)); + REQUIRE_THAT(mergeFinal, WithinAbs(sortFinal, 1e-5)); + REQUIRE(asgBst == asgSort); + REQUIRE(asgMerge == asgSort); +} + +TEST_CASE("Bench AbsoluteError branch assignment: sort vs BST vs merge", + "[.benchmark]") { + struct Case { + const char *name; + size_t bins; + size_t samplesPerBin; + size_t parts; + size_t outputs; + int repeats; + }; + + const Case cases[] = { + {"small_64x20", 64, 20, 4, 1, 5}, + {"medium_128x50", 128, 50, 4, 1, 3}, + {"large_256x100", 256, 100, 8, 1, 2}, + {"multiout_128x40x3", 128, 40, 4, 3, 3}, + }; + + const auto outDir = resultsDir(); + std::filesystem::create_directories(outDir); + const auto csvPath = outDir / "mae_branch_cd_comparison.csv"; + std::ofstream csv(csvPath); + csv << "case,n_bins,samples_per_bin,n_partitions,n_outputs," + "sort_ms,bst_ms,merge_ms," + "bst_speedup_vs_sort,merge_speedup_vs_sort," + "sort_obj,bst_obj,merge_obj\n"; + + std::cout << '\n' + << "AbsoluteError CD bench: sort | bst | merge (default)\n" + << "CSV -> " << csvPath << '\n' + << "--------------------------------------------------------------" + "--------\n"; + + for (const Case &c : cases) { + auto scenario = + makeScenario(c.bins, c.samplesPerBin, c.parts, c.outputs, 2026); + + const auto sortRes = timeCd(scenario, 7, c.repeats, [&](auto &asg) { + return AbsoluteErrorBranchAssignmentSort( + asg, scenario.numPartitions, scenario.leafYs, scenario.leafWs, + scenario.leafWeights, scenario.leafSampleCounts); + }); + const auto bstRes = timeCd(scenario, 7, c.repeats, [&](auto &asg) { + return AbsoluteErrorBranchAssignmentBst( + asg, scenario.numPartitions, scenario.leafYs, scenario.leafWs, + scenario.leafWeights, scenario.leafSampleCounts); + }); + const auto mergeRes = timeCd(scenario, 7, c.repeats, [&](auto &asg) { + return AbsoluteErrorBranchAssignment( + asg, scenario.numPartitions, scenario.leafYs, scenario.leafWs, + scenario.leafWeights, scenario.leafSampleCounts); + }); + + const double bstSpeedup = + bstRes.ms > 0.0 ? (sortRes.ms / bstRes.ms) : 0.0; + const double mergeSpeedup = + mergeRes.ms > 0.0 ? (sortRes.ms / mergeRes.ms) : 0.0; + + std::cout << std::fixed << std::setprecision(2) << c.name << ": sort " + << sortRes.ms << " ms | bst " << bstRes.ms << " ms (" + << bstSpeedup << "x) | merge " << mergeRes.ms << " ms (" + << mergeSpeedup << "x)\n"; + + csv << std::fixed << std::setprecision(3) << c.name << ',' << c.bins << ',' + << c.samplesPerBin << ',' << c.parts << ',' << c.outputs << ',' + << sortRes.ms << ',' << bstRes.ms << ',' << mergeRes.ms << ',' + << std::setprecision(3) << bstSpeedup << ',' << mergeSpeedup << ',' + << std::setprecision(6) << sortRes.objective << ',' << bstRes.objective + << ',' << mergeRes.objective << '\n'; + + REQUIRE(sortRes.ms > 0.0); + REQUIRE(bstRes.ms > 0.0); + REQUIRE(mergeRes.ms > 0.0); + REQUIRE_THAT(bstRes.objective, WithinAbs(sortRes.objective, 1e-4)); + REQUIRE_THAT(mergeRes.objective, WithinAbs(sortRes.objective, 1e-4)); + } + + csv.flush(); + std::cout << "Wrote " << csvPath << std::endl; +} + +#pragma GCC diagnostic pop +#pragma clang diagnostic pop diff --git a/cpp/tests/test_weighted_mae_tree.cpp b/cpp/tests/test_weighted_mae_tree.cpp new file mode 100644 index 0000000..7853d96 --- /dev/null +++ b/cpp/tests/test_weighted_mae_tree.cpp @@ -0,0 +1,93 @@ +/** + * @file test_weighted_mae_tree.cpp + * @brief Correctness tests for ``WeightedMAETree`` vs ``Criterion::absoluteError``. + */ + +#include +#include + +#include +#include + +#include +#include + +using Catch::Matchers::WithinAbs; + +namespace { + +Criterion::AbsoluteErrorStats brute(const std::vector &ys, + const std::vector &ws) { + return Criterion::absoluteError(ys, ws); +} + +} // namespace + +TEST_CASE("WeightedMAETree median/mae match Criterion on random batches", + "[weighted_mae_tree]") { + std::mt19937 rng(7); + std::uniform_real_distribution yDist(-50.0F, 50.0F); + std::uniform_real_distribution wDist(0.1F, 3.0F); + + for (int n : {1, 2, 3, 10, 64, 257}) { + std::vector ys(static_cast(n)); + std::vector ws(static_cast(n)); + for (int i = 0; i < n; ++i) { + ys[static_cast(i)] = yDist(rng); + ws[static_cast(i)] = wDist(rng); + } + + WeightedMAETree tree; + tree.insert_batch(ys, ws); + const auto ref = brute(ys, ws); + const auto got = tree.medianAndMae(); + REQUIRE_THAT(got.first, WithinAbs(ref.median, 1e-5)); + REQUIRE_THAT(got.second, WithinAbs(ref.mae, 1e-5)); + REQUIRE_THAT(tree.totalWeight(), WithinAbs(ref.totalWeight, 1e-6)); + } +} + +TEST_CASE("WeightedMAETree supports remove_batch round-trip", + "[weighted_mae_tree]") { + std::vector ys = {1.F, 2.F, 3.F, 4.F, 5.F, 2.F}; + std::vector ws = {1.F, 1.F, 2.F, 1.F, 1.F, 0.5F}; + std::vector dropY = {2.F, 4.F}; + std::vector dropW = {1.F, 1.F}; + + WeightedMAETree tree; + tree.insert_batch(ys, ws); + tree.remove_batch(dropY, dropW); + + std::vector remainY = {1.F, 3.F, 5.F, 2.F}; + std::vector remainW = {1.F, 2.F, 1.F, 0.5F}; + const auto ref = brute(remainY, remainW); + const auto got = tree.medianAndMae(); + REQUIRE_THAT(got.first, WithinAbs(ref.median, 1e-6)); + REQUIRE_THAT(got.second, WithinAbs(ref.mae, 1e-6)); +} + +TEST_CASE("WeightedMAETree half-tie median matches Criterion", + "[weighted_mae_tree]") { + // Equal total weight on each side of the cut → average of adjacent keys. + std::vector ys = {1.F, 3.F}; + std::vector ws = {1.F, 1.F}; + WeightedMAETree tree; + tree.insert_batch(ys, ws); + const auto ref = brute(ys, ws); + REQUIRE_THAT(tree.median(), WithinAbs(ref.median, 1e-12)); + REQUIRE_THAT(tree.mae(), WithinAbs(ref.mae, 1e-12)); + REQUIRE_THAT(tree.median(), WithinAbs(2.0, 1e-12)); +} + +TEST_CASE("WeightedMAETree duplicate keys", "[weighted_mae_tree]") { + std::vector ys(100, 4.5F); + std::vector ws(100, 0.25F); + ys[0] = -10.F; + ws[0] = 1.F; + + WeightedMAETree tree; + tree.insert_batch(ys, ws); + const auto ref = brute(ys, ws); + REQUIRE_THAT(tree.median(), WithinAbs(ref.median, 1e-6)); + REQUIRE_THAT(tree.mae(), WithinAbs(ref.mae, 1e-6)); +} diff --git a/sgtlearn/__init__.py b/sgtlearn/__init__.py index 065eccf..e3ef573 100644 --- a/sgtlearn/__init__.py +++ b/sgtlearn/__init__.py @@ -4,29 +4,29 @@ ``Discretizers``). Import ``SGTClassifier`` from this package for the sklearn-style API. """ +from sgtlearn import tao +from sgtlearn._export import export_graphviz, export_text, plot_tree from sgtlearn.base import ( BaseShapeCART, + ProcessedFeatures, SGTClassifier, SGTRegressor, - ProcessedFeatures, configure_feature_dict, ) -from sgtlearn.ensemble import RandomSGForestClassifier, RandomSGForestRegressor -from sgtlearn._export import export_graphviz, export_text, plot_tree from sgtlearn.datasets import make_plus -from sgtlearn import tao +from sgtlearn.ensemble import RandomSGForestClassifier, RandomSGForestRegressor __all__ = [ "BaseShapeCART", - "SGTClassifier", - "SGTRegressor", "ProcessedFeatures", - "configure_feature_dict", "RandomSGForestClassifier", "RandomSGForestRegressor", + "SGTClassifier", + "SGTRegressor", + "configure_feature_dict", "export_graphviz", "export_text", - "plot_tree", "make_plus", + "plot_tree", "tao", ] diff --git a/sgtlearn/_export.py b/sgtlearn/_export.py index 41545fd..66d59cd 100644 --- a/sgtlearn/_export.py +++ b/sgtlearn/_export.py @@ -9,11 +9,13 @@ from __future__ import annotations -from typing import Any, Optional, Sequence, Union +from collections.abc import Sequence +from typing import Any + import numpy as np from matplotlib.patches import FancyArrowPatch -__all__ = ["plot_tree", "export_graphviz", "export_text"] +__all__ = ["export_graphviz", "export_text", "plot_tree"] import matplotlib.pyplot as plt from sklearn.utils.validation import check_is_fitted @@ -198,7 +200,7 @@ def _merge_routing_regions( return regions -def _route_samples(tree: dict, X) -> "dict[int, Any]": +def _route_samples(tree: dict, X) -> dict[int, Any]: """Route ``X`` through the tree; return ``{node_id: column-indices}``. The returned array for each node lists the row indices of ``X`` that @@ -273,7 +275,7 @@ def _route_samples(tree: dict, X) -> "dict[int, Any]": def _compute_layout_leafcounter( - tree: dict, max_depth: Optional[int] + tree: dict, max_depth: int | None ) -> dict[int, tuple[float, float]]: """Leaf-counter layout in axes coords [0, 1]. @@ -290,9 +292,7 @@ def is_draw_leaf(nid: int, depth: int) -> bool: n = nodes_by_id[nid] if n["is_leaf"]: return True - if max_depth is not None and depth >= max_depth: - return True - return False + return bool(max_depth is not None and depth >= max_depth) x_int: dict[int, float] = {} counter = [0] @@ -388,10 +388,10 @@ def _draw_leaf_text( node: dict, *, is_classifier: bool, - class_names: Optional[list[str]], + class_names: list[str] | None, criterion: str, precision: int, - fontsize: Optional[int], + fontsize: int | None, color, label: str, impurity: bool, @@ -458,7 +458,7 @@ def _draw_internal_panel_categorical( palette, feat_names: list[str], X_rows: np.ndarray | None, - fontsize: Optional[int], + fontsize: int | None, label: str, ) -> list: cx, cy = center @@ -551,7 +551,7 @@ def _draw_internal_panel( feat_names: list[str], n_hist_bins: int, precision: int, - fontsize: Optional[int], + fontsize: int | None, label: str, ) -> list: """Render a single internal node panel: slabs + optional fine histogram.""" @@ -663,16 +663,16 @@ def plot_tree( estimator: Any, *, X=None, - max_depth: Optional[int] = None, - feature_names: Optional[list[str]] = None, - class_names: Union[list[str], bool, None] = None, + max_depth: int | None = None, + feature_names: list[str] | None = None, + class_names: list[str] | bool | None = None, label: str = "feature", impurity: bool = False, proportion: bool = False, precision: int = 2, cmap: Any = _DEFAULT_PALETTE_COLORS, - ax: Optional[plt.Axes] = None, - fontsize: Optional[int] = None, + ax: plt.Axes | None = None, + fontsize: int | None = None, node_aspect_ratio: float = 2.5, n_hist_bins: int = 20, ) -> list[Any]: @@ -719,7 +719,7 @@ def plot_tree( palette = _build_palette(cmap, tree["num_partitions"]) is_classifier = isinstance(estimator, SGTClassifier) - resolved_class_names: Optional[list[str]] + resolved_class_names: list[str] | None if not is_classifier: resolved_class_names = None elif class_names is True: diff --git a/sgtlearn/_features.py b/sgtlearn/_features.py index d66a99f..543654a 100644 --- a/sgtlearn/_features.py +++ b/sgtlearn/_features.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass -from typing import Any, Mapping, MutableMapping, Sequence +from typing import Any FeatureInfoDict = dict[str, Any] FeatureDict = Mapping[int | str, Sequence[int | str]] diff --git a/sgtlearn/_multioutput.py b/sgtlearn/_multioutput.py index 831d8ca..10246d3 100644 --- a/sgtlearn/_multioutput.py +++ b/sgtlearn/_multioutput.py @@ -7,7 +7,8 @@ from __future__ import annotations -from typing import Any, Sequence, Union +from collections.abc import Sequence +from typing import Any import numpy as np from sklearn.preprocessing import LabelEncoder @@ -15,10 +16,10 @@ __all__ = [ "as_output_matrix", "encode_classification_targets", - "unwrap_classifier_public_attrs", "label_encoders_as_list", "native_y_array", "squeeze_outputs", + "unwrap_classifier_public_attrs", ] @@ -44,7 +45,7 @@ def as_output_matrix(y: Any) -> tuple[np.ndarray, int]: def encode_classification_targets( y: Any, *, - encoders: Union[None, LabelEncoder, Sequence[Any]] = None, + encoders: None | LabelEncoder | Sequence[Any] = None, ) -> tuple[np.ndarray, list[Any], list[np.ndarray], list[int]]: """Encode labels with one encoder per output. @@ -75,13 +76,13 @@ def encode_classification_targets( cols.append(le.fit_transform(y2[:, o])) fitted.append(le) classes_list.append(np.asarray(le.classes_)) - n_classes_list.append(int(len(le.classes_))) + n_classes_list.append(len(le.classes_)) return np.column_stack(cols), fitted, classes_list, n_classes_list enc_list = label_encoders_as_list(encoders, n_outputs) cols = [enc_list[o].transform(y2[:, o]) for o in range(n_outputs)] classes_list = [np.asarray(enc_list[o].classes_) for o in range(n_outputs)] - n_classes_list = [int(len(c)) for c in classes_list] + n_classes_list = [len(c) for c in classes_list] return np.column_stack(cols), enc_list, classes_list, n_classes_list @@ -112,9 +113,7 @@ def unwrap_classifier_public_attrs( ) -def label_encoders_as_list( - encoders: Union[Any, Sequence[Any]], n_outputs: int -) -> list[Any]: +def label_encoders_as_list(encoders: Any | Sequence[Any], n_outputs: int) -> list[Any]: """Normalize a scalar encoder or sequence to length ``n_outputs``.""" if isinstance(encoders, (list, tuple)): enc_list = list(encoders) diff --git a/sgtlearn/_weights.py b/sgtlearn/_weights.py index 8b2f5b4..13665ac 100644 --- a/sgtlearn/_weights.py +++ b/sgtlearn/_weights.py @@ -2,16 +2,18 @@ from __future__ import annotations -from collections.abc import Mapping as ABCMapping, Sequence as ABCSequence -from typing import Any, Mapping, Optional, Sequence, Union +from collections.abc import Mapping, Sequence +from collections.abc import Mapping as ABCMapping +from collections.abc import Sequence as ABCSequence +from typing import Any import numpy as np from sgtlearn._multioutput import as_output_matrix __all__ = [ - "normalize_sample_weight", "effective_sample_weight_classification", + "normalize_sample_weight", ] @@ -28,8 +30,8 @@ def _validate_sample_weight_array(sw: np.ndarray, n_samples: int) -> None: def normalize_sample_weight( - sample_weight: Optional[np.ndarray], n_samples: int -) -> Optional[np.ndarray]: + sample_weight: np.ndarray | None, n_samples: int +) -> np.ndarray | None: """Validated float64 weights for tree ``fit``, or ``None`` for uniform weighting.""" if sample_weight is None: return None @@ -60,10 +62,10 @@ def _per_class_multiplier( def effective_sample_weight_classification( - sample_weight: Optional[np.ndarray], + sample_weight: np.ndarray | None, y_enc: np.ndarray, - class_weight: Union[Mapping[Any, float], Sequence[Mapping[Any, float]]], - classes_: Union[np.ndarray, Sequence[np.ndarray]], + class_weight: Mapping[Any, float] | Sequence[Mapping[Any, float]], + classes_: np.ndarray | Sequence[np.ndarray], ) -> np.ndarray: """``sample_weight * class_weight[y]`` in encoded label space. @@ -94,7 +96,7 @@ def effective_sample_weight_classification( f"output ({n_outputs}); got {len(cw_list)}" ) else: - raise ValueError("class_weight must be a mapping or a sequence of mappings") + raise TypeError("class_weight must be a mapping or a sequence of mappings") n = y2.shape[0] if sample_weight is None: diff --git a/sgtlearn/base.py b/sgtlearn/base.py index dd08305..c3d1639 100644 --- a/sgtlearn/base.py +++ b/sgtlearn/base.py @@ -2,11 +2,17 @@ from __future__ import annotations -from typing import Any, Mapping, Optional, Sequence, Union +from collections.abc import Mapping, Sequence +from typing import Any import numpy as np +from ShapeGeneralizedTrees import ( + ClassificationShapeGeneralizedTree, + RegressionShapeGeneralizedTree, +) from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin from sklearn.exceptions import NotFittedError +from sklearn.preprocessing import LabelEncoder from sklearn.utils.validation import check_array, check_is_fitted, check_X_y from sgtlearn._features import ProcessedFeatures, configure_feature_dict @@ -19,20 +25,15 @@ unwrap_classifier_public_attrs, ) from sgtlearn._weights import ( - normalize_sample_weight, effective_sample_weight_classification, + normalize_sample_weight, ) -from ShapeGeneralizedTrees import ( - ClassificationShapeGeneralizedTree, - RegressionShapeGeneralizedTree, -) -from sklearn.preprocessing import LabelEncoder __all__ = [ "BaseShapeCART", + "ProcessedFeatures", "SGTClassifier", "SGTRegressor", - "ProcessedFeatures", "configure_feature_dict", ] @@ -73,7 +74,7 @@ class _IdentityLabelEncoder(LabelEncoder): def __init__(self, classes_: np.ndarray) -> None: self.classes_ = np.asarray(classes_) - def fit(self, y: np.ndarray) -> "_IdentityLabelEncoder": + def fit(self, y: np.ndarray) -> _IdentityLabelEncoder: raise NotImplementedError( "_IdentityLabelEncoder is built with preset classes_; " "fit the enclosing meta-estimator instead." @@ -293,8 +294,8 @@ def __init__( *, criterion: str = "gini", num_partitions: int = 2, - max_depth: Optional[int] = None, - max_leaf_nodes: Optional[int] = None, + max_depth: int | None = None, + max_leaf_nodes: int | None = None, min_samples_leaf: int = 1, min_impurity_decrease: float = 0.0, inner_max_depth: int = 3, @@ -304,11 +305,9 @@ def __init__( coordinate_descent_max_iters: int = 20, coordinate_descent_patience: int = 5, coordinate_descent_smart_init: bool = True, - random_state: Optional[int] = 42, - max_features: Optional[Union[int, float, str]] = None, - class_weight: Optional[ - Union[Mapping[Any, float], Sequence[Mapping[Any, float]]] - ] = None, + random_state: int | None = 42, + max_features: float | str | None = None, + class_weight: Mapping[Any, float] | Sequence[Mapping[Any, float]] | None = None, tao_n_runs: int = 10, tao_lambda: float = 0.0, ) -> None: @@ -334,22 +333,22 @@ def __init__( self.tao_lambda = tao_lambda self._est: Any = None self._le: Any = None - self.classes_: Optional[Any] = None - self.n_classes_: Optional[Any] = None + self.classes_: Any | None = None + self.n_classes_: Any | None = None self.n_outputs_: int = 1 - self.n_features_in_: Optional[int] = None - self.feature_names_in_: Optional[np.ndarray] = None + self.n_features_in_: int | None = None + self.feature_names_in_: np.ndarray | None = None def fit( self, X: np.ndarray, y: np.ndarray, - sample_weight: Optional[np.ndarray] = None, + sample_weight: np.ndarray | None = None, *, - feature_dict: Optional[Mapping[int | str, Sequence[int | str]]] = None, - processed_features: Optional[ProcessedFeatures] = None, + feature_dict: Mapping[int | str, Sequence[int | str]] | None = None, + processed_features: ProcessedFeatures | None = None, check_input: bool = True, - ) -> "SGTClassifier": + ) -> SGTClassifier: """Fit the tree on ``X`` and class labels ``y``. Parameters @@ -438,7 +437,7 @@ def fit( if y_enc.shape[0] != X.shape[0]: raise ValueError("X and y must have the same number of samples.") - sw: Optional[np.ndarray] = None + sw: np.ndarray | None = None if self.class_weight is not None: sw = effective_sample_weight_classification( sample_weight, y_enc, self.class_weight, self.classes_ @@ -449,7 +448,7 @@ def fit( self.n_features_in_ = X.shape[1] if column_names is None: column_names = _column_names_from_X(X) - self.feature_names_in_: Optional[np.ndarray] = ( + self.feature_names_in_: np.ndarray | None = ( np.asarray(column_names, dtype=object) if column_names is not None else None ) @@ -682,8 +681,8 @@ def __init__( *, criterion: str = "squared_error", num_partitions: int = 2, - max_depth: Optional[int] = None, - max_leaf_nodes: Optional[int] = None, + max_depth: int | None = None, + max_leaf_nodes: int | None = None, min_samples_leaf: int = 1, min_impurity_decrease: float = 0.0, inner_max_depth: int = 3, @@ -693,8 +692,8 @@ def __init__( coordinate_descent_max_iters: int = 20, coordinate_descent_patience: int = 5, coordinate_descent_smart_init: bool = True, - random_state: Optional[int] = 42, - max_features: Optional[Union[int, float, str]] = None, + random_state: int | None = 42, + max_features: float | str | None = None, tao_n_runs: int = 10, tao_lambda: float = 0.0, ) -> None: @@ -717,19 +716,19 @@ def __init__( self.tao_lambda = tao_lambda self._est: Any = None self.n_outputs_: int = 1 - self.n_features_in_: Optional[int] = None - self.feature_names_in_: Optional[np.ndarray] = None + self.n_features_in_: int | None = None + self.feature_names_in_: np.ndarray | None = None def fit( self, X: np.ndarray, y: np.ndarray, - sample_weight: Optional[np.ndarray] = None, + sample_weight: np.ndarray | None = None, *, - feature_dict: Optional[Mapping[int | str, Sequence[int | str]]] = None, - processed_features: Optional[ProcessedFeatures] = None, + feature_dict: Mapping[int | str, Sequence[int | str]] | None = None, + processed_features: ProcessedFeatures | None = None, check_input: bool = True, - ) -> "SGTRegressor": + ) -> SGTRegressor: """Fit the tree on ``X`` and continuous targets ``y``. Parameters @@ -776,7 +775,7 @@ def fit( self.n_features_in_ = X.shape[1] if column_names is None: column_names = _column_names_from_X(X) - self.feature_names_in_: Optional[np.ndarray] = ( + self.feature_names_in_: np.ndarray | None = ( np.asarray(column_names, dtype=object) if column_names is not None else None ) diff --git a/sgtlearn/datasets.py b/sgtlearn/datasets.py index aed688b..55724c0 100644 --- a/sgtlearn/datasets.py +++ b/sgtlearn/datasets.py @@ -2,8 +2,6 @@ from __future__ import annotations -from typing import Optional - import numpy as np __all__ = ["make_plus"] @@ -14,7 +12,7 @@ def make_plus( *, grid: int = 3, margin: float = 0.05, - random_state: Optional[int] = None, + random_state: int | None = None, ) -> tuple[np.ndarray, np.ndarray]: """Generate the "Plus Sign" dataset. diff --git a/sgtlearn/ensemble/__init__.py b/sgtlearn/ensemble/__init__.py index 7909437..a1e519c 100644 --- a/sgtlearn/ensemble/__init__.py +++ b/sgtlearn/ensemble/__init__.py @@ -1,4 +1,4 @@ -from sgtlearn.ensemble.RandomSGForestClassifier import RandomSGForestClassifier -from sgtlearn.ensemble.RandomSGForestRegressor import RandomSGForestRegressor +from sgtlearn.ensemble.random_sgforest_classifier import RandomSGForestClassifier +from sgtlearn.ensemble.random_sgforest_regressor import RandomSGForestRegressor __all__ = ["RandomSGForestClassifier", "RandomSGForestRegressor"] diff --git a/sgtlearn/ensemble/_random_sgforest.py b/sgtlearn/ensemble/_random_sgforest.py index 12598dc..27d2067 100644 --- a/sgtlearn/ensemble/_random_sgforest.py +++ b/sgtlearn/ensemble/_random_sgforest.py @@ -3,8 +3,9 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence from numbers import Integral -from typing import Any, Mapping, Optional, Sequence, Union +from typing import Any import numpy as np from joblib import Parallel, delayed, effective_n_jobs @@ -12,14 +13,12 @@ from sklearn.utils import check_random_state from sklearn.utils.validation import check_array, check_is_fitted -from sgtlearn.base import _column_names_from_X, _configure_processed_features from sgtlearn._features import ProcessedFeatures from sgtlearn._weights import normalize_sample_weight +from sgtlearn.base import _column_names_from_X, _configure_processed_features -def _n_samples_bootstrap( - n_samples: int, max_samples: Optional[Union[int, float]] -) -> int: +def _n_samples_bootstrap(n_samples: int, max_samples: float | None) -> int: if max_samples is None: return n_samples if isinstance(max_samples, Integral) and not isinstance(max_samples, bool): @@ -32,7 +31,7 @@ def _n_samples_bootstrap( m = float(max_samples) if not (0.0 < m <= 1.0): raise ValueError("max_samples as float must be in (0.0, 1.0].") - return max(1, int(round(m * n_samples))) + return max(1, round(m * n_samples)) def _parallel_fit_tree( @@ -42,10 +41,10 @@ def _parallel_fit_tree( n_bootstrap: int, X: np.ndarray, y: np.ndarray, - sample_weight: Optional[np.ndarray], + sample_weight: np.ndarray | None, tree_kw: dict[str, Any], tree_factory: Any, - processed_features: Optional[ProcessedFeatures], + processed_features: ProcessedFeatures | None, ) -> Any: """Fit one bootstrapped (or full) base tree; module-level for ``joblib`` workers.""" if bootstrap: @@ -85,8 +84,8 @@ def __init__( *, criterion: str, num_partitions: int = 2, - max_depth: Optional[int] = None, - max_leaf_nodes: Optional[int] = None, + max_depth: int | None = None, + max_leaf_nodes: int | None = None, min_samples_leaf: int = 1, min_impurity_decrease: float = 0.0, inner_max_depth: int = 3, @@ -96,13 +95,13 @@ def __init__( coordinate_descent_max_iters: int = 20, coordinate_descent_patience: int = 5, coordinate_descent_smart_init: bool = True, - max_features: Optional[Union[int, float, str]] = None, + max_features: float | str | None = None, bootstrap: bool = True, - max_samples: Optional[Union[int, float]] = None, - random_state: Optional[Union[int, np.random.RandomState]] = None, + max_samples: float | None = None, + random_state: int | np.random.RandomState | None = None, tao_n_runs: int = 10, tao_lambda: float = 0.0, - n_jobs: Optional[int] = None, + n_jobs: int | None = None, verbose: int = 0, ) -> None: self.n_estimators = int(n_estimators) @@ -129,24 +128,24 @@ def __init__( self.verbose = int(verbose) def _tree_kwargs(self) -> dict[str, Any]: - return dict( - criterion=self.criterion, - num_partitions=self.num_partitions, - max_depth=self.max_depth, - max_leaf_nodes=self.max_leaf_nodes, - min_samples_leaf=self.min_samples_leaf, - min_impurity_decrease=self.min_impurity_decrease, - inner_max_depth=self.inner_max_depth, - inner_max_leaf_nodes=self.inner_max_leaf_nodes, - inner_min_samples_leaf=self.inner_min_samples_leaf, - inner_min_impurity_decrease=self.inner_min_impurity_decrease, - coordinate_descent_max_iters=self.coordinate_descent_max_iters, - coordinate_descent_patience=self.coordinate_descent_patience, - coordinate_descent_smart_init=self.coordinate_descent_smart_init, - max_features=self.max_features, - tao_n_runs=self.tao_n_runs, - tao_lambda=self.tao_lambda, - ) + return { + "criterion": self.criterion, + "num_partitions": self.num_partitions, + "max_depth": self.max_depth, + "max_leaf_nodes": self.max_leaf_nodes, + "min_samples_leaf": self.min_samples_leaf, + "min_impurity_decrease": self.min_impurity_decrease, + "inner_max_depth": self.inner_max_depth, + "inner_max_leaf_nodes": self.inner_max_leaf_nodes, + "inner_min_samples_leaf": self.inner_min_samples_leaf, + "inner_min_impurity_decrease": self.inner_min_impurity_decrease, + "coordinate_descent_max_iters": self.coordinate_descent_max_iters, + "coordinate_descent_patience": self.coordinate_descent_patience, + "coordinate_descent_smart_init": self.coordinate_descent_smart_init, + "max_features": self.max_features, + "tao_n_runs": self.tao_n_runs, + "tao_lambda": self.tao_lambda, + } @abstractmethod def _check_X_y(self, X: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: @@ -159,9 +158,9 @@ def _make_tree(self, tree_seed: int, tree_kw: dict[str, Any]) -> Any: def _prepare_sample_weight( self, y: np.ndarray, - sample_weight: Optional[np.ndarray], + sample_weight: np.ndarray | None, n_samples: int, - ) -> Optional[np.ndarray]: + ) -> np.ndarray | None: """Return per-sample weights for tree fitting (subclasses may apply class weights).""" return normalize_sample_weight(sample_weight, n_samples) @@ -169,10 +168,10 @@ def fit( self, X: np.ndarray, y: np.ndarray, - sample_weight: Optional[np.ndarray] = None, + sample_weight: np.ndarray | None = None, *, - feature_dict: Optional[Mapping[int | str, Sequence[int | str]]] = None, - processed_features: Optional[ProcessedFeatures] = None, + feature_dict: Mapping[int | str, Sequence[int | str]] | None = None, + processed_features: ProcessedFeatures | None = None, ) -> RandomSGForest: """Fit the forest on ``X`` and targets ``y``. diff --git a/sgtlearn/ensemble/RandomSGForestClassifier.py b/sgtlearn/ensemble/random_sgforest_classifier.py similarity index 94% rename from sgtlearn/ensemble/RandomSGForestClassifier.py rename to sgtlearn/ensemble/random_sgforest_classifier.py index 075d344..67f9618 100644 --- a/sgtlearn/ensemble/RandomSGForestClassifier.py +++ b/sgtlearn/ensemble/random_sgforest_classifier.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any, Mapping, Optional, Sequence, Union +from collections.abc import Mapping, Sequence +from typing import Any import numpy as np from sklearn.base import ClassifierMixin @@ -114,8 +115,8 @@ def __init__( *, criterion: str = "gini", num_partitions: int = 2, - max_depth: Optional[int] = None, - max_leaf_nodes: Optional[int] = None, + max_depth: int | None = None, + max_leaf_nodes: int | None = None, min_samples_leaf: int = 1, min_impurity_decrease: float = 0.0, inner_max_depth: int = 3, @@ -125,16 +126,14 @@ def __init__( coordinate_descent_max_iters: int = 20, coordinate_descent_patience: int = 5, coordinate_descent_smart_init: bool = True, - max_features: Optional[Union[int, float, str]] = "sqrt", + max_features: float | str | None = "sqrt", bootstrap: bool = True, - max_samples: Optional[Union[int, float]] = None, - random_state: Optional[Union[int, np.random.RandomState]] = None, - class_weight: Optional[ - Union[Mapping[Any, float], Sequence[Mapping[Any, float]]] - ] = None, + max_samples: float | None = None, + random_state: int | np.random.RandomState | None = None, + class_weight: Mapping[Any, float] | Sequence[Mapping[Any, float]] | None = None, tao_n_runs: int = 10, tao_lambda: float = 0.0, - n_jobs: Optional[int] = None, + n_jobs: int | None = None, verbose: int = 0, ) -> None: self.class_weight = class_weight @@ -192,9 +191,9 @@ def _check_X_y(self, X: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarr def _prepare_sample_weight( self, y: np.ndarray, - sample_weight: Optional[np.ndarray], + sample_weight: np.ndarray | None, n_samples: int, - ) -> Optional[np.ndarray]: + ) -> np.ndarray | None: if self.class_weight is None: return super()._prepare_sample_weight(y, sample_weight, n_samples) return effective_sample_weight_classification( diff --git a/sgtlearn/ensemble/RandomSGForestRegressor.py b/sgtlearn/ensemble/random_sgforest_regressor.py similarity index 95% rename from sgtlearn/ensemble/RandomSGForestRegressor.py rename to sgtlearn/ensemble/random_sgforest_regressor.py index 14f61a0..8daabee 100644 --- a/sgtlearn/ensemble/RandomSGForestRegressor.py +++ b/sgtlearn/ensemble/random_sgforest_regressor.py @@ -2,15 +2,15 @@ from __future__ import annotations -from typing import Any, Optional, Union +from typing import Any import numpy as np from sklearn.base import RegressorMixin from sklearn.utils.validation import check_X_y +from sgtlearn._multioutput import squeeze_outputs from sgtlearn.base import SGTRegressor from sgtlearn.ensemble._random_sgforest import RandomSGForest -from sgtlearn._multioutput import squeeze_outputs class RandomSGForestRegressor(RegressorMixin, RandomSGForest): @@ -102,8 +102,8 @@ def __init__( *, criterion: str = "squared_error", num_partitions: int = 2, - max_depth: Optional[int] = None, - max_leaf_nodes: Optional[int] = None, + max_depth: int | None = None, + max_leaf_nodes: int | None = None, min_samples_leaf: int = 1, min_impurity_decrease: float = 0.0, inner_max_depth: int = 3, @@ -113,13 +113,13 @@ def __init__( coordinate_descent_max_iters: int = 20, coordinate_descent_patience: int = 5, coordinate_descent_smart_init: bool = True, - max_features: Optional[Union[int, float, str]] = "sqrt", + max_features: float | str | None = "sqrt", bootstrap: bool = True, - max_samples: Optional[Union[int, float]] = None, - random_state: Optional[Union[int, np.random.RandomState]] = None, + max_samples: float | None = None, + random_state: int | np.random.RandomState | None = None, tao_n_runs: int = 10, tao_lambda: float = 0.0, - n_jobs: Optional[int] = None, + n_jobs: int | None = None, verbose: int = 0, ) -> None: super().__init__( diff --git a/sgtlearn/tao.py b/sgtlearn/tao.py index bdfa875..84c6b77 100644 --- a/sgtlearn/tao.py +++ b/sgtlearn/tao.py @@ -6,7 +6,7 @@ from __future__ import annotations -from typing import Optional, TypeVar, Union +from typing import TypeVar import numpy as np from joblib import Parallel, delayed, effective_n_jobs @@ -26,12 +26,12 @@ ) from sgtlearn.base import BaseShapeCART, SGTClassifier, SGTRegressor from sgtlearn.ensemble._random_sgforest import RandomSGForest -from sgtlearn.ensemble.RandomSGForestClassifier import RandomSGForestClassifier -from sgtlearn.ensemble.RandomSGForestRegressor import RandomSGForestRegressor +from sgtlearn.ensemble.random_sgforest_classifier import RandomSGForestClassifier +from sgtlearn.ensemble.random_sgforest_regressor import RandomSGForestRegressor __all__ = ["TAO_refine"] -TaoModel = TypeVar("TaoModel", bound=Union[BaseShapeCART, RandomSGForest]) +TaoModel = TypeVar("TaoModel", bound=BaseShapeCART | RandomSGForest) def _tao_targets(model: TaoModel) -> list[SGTClassifier | SGTRegressor]: @@ -110,7 +110,7 @@ def _prepare_tao_arrays( model: TaoModel, X: np.ndarray, y: np.ndarray, - sample_weight: Optional[np.ndarray], + sample_weight: np.ndarray | None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Build ``(X32, y_native, sample_weights)`` shared by all trees in ``model``.""" X32 = np.ascontiguousarray(X, dtype=np.float32) @@ -157,11 +157,11 @@ def TAO_refine( X: np.ndarray, y: np.ndarray, *, - sample_weight: Optional[np.ndarray] = None, + sample_weight: np.ndarray | None = None, n_runs: int = 10, lambda_: float = 0.0, check_input: bool = True, - n_jobs: Optional[int] = None, + n_jobs: int | None = None, ) -> TaoModel: """Refine a fitted shape-generalized tree or forest in place with TAO.