Ignore nan trees in regression random forest

src/Forest.cpp

@@ -10,6 +10,7 @@
  #-------------------------------------------------------------------------------*/
 
 #include <math.h>
+#include <cmath>
 #include <algorithm>
 #include <stdexcept>
 #include <string>
@@ -615,18 +616,25 @@ void Forest::computePermutationImportance() {
   std::vector<std::vector<double>> variance_threads(num_threads);
   std::vector<std::vector<double>> variable_importance_casewise_threads(num_threads);
 
+  // Count trees with NaN importance so they can be dropped from the average
+  std::vector<std::vector<double>> variable_importance_nan_threads(num_threads);
+  std::vector<std::vector<double>> variable_importance_casewise_nan_threads(num_threads);
+
   // Compute importance
   for (uint i = 0; i < num_threads; ++i) {
     variable_importance_threads[i].resize(num_independent_variables, 0);
+    variable_importance_nan_threads[i].resize(num_independent_variables, 0);
     if (importance_mode == IMP_PERM_BREIMAN || importance_mode == IMP_PERM_LIAW) {
       variance_threads[i].resize(num_independent_variables, 0);
     }
     if (importance_mode == IMP_PERM_CASEWISE) {
       variable_importance_casewise_threads[i].resize(num_independent_variables * num_samples, 0);
+      variable_importance_casewise_nan_threads[i].resize(num_independent_variables * num_samples, 0);
     }
     threads.emplace_back(&Forest::computeTreePermutationImportanceInThread, this, i,
         std::ref(variable_importance_threads[i]), std::ref(variance_threads[i]),
-        std::ref(variable_importance_casewise_threads[i]));
+        std::ref(variable_importance_casewise_threads[i]), std::ref(variable_importance_nan_threads[i]),
+        std::ref(variable_importance_casewise_nan_threads[i]));
   }
   showProgress("Computing permutation importance..", num_trees);
   for (auto &thread : threads) {
@@ -648,6 +656,15 @@ void Forest::computePermutationImportance() {
   }
   variable_importance_threads.clear();
 
+  // Sum thread NaN counts
+  std::vector<double> variable_importance_nan(num_independent_variables, 0);
+  for (size_t i = 0; i < num_independent_variables; ++i) {
+    for (uint j = 0; j < num_threads; ++j) {
+      variable_importance_nan[i] += variable_importance_nan_threads[j][i];
+    }
+  }
+  variable_importance_nan_threads.clear();
+
   // Sum thread variances
   std::vector<double> variance(num_independent_variables, 0);
   if (importance_mode == IMP_PERM_BREIMAN || importance_mode == IMP_PERM_LIAW) {
@@ -660,6 +677,7 @@ void Forest::computePermutationImportance() {
   }
 
   // Sum thread casewise importances
+  std::vector<double> variable_importance_casewise_nan;
   if (importance_mode == IMP_PERM_CASEWISE) {
     variable_importance_casewise.resize(num_independent_variables * num_samples, 0);
     for (size_t i = 0; i < variable_importance_casewise.size(); ++i) {
@@ -668,23 +686,37 @@ void Forest::computePermutationImportance() {
       }
     }
     variable_importance_casewise_threads.clear();
+
+    // Sum thread casewise NaN counts
+    variable_importance_casewise_nan.resize(num_independent_variables * num_samples, 0);
+    for (size_t i = 0; i < variable_importance_casewise_nan.size(); ++i) {
+      for (uint j = 0; j < num_threads; ++j) {
+        variable_importance_casewise_nan[i] += variable_importance_casewise_nan_threads[j][i];
+      }
+    }
+    variable_importance_casewise_nan_threads.clear();
   }
 
   for (size_t i = 0; i < variable_importance.size(); ++i) {
-    variable_importance[i] /= num_trees;
+    // Mean over trees, ignoring trees with NaN importance
+    double num_valid = (double) num_trees - variable_importance_nan[i];
+    // If no valid trees, num_valid == 0 yields NaN (0.0 / 0), as intended
+    variable_importance[i] /= num_valid;
 
     // Normalize by variance for scaled permutation importance
     if (importance_mode == IMP_PERM_BREIMAN || importance_mode == IMP_PERM_LIAW) {
       if (variance[i] != 0) {
-        variance[i] = variance[i] / num_trees - variable_importance[i] * variable_importance[i];
-        variable_importance[i] /= sqrt(variance[i] / num_trees);
+        variance[i] = variance[i] / num_valid - variable_importance[i] * variable_importance[i];
+        variable_importance[i] /= sqrt(variance[i] / num_valid);
       }
     }
   }
 
   if (importance_mode == IMP_PERM_CASEWISE) {
     for (size_t i = 0; i < variable_importance_casewise.size(); ++i) {
-      variable_importance_casewise[i] /= num_trees;
+      // Mean over trees, ignoring trees with NaN importance
+      double num_valid = (double) num_trees - variable_importance_casewise_nan[i];
+      variable_importance_casewise[i] /= num_valid;
     }
   }
 }
@@ -763,10 +795,12 @@ void Forest::predictInternalInThread(uint thread_idx) {
 }
 
 void Forest::computeTreePermutationImportanceInThread(uint thread_idx, std::vector<double>& importance,
-    std::vector<double>& variance, std::vector<double>& importance_casewise) {
+    std::vector<double>& variance, std::vector<double>& importance_casewise,
+    std::vector<double>& importance_nan, std::vector<double>& importance_casewise_nan) {
   if (thread_ranges.size() > thread_idx + 1) {
     for (size_t i = thread_ranges[thread_idx]; i < thread_ranges[thread_idx + 1]; ++i) {
-      trees[i]->computePermutationImportance(importance, variance, importance_casewise);
+      trees[i]->computePermutationImportance(importance, variance, importance_casewise, importance_nan,
+          importance_casewise_nan);
 
       // Check for user interrupt
 #ifdef R_BUILD

src/Forest.h

@@ -189,7 +189,8 @@ class Forest {
   void predictTreesInThread(uint thread_idx, const Data* prediction_data, bool oob_prediction);
   void predictInternalInThread(uint thread_idx);
   void computeTreePermutationImportanceInThread(uint thread_idx, std::vector<double>& importance,
-      std::vector<double>& variance, std::vector<double>& importance_casewise);
+      std::vector<double>& variance, std::vector<double>& importance_casewise,
+      std::vector<double>& importance_nan, std::vector<double>& importance_casewise_nan);
 
   // Load forest from file
   void loadFromFile(std::string filename);

src/ForestRegression.cpp

@@ -10,6 +10,7 @@
  #-------------------------------------------------------------------------------*/
 
 #include <algorithm>
+#include <cmath>
 #include <stdexcept>
 #include <string>
 
@@ -117,12 +118,18 @@ void ForestRegression::predictInternal(size_t sample_idx) {
       }
     }
   } else {
-    // Mean over trees
+    // Mean over trees, ignoring NaN tree predictions
     double prediction_sum = 0;
+    size_t num_valid = 0;
     for (size_t tree_idx = 0; tree_idx < num_trees; ++tree_idx) {
-      prediction_sum += getTreePrediction(tree_idx, sample_idx);
+      double value = getTreePrediction(tree_idx, sample_idx);
+      if (!std::isnan(value)) {
+        prediction_sum += value;
+        ++num_valid;
+      }
     }
-    predictions[0][0][sample_idx] = prediction_sum / num_trees;
+    // If no valid predictions, num_valid == 0 yields NaN (0.0 / 0), as intended
+    predictions[0][0][sample_idx] = prediction_sum / (double) num_valid;
   }
 }
 
@@ -138,8 +145,11 @@ void ForestRegression::computePredictionErrorInternal() {
       size_t sampleID = trees[tree_idx]->getOobSampleIDs()[sample_idx];
       double value = getTreePrediction(tree_idx, sample_idx);
 
-      predictions[0][0][sampleID] += value;
-      ++samples_oob_count[sampleID];
+      // Ignore NaN tree predictions so they do not poison the OOB mean/MSE
+      if (!std::isnan(value)) {
+        predictions[0][0][sampleID] += value;
+        ++samples_oob_count[sampleID];
+      }
     }
   }
 

