random_effects.h

 
#ifndef STOCHTREE_RANDOM_EFFECTS_H_
#define STOCHTREE_RANDOM_EFFECTS_H_
 
#include <stochtree/category_tracker.h>
#include <stochtree/cutpoint_candidates.h>
#include <stochtree/data.h>
#include <stochtree/ensemble.h>
#include <stochtree/ig_sampler.h>
#include <stochtree/log.h>
#include <stochtree/normal_sampler.h>
#include <stochtree/partition_tracker.h>
#include <stochtree/prior.h>
#include <nlohmann/json.hpp>
#include <Eigen/Dense>
 
#include <fstream>
#include <map>
#include <memory>
#include <random>
#include <string>
#include <vector>
 
namespace StochTree {
 
class LabelMapper;
class MultivariateRegressionRandomEffectsModel;
class RandomEffectsContainer;
 
class RandomEffectsTracker {
 public:
  RandomEffectsTracker(std::vector<int>& group_indices);
  RandomEffectsTracker(int* group_indices, int num_observations);
  ~RandomEffectsTracker() {}
  inline data_size_t GetCategoryId(int observation_num) { return sample_category_mapper_->GetCategoryId(observation_num); }
  inline data_size_t CategoryBegin(int category_id) { return category_sample_tracker_->CategoryBegin(category_id); }
  inline data_size_t CategoryEnd(int category_id) { return category_sample_tracker_->CategoryEnd(category_id); }
  inline data_size_t CategorySize(int category_id) { return category_sample_tracker_->CategorySize(category_id); }
  inline int NumCategories() { return num_categories_; }
  inline int CategoryNumber(int category_id) { return category_sample_tracker_->CategoryNumber(category_id); }
  SampleCategoryMapper* GetSampleCategoryMapper() { return sample_category_mapper_.get(); }
  CategorySampleTracker* GetCategorySampleTracker() { return category_sample_tracker_.get(); }
  std::vector<data_size_t>::iterator UnsortedNodeBeginIterator(int category_id);
  std::vector<data_size_t>::iterator UnsortedNodeEndIterator(int category_id);
  std::map<int, int>& GetLabelMap() { return category_sample_tracker_->GetLabelMap(); }
  std::vector<int>& GetUniqueGroupIds() { return category_sample_tracker_->GetUniqueGroupIds(); }
  std::vector<data_size_t>& NodeIndices(int category_id) { return category_sample_tracker_->NodeIndices(category_id); }
  std::vector<data_size_t>& NodeIndicesInternalIndex(int internal_category_id) { return category_sample_tracker_->NodeIndicesInternalIndex(internal_category_id); }
  double* GetPredictions() { return rfx_predictions_.data(); }
  double GetPrediction(data_size_t observation_num) { return rfx_predictions_.at(observation_num); }
  void SetPrediction(data_size_t observation_num, double pred) { rfx_predictions_.at(observation_num) = pred; }
  void ResetFromSample(MultivariateRegressionRandomEffectsModel& rfx_model,
                       RandomEffectsDataset& rfx_dataset, ColumnVector& residual);
  void RootReset(MultivariateRegressionRandomEffectsModel& rfx_model,
                 RandomEffectsDataset& rfx_dataset, ColumnVector& residual);
 
