Median branch assignment

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

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 <cstddef>
 #include "AbsoluteErrorBranchAssignment.h"
 
+#include "AbsoluteErrorBranchAssignmentCommon.h"
 #include "Criterion.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <stdexcept>
 
 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<std::vector<float>>(nOutputs_));
+  partWs_.assign(numPartitions, std::vector<std::vector<float>>(nOutputs_));
+  partSrcBin_.assign(numPartitions,
+                     std::vector<std::vector<size_t>>(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<size_t> 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<float> outY;
+    std::vector<float> outW;
+    std::vector<size_t> 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<float> outY;
+    std::vector<float> outW;
+    std::vector<size_t> 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<double>(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<float> &ys,
-    std::vector<float> &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<float> ys;
-  std::vector<float> 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<size_t> 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<float> ysSorted;
-    std::vector<float> 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;
 }