leaf_model.h

 
#ifndef STOCHTREE_LEAF_MODEL_H_
#define STOCHTREE_LEAF_MODEL_H_
 
#include <Eigen/Dense>
#include <stochtree/cutpoint_candidates.h>
#include <stochtree/data.h>
#include <stochtree/gamma_sampler.h>
#include <stochtree/ig_sampler.h>
#include <stochtree/log.h>
#include <stochtree/meta.h>
#include <stochtree/normal_sampler.h>
#include <stochtree/openmp_utils.h>
#include <stochtree/partition_tracker.h>
#include <stochtree/prior.h>
#include <stochtree/tree.h>
 
#include <random>
#include <variant>
 
namespace StochTree {
 
enum ModelType {
  kConstantLeafGaussian,
  kUnivariateRegressionLeafGaussian,
  kMultivariateRegressionLeafGaussian,
  kLogLinearVariance,
  kCloglogOrdinal
};
 
class GaussianConstantSuffStat {
 public:
  data_size_t n;
  double sum_w;
  double sum_yw;
  GaussianConstantSuffStat() {
    n = 0;
    sum_w = 0.0;
    sum_yw = 0.0;
  }
  void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) {
    n += 1;
    if (dataset.HasVarWeights()) {
      sum_w += 1 / dataset.VarWeightValue(row_idx);
      sum_yw += outcome(row_idx, 0) / dataset.VarWeightValue(row_idx);
    } else {
      sum_w += 1.0;
      sum_yw += outcome(row_idx, 0);
    }
  }
  void ResetSuffStat() {
    n = 0;
    sum_w = 0.0;
    sum_yw = 0.0;
  }
  void AddSuffStatInplace(GaussianConstantSuffStat& suff_stat) {
    n += suff_stat.n;
    sum_w += suff_stat.sum_w;
    sum_yw += suff_stat.sum_yw;
  }
  void AddSuffStat(GaussianConstantSuffStat& lhs, GaussianConstantSuffStat& rhs) {
    n = lhs.n + rhs.n;
    sum_w = lhs.sum_w + rhs.sum_w;
    sum_yw = lhs.sum_yw + rhs.sum_yw;
  }
  void SubtractSuffStat(GaussianConstantSuffStat& lhs, GaussianConstantSuffStat& rhs) {
    n = lhs.n - rhs.n;
    sum_w = lhs.sum_w - rhs.sum_w;
    sum_yw = lhs.sum_yw - rhs.sum_yw;
  }
  bool SampleGreaterThan(data_size_t threshold) {
    return n > threshold;
  }
  bool SampleGreaterThanEqual(data_size_t threshold) {
    return n >= threshold;
  }
  data_size_t SampleSize() {
    return n;
  }
};
 
class GaussianConstantLeafModel {
 public:
  GaussianConstantLeafModel(double tau) {
    tau_ = tau;
    normal_sampler_ = UnivariateNormalSampler();
  }
  ~GaussianConstantLeafModel() {}
  double SplitLogMarginalLikelihood(GaussianConstantSuffStat& left_stat, GaussianConstantSuffStat& right_stat, double global_variance);
  double NoSplitLogMarginalLikelihood(GaussianConstantSuffStat& suff_stat, double global_variance);
  double PosteriorParameterMean(GaussianConstantSuffStat& suff_stat, double global_variance);
  double PosteriorParameterVariance(GaussianConstantSuffStat& suff_stat, double global_variance);
  void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen);
  void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value);
  void SetScale(double tau) { tau_ = tau; }
  inline bool RequiresBasis() { return false; }
 
 private:
  double tau_;
  UnivariateNormalSampler normal_sampler_;
};
 