 private:
  std::unique_ptr<SampleCategoryMapper> sample_category_mapper_;
  std::unique_ptr<CategorySampleTracker> category_sample_tracker_;
  std::vector<double> rfx_predictions_;
  int num_categories_;
  int num_observations_;
};
 
class LabelMapper {
 public:
  LabelMapper() {}
  LabelMapper(std::map<int, int> label_map) {
    label_map_ = label_map;
    for (const auto& [key, value] : label_map) keys_.push_back(key);
  }
  ~LabelMapper() {}
  void LoadFromLabelMap(std::map<int, int> label_map) {
    label_map_ = label_map;
    for (const auto& [key, value] : label_map) keys_.push_back(key);
  }
  bool ContainsLabel(int category_id) {
    auto pos = label_map_.find(category_id);
    return pos != label_map_.end();
  }
  int CategoryNumber(int category_id) {
    return label_map_[category_id];
  }
  void CopyFromOther(LabelMapper& other) {
    keys_ = other.Keys();
    label_map_ = other.Map();
  }
  void SaveToJsonFile(std::string filename) {
    nlohmann::json model_json = this->to_json();
    std::ofstream output_file(filename);
    output_file << model_json << std::endl;
  }
  void LoadFromJsonFile(std::string filename) {
    std::ifstream f(filename);
    nlohmann::json rfx_label_mapper_json = nlohmann::json::parse(f);
    this->Reset();
    this->from_json(rfx_label_mapper_json);
  }
  std::string DumpJsonString() {
    nlohmann::json model_json = this->to_json();
    return model_json.dump();
  }
  void LoadFromJsonString(std::string& json_string) {
    nlohmann::json rfx_label_mapper_json = nlohmann::json::parse(json_string);
    this->Reset();
    this->from_json(rfx_label_mapper_json);
  }
  std::vector<int>& Keys() { return keys_; }
  std::map<int, int>& Map() { return label_map_; }
  void Reset() {
    label_map_.clear();
    keys_.clear();
  }
  nlohmann::json to_json();
  void from_json(const nlohmann::json& rfx_label_mapper_json);
 
 private:
  std::map<int, int> label_map_;
  std::vector<int> keys_;
};
 
class MultivariateRegressionRandomEffectsModel {
 public:
  MultivariateRegressionRandomEffectsModel(int num_components, int num_groups) {
    normal_sampler_ = MultivariateNormalSampler();
    ig_sampler_ = InverseGammaSampler();
    num_components_ = num_components;
    num_groups_ = num_groups;
    working_parameter_ = Eigen::VectorXd(num_components_);
    group_parameters_ = Eigen::MatrixXd(num_components_, num_groups_);
    group_parameter_covariance_ = Eigen::MatrixXd(num_components_, num_components_);
    working_parameter_covariance_ = Eigen::MatrixXd(num_components_, num_components_);
  }
  ~MultivariateRegressionRandomEffectsModel() {}
 
  void ResetFromSample(RandomEffectsContainer& rfx_container, int sample_num);
 
  void SampleRandomEffects(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen);
  void SampleWorkingParameter(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen);
  void SampleGroupParameters(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen);
  void SampleVarianceComponents(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& tracker, double global_variance, std::mt19937& gen);
 
  void SetWorkingParameter(Eigen::VectorXd& working_parameter) {
    working_parameter_ = working_parameter;
  }
  void SetGroupParameters(Eigen::MatrixXd& group_parameters) {
    group_parameters_ = group_parameters;
  }
  void SetGroupParameter(Eigen::VectorXd& group_parameter, int group_id) {
    group_parameters_(Eigen::all, group_id) = group_parameter;
  }
  void SetWorkingParameterCovariance(Eigen::MatrixXd& working_parameter_covariance) {
    working_parameter_covariance_ = working_parameter_covariance;
  }
  void SetGroupParameterCovariance(Eigen::MatrixXd& group_parameter_covariance) {
    group_parameter_covariance_ = group_parameter_covariance;
  }
  void SetGroupParameterVarianceComponent(double value, int component_id) {
    group_parameter_covariance_(component_id, component_id) = value;
  }
  void SetVariancePriorShape(double value) {
    variance_prior_shape_ = value;
  }
  void SetVariancePriorScale(double value) {
    variance_prior_scale_ = value;
  }
 
  Eigen::VectorXd& GetWorkingParameter() {
    return working_parameter_;
  }
  Eigen::MatrixXd& GetGroupParameters() {
    return group_parameters_;
  }
  Eigen::MatrixXd& GetWorkingParameterCovariance() {
    return working_parameter_covariance_;
  }
  Eigen::MatrixXd& GetGroupParameterCovariance() {
    return group_parameter_covariance_;
  }
  double GetVariancePriorShape() {
    return variance_prior_shape_;
  }
  double GetVariancePriorScale() {
    return variance_prior_scale_;
  }
  inline int NumComponents() { return num_components_; }
  inline int NumGroups() { return num_groups_; }
 
