Initialize relabeled_failures in SurvivalSplittingRule ctor (#1513)

erikcs · web-flow · commit a7cbe04b2c8b · 2025-10-03T05:43:13.000+10:00
diff --git a/core/src/splitting/SurvivalSplittingRule.cpp b/core/src/splitting/SurvivalSplittingRule.cpp
@@ -23,8 +23,8 @@
 
 namespace grf {
 
-SurvivalSplittingRule::SurvivalSplittingRule(double alpha):
-    alpha(alpha) {
+SurvivalSplittingRule::SurvivalSplittingRule(size_t num_data_rows, double alpha):
+    relabeled_failures(num_data_rows, 0), alpha(alpha) {
 }
 
 bool SurvivalSplittingRule::find_best_split(const Data& data,
@@ -99,9 +99,6 @@ void SurvivalSplittingRule::find_best_split_internal(const Data& data,
   std::vector<double> at_risk(num_failures + 1);
   at_risk[0] = static_cast<double>(size_node);
 
-  // allocating an N-sized (full data set size) array is faster than a hash table
-  std::vector<size_t> relabeled_failures(data.get_num_rows());
-
   std::vector<double> numerator_weights(num_failures + 1);
   std::vector<double> denominator_weights(num_failures + 1);
 
@@ -135,7 +132,7 @@ void SurvivalSplittingRule::find_best_split_internal(const Data& data,
 
   for (auto& var : possible_split_vars) {
     find_best_split_value(data, var, size_node, min_child_size, num_failures_node, num_failures,
-                          best_value, best_var, best_logrank, best_send_missing_left, samples, relabeled_failures,
+                          best_value, best_var, best_logrank, best_send_missing_left, samples,
                           count_failure, at_risk, numerator_weights, denominator_weights);
   }
 }
@@ -151,7 +148,6 @@ void SurvivalSplittingRule::find_best_split_value(const Data& data,
                                                   double& best_logrank,
                                                   bool& best_send_missing_left,
                                                   const std::vector<size_t>& samples,
-                                                  const std::vector<size_t>& relabeled_failures,
                                                   const std::vector<double>& count_failure,
                                                   const std::vector<double>& at_risk,
                                                   const std::vector<double>& numerator_weights,
diff --git a/core/src/splitting/SurvivalSplittingRule.h b/core/src/splitting/SurvivalSplittingRule.h
@@ -30,7 +30,7 @@ namespace grf {
 
 class SurvivalSplittingRule final: public SplittingRule {
 public:
-  SurvivalSplittingRule(double alpha);
+  SurvivalSplittingRule(size_t num_data_rows, double alpha);
 
   bool find_best_split(const Data& data,
                        size_t node,
@@ -66,7 +66,6 @@ class SurvivalSplittingRule final: public SplittingRule {
                              double& best_logrank,
                              bool& best_send_missing_left,
                              const std::vector<size_t>& samples,
-                             const std::vector<size_t>& relabeled_failures,
                              const std::vector<double>& count_failure,
                              const std::vector<double>& at_risk,
                              const std::vector<double>& numerator_weights,
@@ -81,6 +80,7 @@ class SurvivalSplittingRule final: public SplittingRule {
                                 const std::vector<double>& numerator_weights,
                                 const std::vector<double>& denominator_weights);
 
+  std::vector<size_t> relabeled_failures;
   double alpha;
 
   DISALLOW_COPY_AND_ASSIGN(SurvivalSplittingRule);
diff --git a/core/src/splitting/factory/CausalSurvivalSplittingRuleFactory.cpp b/core/src/splitting/factory/CausalSurvivalSplittingRuleFactory.cpp
@@ -23,7 +23,8 @@
 namespace grf {
 
 std::unique_ptr<SplittingRule> CausalSurvivalSplittingRuleFactory::create(size_t max_num_unique_values,
-                                                                        const TreeOptions& options) const {
+                                                                          const Data& data,
+                                                                          const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new CausalSurvivalSplittingRule(
       max_num_unique_values,
       options.get_min_node_size(),
diff --git a/core/src/splitting/factory/CausalSurvivalSplittingRuleFactory.h b/core/src/splitting/factory/CausalSurvivalSplittingRuleFactory.h
@@ -29,6 +29,7 @@ class CausalSurvivalSplittingRuleFactory final: public SplittingRuleFactory {
 public:
   CausalSurvivalSplittingRuleFactory() = default;
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   DISALLOW_COPY_AND_ASSIGN(CausalSurvivalSplittingRuleFactory);
diff --git a/core/src/splitting/factory/InstrumentalSplittingRuleFactory.cpp b/core/src/splitting/factory/InstrumentalSplittingRuleFactory.cpp
@@ -23,6 +23,7 @@
 namespace grf {
 
 std::unique_ptr<SplittingRule> InstrumentalSplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                        const Data& data,
                                                                         const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new InstrumentalSplittingRule(
       max_num_unique_values,
diff --git a/core/src/splitting/factory/InstrumentalSplittingRuleFactory.h b/core/src/splitting/factory/InstrumentalSplittingRuleFactory.h
@@ -37,6 +37,7 @@ class InstrumentalSplittingRuleFactory final: public SplittingRuleFactory {
 public:
   InstrumentalSplittingRuleFactory() = default;
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   DISALLOW_COPY_AND_ASSIGN(InstrumentalSplittingRuleFactory);
diff --git a/core/src/splitting/factory/MultiCausalSplittingRuleFactory.cpp b/core/src/splitting/factory/MultiCausalSplittingRuleFactory.cpp
@@ -28,6 +28,7 @@ MultiCausalSplittingRuleFactory::MultiCausalSplittingRuleFactory(size_t response
   num_treatments(num_treatments) {}
 
 std::unique_ptr<SplittingRule> MultiCausalSplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                       const Data& data,
                                                                        const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new MultiCausalSplittingRule(
       max_num_unique_values,
diff --git a/core/src/splitting/factory/MultiCausalSplittingRuleFactory.h b/core/src/splitting/factory/MultiCausalSplittingRuleFactory.h
@@ -30,6 +30,7 @@ class MultiCausalSplittingRuleFactory final: public SplittingRuleFactory {
                                   size_t num_treatments);
 
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   size_t response_length;
diff --git a/core/src/splitting/factory/MultiRegressionSplittingRuleFactory.cpp b/core/src/splitting/factory/MultiRegressionSplittingRuleFactory.cpp
@@ -26,6 +26,7 @@ MultiRegressionSplittingRuleFactory::MultiRegressionSplittingRuleFactory(size_t
   num_outcomes(num_outcomes) {}
 
 std::unique_ptr<SplittingRule> MultiRegressionSplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                           const Data& data,
                                                                            const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new MultiRegressionSplittingRule(
       max_num_unique_values,
diff --git a/core/src/splitting/factory/MultiRegressionSplittingRuleFactory.h b/core/src/splitting/factory/MultiRegressionSplittingRuleFactory.h
@@ -36,6 +36,7 @@ class MultiRegressionSplittingRuleFactory final: public SplittingRuleFactory {
   MultiRegressionSplittingRuleFactory(size_t num_outcomes);
 
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   size_t num_outcomes;
diff --git a/core/src/splitting/factory/ProbabilitySplittingRuleFactory.cpp b/core/src/splitting/factory/ProbabilitySplittingRuleFactory.cpp
@@ -27,6 +27,7 @@ ProbabilitySplittingRuleFactory::ProbabilitySplittingRuleFactory(size_t num_clas
     num_classes(num_classes) {}
 
 std::unique_ptr<SplittingRule> ProbabilitySplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                       const Data& data,
                                                                        const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new ProbabilitySplittingRule(
       max_num_unique_values,
diff --git a/core/src/splitting/factory/ProbabilitySplittingRuleFactory.h b/core/src/splitting/factory/ProbabilitySplittingRuleFactory.h
@@ -37,6 +37,7 @@ class ProbabilitySplittingRuleFactory final: public SplittingRuleFactory {
 public:
   ProbabilitySplittingRuleFactory(size_t num_classes);
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 
 private:
diff --git a/core/src/splitting/factory/RegressionSplittingRuleFactory.cpp b/core/src/splitting/factory/RegressionSplittingRuleFactory.cpp
@@ -23,6 +23,7 @@
 namespace grf {
 
 std::unique_ptr<SplittingRule> RegressionSplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                      const Data& data,
                                                                       const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new RegressionSplittingRule(
       max_num_unique_values,
diff --git a/core/src/splitting/factory/RegressionSplittingRuleFactory.h b/core/src/splitting/factory/RegressionSplittingRuleFactory.h
@@ -35,6 +35,7 @@ class RegressionSplittingRuleFactory final: public SplittingRuleFactory {
 public:
   RegressionSplittingRuleFactory() = default;
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   DISALLOW_COPY_AND_ASSIGN(RegressionSplittingRuleFactory);
diff --git a/core/src/splitting/factory/SplittingRuleFactory.h b/core/src/splitting/factory/SplittingRuleFactory.h
@@ -34,6 +34,7 @@ class SplittingRuleFactory {
   virtual ~SplittingRuleFactory() = default;
 
   virtual std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                                const Data& data,
                                                 const TreeOptions& options) const = 0;
 };
 
diff --git a/core/src/splitting/factory/SurvivalSplittingRuleFactory.cpp b/core/src/splitting/factory/SurvivalSplittingRuleFactory.cpp
@@ -23,9 +23,10 @@
 namespace grf {
 
 std::unique_ptr<SplittingRule> SurvivalSplittingRuleFactory::create(size_t max_num_unique_values,
+                                                                    const Data& data,
                                                                     const TreeOptions& options) const {
   return std::unique_ptr<SplittingRule>(new SurvivalSplittingRule(
-      options.get_alpha()));
+      data.get_num_rows(), options.get_alpha()));
 }
 
 } // namespace grf
diff --git a/core/src/splitting/factory/SurvivalSplittingRuleFactory.h b/core/src/splitting/factory/SurvivalSplittingRuleFactory.h
@@ -33,6 +33,7 @@ class SurvivalSplittingRuleFactory final: public SplittingRuleFactory {
   SurvivalSplittingRuleFactory() = default;
 
   std::unique_ptr<SplittingRule> create(size_t max_num_unique_values,
+                                        const Data& data,
                                         const TreeOptions& options) const;
 private:
   DISALLOW_COPY_AND_ASSIGN(SurvivalSplittingRuleFactory);
diff --git a/core/src/tree/TreeTrainer.cpp b/core/src/tree/TreeTrainer.cpp
@@ -61,7 +61,7 @@ std::unique_ptr<Tree> TreeTrainer::train(const Data& data,
 
   // nodes[0].size() is the number of samples subsampled for this tree.
   std::unique_ptr<SplittingRule> splitting_rule = splitting_rule_factory->create(
-      nodes[0].size(), options);
+      nodes[0].size(), data, options);
 
   size_t num_open_nodes = 1;
   size_t i = 0;
diff --git a/core/test/splitting/SplittingRuleInvarianceTest.cpp b/core/test/splitting/SplittingRuleInvarianceTest.cpp
@@ -41,7 +41,7 @@ void run_one_split(const Data& data,
                    size_t num_features,
                    size_t& split_var,
                    double& split_value) {
-  std::unique_ptr<SplittingRule> splitting_rule = splitting_rule_factory->create(data.get_num_rows(), options);
+  std::unique_ptr<SplittingRule> splitting_rule = splitting_rule_factory->create(data.get_num_rows(), data, options);
   std::vector<size_t> possible_split_vars(num_features - 1);
   // Fill with {0, 1, 2, ..., Xj}
   std::iota(possible_split_vars.begin(), possible_split_vars.end(), 0);
diff --git a/core/test/splitting/SurvivalSplittingRuleTest.cpp b/core/test/splitting/SurvivalSplittingRuleTest.cpp
@@ -77,7 +77,7 @@ TEST_CASE("survival splitting logrank calculation is correct", "[survival], [spl
   TreeOptions options = ForestTestUtilities::default_options().get_tree_options();
 
   std::unique_ptr<RelabelingStrategy> relabeling_strategy(new NoopRelabelingStrategy());
-  std::unique_ptr<SurvivalSplittingRule> surv_splitting_rule(new SurvivalSplittingRule(options.get_alpha()));
+  std::unique_ptr<SurvivalSplittingRule> surv_splitting_rule(new SurvivalSplittingRule(data.get_num_rows(), options.get_alpha()));
 
   std::vector<double> logranks = run_splits(data, options, surv_splitting_rule, relabeling_strategy, num_features);
 