class GaussianUnivariateRegressionSuffStat {
 public:
  data_size_t n;
  double sum_xxw;
  double sum_yxw;
  GaussianUnivariateRegressionSuffStat() {
    n = 0;
    sum_xxw = 0.0;
    sum_yxw = 0.0;
  }
  void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) {
    n += 1;
    if (dataset.HasVarWeights()) {
      sum_xxw += dataset.BasisValue(row_idx, 0) * dataset.BasisValue(row_idx, 0) / dataset.VarWeightValue(row_idx);
      sum_yxw += outcome(row_idx, 0) * dataset.BasisValue(row_idx, 0) / dataset.VarWeightValue(row_idx);
    } else {
      sum_xxw += dataset.BasisValue(row_idx, 0) * dataset.BasisValue(row_idx, 0);
      sum_yxw += outcome(row_idx, 0) * dataset.BasisValue(row_idx, 0);
    }
  }
  void ResetSuffStat() {
    n = 0;
    sum_xxw = 0.0;
    sum_yxw = 0.0;
  }
  void AddSuffStatInplace(GaussianUnivariateRegressionSuffStat& suff_stat) {
    n += suff_stat.n;
    sum_xxw += suff_stat.sum_xxw;
    sum_yxw += suff_stat.sum_yxw;
  }
  void AddSuffStat(GaussianUnivariateRegressionSuffStat& lhs, GaussianUnivariateRegressionSuffStat& rhs) {
    n = lhs.n + rhs.n;
    sum_xxw = lhs.sum_xxw + rhs.sum_xxw;
    sum_yxw = lhs.sum_yxw + rhs.sum_yxw;
  }
  void SubtractSuffStat(GaussianUnivariateRegressionSuffStat& lhs, GaussianUnivariateRegressionSuffStat& rhs) {
    n = lhs.n - rhs.n;
    sum_xxw = lhs.sum_xxw - rhs.sum_xxw;
    sum_yxw = lhs.sum_yxw - rhs.sum_yxw;
  }
  bool SampleGreaterThan(data_size_t threshold) {
    return n > threshold;
  }
  bool SampleGreaterThanEqual(data_size_t threshold) {
    return n >= threshold;
  }
  data_size_t SampleSize() {
    return n;
  }
};
 
class GaussianUnivariateRegressionLeafModel {
 public:
  GaussianUnivariateRegressionLeafModel(double tau) {
    tau_ = tau;
    normal_sampler_ = UnivariateNormalSampler();
  }
  ~GaussianUnivariateRegressionLeafModel() {}
  double SplitLogMarginalLikelihood(GaussianUnivariateRegressionSuffStat& left_stat, GaussianUnivariateRegressionSuffStat& right_stat, double global_variance);
  double NoSplitLogMarginalLikelihood(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance);
  double PosteriorParameterMean(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance);
  double PosteriorParameterVariance(GaussianUnivariateRegressionSuffStat& suff_stat, double global_variance);
  void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen);
  void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value);
  void SetScale(double tau) { tau_ = tau; }
  inline bool RequiresBasis() { return true; }
 
 private:
  double tau_;
  UnivariateNormalSampler normal_sampler_;
};
 
