prediction.h

 
#ifndef STOCHTREE_PREDICTION_H_
#define STOCHTREE_PREDICTION_H_
 
#include <vector>
#include <stochtree/bart.h>
#include <stochtree/bcf.h>
#include <stochtree/container.h>
#include <stochtree/meta.h>
#include <stochtree/random_effects.h>
 
namespace StochTree {
 
enum class PredType {
  kPosterior,
  kMean
};
 
enum class PredScale {
  kLinear,
  kProbability,
  kClass
};
 
struct BARTPredTerms {
  bool y_hat = true;
  bool mean_forest = false;
  bool variance_forest = false;
  bool random_effects = false;
};
 
struct BARTPredictionResult {
  // Outcome conditional mean
  std::vector<double> y_hat;
 
  // Covariate-dependent prognostic term (mu(x))
  std::vector<double> mean_forest_predictions;
 
  // Conditional variance term
  std::vector<double> variance_forest_predictions;
 
  // Random effects predictions
  std::vector<double> rfx_predictions;
};
 
struct BARTPredictionMetadata {
  // Metadata about the samples / model (e.g., number of samples, burn-in, etc.)
  int num_samples = 0;
  int num_obs = 0;
  int num_basis = 0;
  double y_bar = 0.0;
  double y_std = 0.0;
  bool has_variance_forest = false;
  bool has_rfx = false;
  BARTRFXModelSpec rfx_model_spec;
  PredType pred_type = PredType::kPosterior;
  BARTPredTerms pred_terms;
  PredScale pred_scale = PredScale::kLinear;
  LinkFunction link_function = LinkFunction::Identity;
  OutcomeType outcome_type = OutcomeType::Continuous;
  int cloglog_num_classes = 0;
};
 
BARTPredictionResult predict_bart_model(BARTData& data, BARTSamples& samples, BARTPredictionMetadata& metadata);
 
struct BCFPredTerms {
  bool y_hat = true;
  bool mu_x = false;
  bool tau_x = false;
  bool prognostic_function = false;
  bool cate = false;
  bool conditional_variance = false;
  bool random_effects = false;
};
 
struct BCFPredictionResult {
  // Outcome conditional mean
  std::vector<double> y_hat;
 
  // Covariate-dependent prognostic term (mu(x))
  std::vector<double> mu_x;
 
  // Covariate-dependent treatment effect term (tau(x))
  std::vector<double> tau_x;
 
  // Prognostic function (includes mu(x) and any random intercepts, provided random effects
  // were estimated with `intercept_only` or `intercept_plus_treatment` specification)
  std::vector<double> prognostic_function;
 
  // CATE function (includes tau(x) and any random slopes on treatment, provided
  // random effects were estimated with `intercept_plus_treatment` specification)
  std::vector<double> cate;
 
  // Conditional variance term
  std::vector<double> conditional_variance;
 
  // Random effects predictions
  std::vector<double> random_effects;
};
 
struct BCFPredictionMetadata {
  // Metadata about the samples / model (e.g., number of samples, burn-in, etc.) could be added here as needed
  int num_samples = 0;
  int num_obs = 0;
  int treatment_dim = 0;
  double y_bar = 0.0;
  double y_std = 0.0;
  bool has_variance_forest = false;
  bool has_rfx = false;
  BCFRFXModelSpec rfx_model_spec;
  bool adaptive_coding = false;
  bool sample_tau_0 = false;
  PredType pred_type = PredType::kPosterior;
  BCFPredTerms pred_terms;
  PredScale pred_scale = PredScale::kLinear;
};
 
BCFPredictionResult predict_bcf_model(BCFData& data, BCFSamples& samples, BCFPredictionMetadata& metadata);
 
}  // namespace StochTree
 
#endif  // STOCHTREE_PREDICTION_H_