  std::vector<double> Predict(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker) {
    std::vector<double> output(dataset.NumObservations());
    PredictInplace(dataset, tracker, output);
    return output;
  }
 
  void PredictInplace(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker, std::vector<double>& output) {
    Eigen::MatrixXd X = dataset.GetBasis();
    std::vector<int> group_labels = dataset.GetGroupLabels();
    CHECK_EQ(X.rows(), group_labels.size());
    int n = X.rows();
    CHECK_EQ(n, output.size());
    Eigen::MatrixXd alpha_diag = working_parameter_.asDiagonal().toDenseMatrix();
    int group_ind;
    for (int i = 0; i < n; i++) {
      group_ind = tracker.CategoryNumber(group_labels[i]);
      output[i] = X(i, Eigen::all) * alpha_diag * group_parameters_(Eigen::all, group_ind);
    }
  }
 
  void AddCurrentPredictionToResidual(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker, ColumnVector& residual) {
    data_size_t n = dataset.NumObservations();
    CHECK_EQ(n, residual.NumRows());
    double current_pred;
    double new_resid;
    for (data_size_t i = 0; i < n; i++) {
      current_pred = tracker.GetPrediction(i);
      new_resid = residual.GetElement(i) + current_pred;
      residual.SetElement(i, new_resid);
    }
  }
 
  void SubtractNewPredictionFromResidual(RandomEffectsDataset& dataset, RandomEffectsTracker& tracker, ColumnVector& residual) {
    Eigen::MatrixXd X = dataset.GetBasis();
    std::vector<int> group_labels = dataset.GetGroupLabels();
    CHECK_EQ(X.rows(), group_labels.size());
    int n = X.rows();
    double new_pred;
    double new_resid;
    Eigen::MatrixXd alpha_diag = working_parameter_.asDiagonal().toDenseMatrix();
    int group_ind;
    for (int i = 0; i < n; i++) {
      group_ind = tracker.CategoryNumber(group_labels[i]);
      new_pred = X(i, Eigen::all) * alpha_diag * group_parameters_(Eigen::all, group_ind);
      new_resid = residual.GetElement(i) - new_pred;
      residual.SetElement(i, new_resid);
      tracker.SetPrediction(i, new_pred);
    }
  }
 
  Eigen::VectorXd WorkingParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance);
  Eigen::MatrixXd WorkingParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance);
  Eigen::VectorXd GroupParameterMean(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id);
  Eigen::MatrixXd GroupParameterVariance(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int group_id);
  double VarianceComponentShape(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id);
  double VarianceComponentScale(RandomEffectsDataset& dataset, ColumnVector& residual, RandomEffectsTracker& rfx_tracker, double global_variance, int component_id);
 
 private:
  MultivariateNormalSampler normal_sampler_;
  InverseGammaSampler ig_sampler_;
 
  int num_components_;
  int num_groups_;
 
  Eigen::VectorXd working_parameter_;
  Eigen::MatrixXd group_parameters_;
 
  Eigen::MatrixXd group_parameter_covariance_;
 
  Eigen::MatrixXd working_parameter_covariance_;
 
  double variance_prior_shape_;
  double variance_prior_scale_;
};
 