class GaussianMultivariateRegressionSuffStat {
 public:
  data_size_t n;
  int p;
  Eigen::MatrixXd XtWX;
  Eigen::MatrixXd ytWX;
  GaussianMultivariateRegressionSuffStat(int basis_dim) {
    n = 0;
    XtWX = Eigen::MatrixXd::Zero(basis_dim, basis_dim);
    ytWX = Eigen::MatrixXd::Zero(1, basis_dim);
    p = basis_dim;
  }
  void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) {
    n += 1;
    if (dataset.HasVarWeights()) {
      for (int i = 0; i < p; i++) {
        ytWX(0, i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i) / dataset.VarWeightValue(row_idx);
        for (int j = 0; j < p; j++) {
          XtWX(i, j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j) / dataset.VarWeightValue(row_idx);
        }
      }
    } else {
      for (int i = 0; i < p; i++) {
        ytWX(0, i) += outcome(row_idx, 0) * dataset.BasisValue(row_idx, i);
        for (int j = 0; j < p; j++) {
          XtWX(i, j) += dataset.BasisValue(row_idx, i) * dataset.BasisValue(row_idx, j);
        }
      }
    }
  }
  void ResetSuffStat() {
    n = 0;
    for (int i = 0; i < p; i++) {
      ytWX(0, i) = 0.0;
      for (int j = 0; j < p; j++) {
        XtWX(i, j) = 0.0;
      }
    }
  }
  void AddSuffStatInplace(GaussianMultivariateRegressionSuffStat& suff_stat) {
    n += suff_stat.n;
    XtWX += suff_stat.XtWX;
    ytWX += suff_stat.ytWX;
  }
  void AddSuffStat(GaussianMultivariateRegressionSuffStat& lhs, GaussianMultivariateRegressionSuffStat& rhs) {
    n = lhs.n + rhs.n;
    XtWX = lhs.XtWX + rhs.XtWX;
    ytWX = lhs.ytWX + rhs.ytWX;
  }
  void SubtractSuffStat(GaussianMultivariateRegressionSuffStat& lhs, GaussianMultivariateRegressionSuffStat& rhs) {
    n = lhs.n - rhs.n;
    XtWX = lhs.XtWX - rhs.XtWX;
    ytWX = lhs.ytWX - rhs.ytWX;
  }
  bool SampleGreaterThan(data_size_t threshold) {
    return n > threshold;
  }
  bool SampleGreaterThanEqual(data_size_t threshold) {
    return n >= threshold;
  }
  data_size_t SampleSize() {
    return n;
  }
};
 
class GaussianMultivariateRegressionLeafModel {
 public:
  GaussianMultivariateRegressionLeafModel(Eigen::MatrixXd& Sigma_0) {
    Sigma_0_ = Sigma_0;
    multivariate_normal_sampler_ = MultivariateNormalSampler();
  }
  ~GaussianMultivariateRegressionLeafModel() {}
  double SplitLogMarginalLikelihood(GaussianMultivariateRegressionSuffStat& left_stat, GaussianMultivariateRegressionSuffStat& right_stat, double global_variance);
  double NoSplitLogMarginalLikelihood(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance);
  Eigen::VectorXd PosteriorParameterMean(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance);
  Eigen::MatrixXd PosteriorParameterVariance(GaussianMultivariateRegressionSuffStat& suff_stat, double global_variance);
  void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen);
  void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value);
  void SetScale(Eigen::MatrixXd& Sigma_0) { Sigma_0_ = Sigma_0; }
  inline bool RequiresBasis() { return true; }
 
 private:
  Eigen::MatrixXd Sigma_0_;
  MultivariateNormalSampler multivariate_normal_sampler_;
};
 
class LogLinearVarianceSuffStat {
 public:
  data_size_t n;
  double weighted_sum_ei;
  LogLinearVarianceSuffStat() {
    n = 0;
    weighted_sum_ei = 0.0;
  }
  void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) {
    n += 1;
    weighted_sum_ei += std::exp(std::log(outcome(row_idx) * outcome(row_idx)) - tracker.GetSamplePrediction(row_idx) + tracker.GetTreeSamplePrediction(row_idx, tree_idx));
  }
  void ResetSuffStat() {
    n = 0;
    weighted_sum_ei = 0.0;
  }
  void AddSuffStatInplace(LogLinearVarianceSuffStat& suff_stat) {
    n += suff_stat.n;
    weighted_sum_ei += suff_stat.weighted_sum_ei;
  }
  void AddSuffStat(LogLinearVarianceSuffStat& lhs, LogLinearVarianceSuffStat& rhs) {
    n = lhs.n + rhs.n;
    weighted_sum_ei = lhs.weighted_sum_ei + rhs.weighted_sum_ei;
  }
  void SubtractSuffStat(LogLinearVarianceSuffStat& lhs, LogLinearVarianceSuffStat& rhs) {
    n = lhs.n - rhs.n;
    weighted_sum_ei = lhs.weighted_sum_ei - rhs.weighted_sum_ei;
  }
  bool SampleGreaterThan(data_size_t threshold) {
    return n > threshold;
  }
  bool SampleGreaterThanEqual(data_size_t threshold) {
    return n >= threshold;
  }
  data_size_t SampleSize() {
    return n;
  }
};
 