src/Tree.cpp

@@ -9,6 +9,7 @@
  R package "ranger" under GPL3 license.
  #-------------------------------------------------------------------------------*/
 
+#include <cmath>
 #include <iterator>
 
 #include "Tree.h"
@@ -223,7 +224,8 @@ void Tree::predict(const Data* prediction_data, bool oob_prediction) {
 }
 
 void Tree::computePermutationImportance(std::vector<double>& forest_importance, std::vector<double>& forest_variance,
-    std::vector<double>& forest_importance_casewise) {
+    std::vector<double>& forest_importance_casewise, std::vector<double>& forest_importance_nan,
+    std::vector<double>& forest_importance_casewise_nan) {
 
   size_t num_independent_variables = data->getNumCols();
 
@@ -267,20 +269,31 @@ void Tree::computePermutationImportance(std::vector<double>& forest_importance,
         accuracy_permuted = computePredictionAccuracyInternal(&prederr_shuf_casewise);
         for (size_t j = 0; j < num_samples_oob; ++j) {
           size_t pos = i * num_samples + oob_sampleIDs[j];
-          forest_importance_casewise[pos] += prederr_shuf_casewise[j] - prederr_normal_casewise[j];
+          // Ignore NaN casewise differences so they do not poison the average
+          double value = prederr_shuf_casewise[j] - prederr_normal_casewise[j];
+          if (!std::isnan(value)) {
+            forest_importance_casewise[pos] += value;
+          } else {
+            ++forest_importance_casewise_nan[pos];
+          }
         }
       } else {
         accuracy_permuted = computePredictionAccuracyInternal(NULL);
       }
 
       double accuracy_difference = accuracy_normal - accuracy_permuted;
-      forest_importance[i] += accuracy_difference;
-
-      // Compute variance
-      if (importance_mode == IMP_PERM_BREIMAN) {
-        forest_variance[i] += accuracy_difference * accuracy_difference;
-      } else if (importance_mode == IMP_PERM_LIAW) {
-        forest_variance[i] += accuracy_difference * accuracy_difference * num_samples_oob;
+      // Ignore NaN importance so it does not poison the average across trees
+      if (!std::isnan(accuracy_difference)) {
+        forest_importance[i] += accuracy_difference;
+
+        // Compute variance
+        if (importance_mode == IMP_PERM_BREIMAN) {
+          forest_variance[i] += accuracy_difference * accuracy_difference;
+        } else if (importance_mode == IMP_PERM_LIAW) {
+          forest_variance[i] += accuracy_difference * accuracy_difference * num_samples_oob;
+        }
+      } else {
+        ++forest_importance_nan[i];
       }
     }
   }

src/Tree.h

@@ -50,7 +50,8 @@ class Tree {
   void predict(const Data* prediction_data, bool oob_prediction);
 
   void computePermutationImportance(std::vector<double>& forest_importance, std::vector<double>& forest_variance,
-      std::vector<double>& forest_importance_casewise);
+      std::vector<double>& forest_importance_casewise, std::vector<double>& forest_importance_nan,
+      std::vector<double>& forest_importance_casewise_nan);
 
   void appendToFile(std::ofstream& file);
   virtual void appendToFileInternal(std::ofstream& file) = 0;