class RandomEffectsContainer {
 public:
  RandomEffectsContainer(int num_components, int num_groups) {
    num_components_ = num_components;
    num_groups_ = num_groups;
    num_samples_ = 0;
  }
  RandomEffectsContainer() {
    num_components_ = 0;
    num_groups_ = 0;
    num_samples_ = 0;
  }
  ~RandomEffectsContainer() {}
  void SaveToJsonFile(std::string filename) {
    nlohmann::json model_json = this->to_json();
    std::ofstream output_file(filename);
    output_file << model_json << std::endl;
  }
  void LoadFromJsonFile(std::string filename) {
    std::ifstream f(filename);
    nlohmann::json rfx_container_json = nlohmann::json::parse(f);
    this->Reset();
    this->from_json(rfx_container_json);
  }
  std::string DumpJsonString() {
    nlohmann::json model_json = this->to_json();
    return model_json.dump();
  }
  void LoadFromJsonString(std::string& json_string) {
    nlohmann::json rfx_container_json = nlohmann::json::parse(json_string);
    this->Reset();
    this->from_json(rfx_container_json);
  }
  void CopyFromOther(RandomEffectsContainer& other) {
    this->Reset();
    num_samples_ = other.NumSamples();
    num_components_ = other.NumComponents();
    num_groups_ = other.NumGroups();
    beta_ = other.GetBeta();
    alpha_ = other.GetAlpha();
    xi_ = other.GetXi();
    sigma_xi_ = other.GetSigma();
  }
  void AddSample(MultivariateRegressionRandomEffectsModel& model);
  void DeleteSample(int sample_num);
  void Predict(RandomEffectsDataset& dataset, LabelMapper& label_mapper, std::vector<double>& output);
  inline int NumSamples() { return num_samples_; }
  inline int NumComponents() { return num_components_; }
  inline int NumGroups() { return num_groups_; }
  inline void SetNumSamples(int num_samples) { num_samples_ = num_samples; }
  inline void SetNumComponents(int num_components) { num_components_ = num_components; }
  inline void SetNumGroups(int num_groups) { num_groups_ = num_groups; }
  void Reset() {
    num_samples_ = 0;
    num_components_ = 0;
    num_groups_ = 0;
    beta_.clear();
    alpha_.clear();
    xi_.clear();
    sigma_xi_.clear();
  }
  std::vector<double>& GetBeta() { return beta_; }
  std::vector<double>& GetAlpha() { return alpha_; }
  std::vector<double>& GetXi() { return xi_; }
  std::vector<double>& GetSigma() { return sigma_xi_; }
  nlohmann::json to_json();
  void from_json(const nlohmann::json& rfx_container_json);
  void append_from_json(const nlohmann::json& rfx_container_json);
 
 private:
  int num_samples_;
  int num_components_;
  int num_groups_;
  std::vector<double> beta_;
  std::vector<double> alpha_;
  std::vector<double> xi_;
  std::vector<double> sigma_xi_;
  void AddAlpha(MultivariateRegressionRandomEffectsModel& model);
  void AddXi(MultivariateRegressionRandomEffectsModel& model);
  void AddSigma(MultivariateRegressionRandomEffectsModel& model);
};
 
inline void AppendRandomEffectsContainerSamples(std::unique_ptr<RandomEffectsContainer>& dst,
                                                const std::unique_ptr<RandomEffectsContainer>& src) {
  if (src == nullptr && dst == nullptr) return;
  if (src == nullptr || dst == nullptr) {
    Log::Fatal("Cannot merge samples: random effects container present in one chain but not the other");
  }
  // Check that group and basis dimensions match between the two containers before appending
  if (src->NumComponents() != dst->NumComponents()) {
    Log::Fatal("Cannot merge samples: random effects container has %d components in one chain but %d in the other", src->NumComponents(), dst->NumComponents());
  }
  if (src->NumGroups() != dst->NumGroups()) {
    Log::Fatal("Cannot merge samples: random effects container has %d groups in one chain but %d in the other", src->NumGroups(), dst->NumGroups());
  }
  dst->SetNumSamples(dst->NumSamples() + src->NumSamples());
  std::vector<double>& dst_beta = dst->GetBeta();
  std::vector<double>& dst_alpha = dst->GetAlpha();
  std::vector<double>& dst_xi = dst->GetXi();
  std::vector<double>& dst_sigma_xi = dst->GetSigma();
  std::vector<double>& src_beta = src->GetBeta();
  std::vector<double>& src_alpha = src->GetAlpha();
  std::vector<double>& src_xi = src->GetXi();
  std::vector<double>& src_sigma_xi = src->GetSigma();
  dst_beta.insert(dst_beta.end(), src_beta.begin(), src_beta.end());
  dst_alpha.insert(dst_alpha.end(), src_alpha.begin(), src_alpha.end());
  dst_xi.insert(dst_xi.end(), src_xi.begin(), src_xi.end());
  dst_sigma_xi.insert(dst_sigma_xi.end(), src_sigma_xi.begin(), src_sigma_xi.end());
}
 
}  // namespace StochTree
 
#endif  // STOCHTREE_RANDOM_EFFECTS_H_