class LogLinearVarianceLeafModel {
 public:
  LogLinearVarianceLeafModel(double a, double b) {
    a_ = a;
    b_ = b;
    gamma_sampler_ = GammaSampler();
  }
  ~LogLinearVarianceLeafModel() {}
  double SplitLogMarginalLikelihood(LogLinearVarianceSuffStat& left_stat, LogLinearVarianceSuffStat& right_stat, double global_variance);
  double NoSplitLogMarginalLikelihood(LogLinearVarianceSuffStat& suff_stat, double global_variance);
  double SuffStatLogMarginalLikelihood(LogLinearVarianceSuffStat& suff_stat, double global_variance);
  double PosteriorParameterShape(LogLinearVarianceSuffStat& suff_stat, double global_variance);
  double PosteriorParameterScale(LogLinearVarianceSuffStat& suff_stat, double global_variance);
  void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen);
  void SetEnsembleRootPredictedValue(ForestDataset& dataset, TreeEnsemble* ensemble, double root_pred_value);
  void SetPriorShape(double a) { a_ = a; }
  void SetPriorRate(double b) { b_ = b; }
  inline bool RequiresBasis() { return false; }
 
 private:
  double a_;
  double b_;
  GammaSampler gamma_sampler_;
};
 
class CloglogOrdinalSuffStat {
 public:
  data_size_t n;
  double sum_Y_less_K;
  double other_sum;
 
  CloglogOrdinalSuffStat() {
    n = 0;
    sum_Y_less_K = 0.0;
    other_sum = 0.0;
  }
 
  void IncrementSuffStat(ForestDataset& dataset, Eigen::VectorXd& outcome, ForestTracker& tracker, data_size_t row_idx, int tree_idx) {
    n += 1;
 
    // Get ordinal outcome value for this observation
    unsigned int y = static_cast<unsigned int>(outcome(row_idx));
 
    // Get auxiliary data from tracker (assuming types: 0=latents Z, 1=forest predictions, 2=cutpoints gamma, 3=cumsum exp of gamma)
    double Z = dataset.GetAuxiliaryDataValue(0, row_idx);             // latent variables Z
    double lambda_minus = dataset.GetAuxiliaryDataValue(1, row_idx);  // forest predictions excluding current tree
 
    // Get cutpoints gamma and cumulative sum of exp(gamma)
    const std::vector<double>& gamma = dataset.GetAuxiliaryDataVectorConst(2);  // cutpoints gamma
    const std::vector<double>& seg = dataset.GetAuxiliaryDataVectorConst(3);    // cumsum exp of gamma
 
    int K = gamma.size() + 1;  // Number of ordinal categories
 
    if (y == K - 1) {
      other_sum += std::exp(lambda_minus) * seg[y];  // checked and it's correct
    } else {
      sum_Y_less_K += 1.0;
      other_sum += std::exp(lambda_minus) * (Z * std::exp(gamma[y]) + seg[y]);  // checked and it's correct
    }
  }
 
  void ResetSuffStat() {
    n = 0;
    sum_Y_less_K = 0.0;
    other_sum = 0.0;
  }
 
  void AddSuffStatInplace(CloglogOrdinalSuffStat& suff_stat) {
    n += suff_stat.n;
    sum_Y_less_K += suff_stat.sum_Y_less_K;
    other_sum += suff_stat.other_sum;
  }
 
  void AddSuffStat(CloglogOrdinalSuffStat& lhs, CloglogOrdinalSuffStat& rhs) {
    n = lhs.n + rhs.n;
    sum_Y_less_K = lhs.sum_Y_less_K + rhs.sum_Y_less_K;
    other_sum = lhs.other_sum + rhs.other_sum;
  }
 
  void SubtractSuffStat(CloglogOrdinalSuffStat& lhs, CloglogOrdinalSuffStat& rhs) {
    n = lhs.n - rhs.n;
    sum_Y_less_K = lhs.sum_Y_less_K - rhs.sum_Y_less_K;
    other_sum = lhs.other_sum - rhs.other_sum;
  }
 
  bool SampleGreaterThan(data_size_t threshold) {
    return n > threshold;
  }
 
  bool SampleGreaterThanEqual(data_size_t threshold) {
    return n >= threshold;
  }
 
  data_size_t SampleSize() {
    return n;
  }
};
 
class CloglogOrdinalLeafModel {
 public:
  CloglogOrdinalLeafModel(double a, double b) {
    a_ = a;
    b_ = b;
    gamma_sampler_ = GammaSampler();
  }
  ~CloglogOrdinalLeafModel() {}
 
  double SplitLogMarginalLikelihood(CloglogOrdinalSuffStat& left_stat, CloglogOrdinalSuffStat& right_stat, double global_variance);
 
  double NoSplitLogMarginalLikelihood(CloglogOrdinalSuffStat& suff_stat, double global_variance);
 
  double SuffStatLogMarginalLikelihood(CloglogOrdinalSuffStat& suff_stat, double global_variance);
 
  double PosteriorParameterShape(CloglogOrdinalSuffStat& suff_stat, double global_variance);
 
  double PosteriorParameterRate(CloglogOrdinalSuffStat& suff_stat, double global_variance);
 
  void SampleLeafParameters(ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, Tree* tree, int tree_num, double global_variance, std::mt19937& gen);
  inline bool RequiresBasis() { return false; }
 
 private:
  double a_;
  double b_;
  GammaSampler gamma_sampler_;
};
 
using SuffStatVariant = std::variant<GaussianConstantSuffStat,
                                     GaussianUnivariateRegressionSuffStat,
                                     GaussianMultivariateRegressionSuffStat,
                                     LogLinearVarianceSuffStat,
                                     CloglogOrdinalSuffStat>;
 
using LeafModelVariant = std::variant<GaussianConstantLeafModel,
                                      GaussianUnivariateRegressionLeafModel,
                                      GaussianMultivariateRegressionLeafModel,
                                      LogLinearVarianceLeafModel,
                                      CloglogOrdinalLeafModel>;
 
template <typename SuffStatType, typename... SuffStatConstructorArgs>
static inline SuffStatVariant createSuffStat(SuffStatConstructorArgs... leaf_suff_stat_args) {
  return SuffStatType(leaf_suff_stat_args...);
}
 
template <typename LeafModelType, typename... LeafModelConstructorArgs>
static inline LeafModelVariant createLeafModel(LeafModelConstructorArgs... leaf_model_args) {
  return LeafModelType(leaf_model_args...);
}
 
static inline SuffStatVariant suffStatFactory(ModelType model_type, int basis_dim = 0) {
  if (model_type == kConstantLeafGaussian) {
    return createSuffStat<GaussianConstantSuffStat>();
  } else if (model_type == kUnivariateRegressionLeafGaussian) {
    return createSuffStat<GaussianUnivariateRegressionSuffStat>();
  } else if (model_type == kMultivariateRegressionLeafGaussian) {
    return createSuffStat<GaussianMultivariateRegressionSuffStat, int>(basis_dim);
  } else if (model_type == kLogLinearVariance) {
    return createSuffStat<LogLinearVarianceSuffStat>();
  } else {
    return createSuffStat<CloglogOrdinalSuffStat>();
  }
}
 
static inline LeafModelVariant leafModelFactory(ModelType model_type, double tau, Eigen::MatrixXd& Sigma0, double a, double b) {
  if (model_type == kConstantLeafGaussian) {
    return createLeafModel<GaussianConstantLeafModel, double>(tau);
  } else if (model_type == kUnivariateRegressionLeafGaussian) {
    return createLeafModel<GaussianUnivariateRegressionLeafModel, double>(tau);
  } else if (model_type == kMultivariateRegressionLeafGaussian) {
    return createLeafModel<GaussianMultivariateRegressionLeafModel, Eigen::MatrixXd>(Sigma0);
  } else if (model_type == kLogLinearVariance) {
    return createLeafModel<LogLinearVarianceLeafModel, double, double>(a, b);
  } else {
    return createLeafModel<CloglogOrdinalLeafModel, double, double>(a, b);
  }
}
 
template <typename SuffStatType, typename... SuffStatConstructorArgs>
static inline void AccumulateSuffStatProposed(
    SuffStatType& node_suff_stat, SuffStatType& left_suff_stat, SuffStatType& right_suff_stat, ForestDataset& dataset, ForestTracker& tracker,
    ColumnVector& residual, double global_variance, TreeSplit& split, int tree_num, int leaf_num, int split_feature, int num_threads,
    SuffStatConstructorArgs&... suff_stat_args) {
  // Determine the position of the node's indices in the forest tracking data structure
  int node_begin_index = tracker.UnsortedNodeBegin(tree_num, leaf_num);
  int node_end_index = tracker.UnsortedNodeEnd(tree_num, leaf_num);
 
  // Extract pointer to the feature partition for tree_num
  UnsortedNodeSampleTracker* unsorted_node_sample_tracker = tracker.GetUnsortedNodeSampleTracker();
  FeatureUnsortedPartition* feature_partition = unsorted_node_sample_tracker->GetFeaturePartition(tree_num);
 
  // Determine the number of threads to use
  int chunk_size = (node_end_index - node_begin_index) / num_threads;
  if (chunk_size < 100) {
    num_threads = 1;
    chunk_size = node_end_index - node_begin_index;
  }
 
  if (num_threads > 1) {
    // Split the work into num_threads chunks
    std::vector<std::pair<int, int>> thread_ranges(num_threads);
    std::vector<SuffStatType> thread_suff_stats_node;
    std::vector<SuffStatType> thread_suff_stats_left;
    std::vector<SuffStatType> thread_suff_stats_right;
    for (int i = 0; i < num_threads; i++) {
      thread_ranges[i] = std::make_pair(node_begin_index + i * chunk_size,
                                        node_begin_index + (i + 1) * chunk_size);
      thread_suff_stats_node.emplace_back(suff_stat_args...);
      thread_suff_stats_left.emplace_back(suff_stat_args...);
      thread_suff_stats_right.emplace_back(suff_stat_args...);
    }
 
    // Accumulate sufficient statistics
    StochTree::ParallelFor(0, num_threads, num_threads, [&](int i) {
      int start_idx = thread_ranges[i].first;
      int end_idx = thread_ranges[i].second;
      for (int idx = start_idx; idx < end_idx; idx++) {
        int obs_num = feature_partition->indices_[idx];
        double feature_value = dataset.CovariateValue(obs_num, split_feature);
        thread_suff_stats_node[i].IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
        if (split.SplitTrue(feature_value)) {
          thread_suff_stats_left[i].IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
        } else {
          thread_suff_stats_right[i].IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
        }
      }
    });
 
    // Combine the thread-local sufficient statistics
    for (int i = 0; i < num_threads; i++) {
      node_suff_stat.AddSuffStatInplace(thread_suff_stats_node[i]);
      left_suff_stat.AddSuffStatInplace(thread_suff_stats_left[i]);
      right_suff_stat.AddSuffStatInplace(thread_suff_stats_right[i]);
    }
  } else {
    for (int idx = node_begin_index; idx < node_end_index; idx++) {
      int obs_num = feature_partition->indices_[idx];
      double feature_value = dataset.CovariateValue(obs_num, split_feature);
      node_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
      if (split.SplitTrue(feature_value)) {
        left_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
      } else {
        right_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, obs_num, tree_num);
      }
    }
  }
}
 
template <typename SuffStatType>
static inline void AccumulateSuffStatExisting(SuffStatType& node_suff_stat, SuffStatType& left_suff_stat, SuffStatType& right_suff_stat, ForestDataset& dataset, ForestTracker& tracker,
                                              ColumnVector& residual, double global_variance, int tree_num, int split_node_id, int left_node_id, int right_node_id) {
  // Acquire iterators
  auto left_node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, left_node_id);
  auto left_node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, left_node_id);
  auto right_node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, right_node_id);
  auto right_node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, right_node_id);
 
  // Accumulate sufficient statistics for the left and split nodes
  for (auto i = left_node_begin_iter; i != left_node_end_iter; i++) {
    auto idx = *i;
    left_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
    node_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
  }
 
  // Accumulate sufficient statistics for the right and split nodes
  for (auto i = right_node_begin_iter; i != right_node_end_iter; i++) {
    auto idx = *i;
    right_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
    node_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
  }
}
 
template <typename SuffStatType, bool sorted>
static inline void AccumulateSingleNodeSuffStat(SuffStatType& node_suff_stat, ForestDataset& dataset, ForestTracker& tracker, ColumnVector& residual, int tree_num, int node_id) {
  // Acquire iterators
  std::vector<data_size_t>::iterator node_begin_iter;
  std::vector<data_size_t>::iterator node_end_iter;
  if (sorted) {
    // Default to the first feature if we're using the presort tracker
    node_begin_iter = tracker.SortedNodeBeginIterator(node_id, 0);
    node_end_iter = tracker.SortedNodeEndIterator(node_id, 0);
  } else {
    node_begin_iter = tracker.UnsortedNodeBeginIterator(tree_num, node_id);
    node_end_iter = tracker.UnsortedNodeEndIterator(tree_num, node_id);
  }
 
  // Accumulate sufficient statistics
  for (auto i = node_begin_iter; i != node_end_iter; i++) {
    auto idx = *i;
    node_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
  }
}
 
template <typename SuffStatType>
static inline void AccumulateCutpointBinSuffStat(SuffStatType& left_suff_stat, ForestTracker& tracker, CutpointGridContainer& cutpoint_grid_container,
                                                 ForestDataset& dataset, ColumnVector& residual, double global_variance, int tree_num, int node_id,
                                                 int feature_num, int cutpoint_num) {
  // Acquire iterators
  auto node_begin_iter = tracker.SortedNodeBeginIterator(node_id, feature_num);
  auto node_end_iter = tracker.SortedNodeEndIterator(node_id, feature_num);
 
  // Determine node start point
  data_size_t node_begin = tracker.SortedNodeBegin(node_id, feature_num);
 
  // Determine cutpoint bin start and end points
  data_size_t current_bin_begin = cutpoint_grid_container.BinStartIndex(cutpoint_num, feature_num);
  data_size_t current_bin_size = cutpoint_grid_container.BinLength(cutpoint_num, feature_num);
  data_size_t next_bin_begin = cutpoint_grid_container.BinStartIndex(cutpoint_num + 1, feature_num);
 
  // Cutpoint specific iterators
  // TODO: fix the hack of having to subtract off node_begin, probably by cleaning up the CutpointGridContainer interface
  auto cutpoint_begin_iter = node_begin_iter + (current_bin_begin - node_begin);
  auto cutpoint_end_iter = node_begin_iter + (next_bin_begin - node_begin);
 
  // Accumulate sufficient statistics
  for (auto i = cutpoint_begin_iter; i != cutpoint_end_iter; i++) {
    auto idx = *i;
    left_suff_stat.IncrementSuffStat(dataset, residual.GetData(), tracker, idx, tree_num);
  }
}
 
 
}  // namespace StochTree
 
#endif  // STOCHTREE_LEAF_MODEL_H_