bart_sampler.h

 
#ifndef STOCHTREE_BART_SAMPLER_H_
#define STOCHTREE_BART_SAMPLER_H_
 
#include <stochtree/bart.h>
#include <stochtree/container.h>
#include <stochtree/data.h>
#include <stochtree/ensemble.h>
#include <stochtree/leaf_model.h>
#include <stochtree/linear_regression.h>
#include <stochtree/ordinal_sampler.h>
#include <stochtree/partition_tracker.h>
#include <stochtree/prior.h>
#include <stochtree/probit.h>
#include <stochtree/random_effects.h>
#include <stochtree/tree_sampler.h>
#include <stochtree/variance_model.h>
#include <iomanip>
#include <memory>
#include <sstream>
#include <string>
#include <type_traits>
#include <variant>
#include <vector>
 
namespace StochTree {
 
class BARTSampler {
 public:
  // If continuation is true, the sampler warm-starts from the last retained sample
  // already present in `samples` (rather than initializing forests to root), and the
  // existing forest container is preserved so that new samples are appended to it.
  //
  // If warmstart_source is non-null (and continuation is false), this is a FRESH run whose forests /
  // scalars / rfx are seeded from an EXTERNAL model's samples (`*warmstart_source`) at
  // `warmstart_sample_num` (1-indexed). This backs bart(previous_model_json=...): the destination
  // `samples` container is fresh, but the active forest is reconstituted from the previous model's
  // retained sample. Same-scale is assumed (no leaf rescale); the previous model's y_std is used to
  // convert its stored global variance back to standardized space, matching the R main-branch flow.
  // The borrowed pointer need only outlive this constructor call.
  BARTSampler(BARTSamples& samples, BARTConfig& config, BARTData& data, bool continuation = false,
              BARTSamples* warmstart_source = nullptr, int warmstart_sample_num = 0);
 
  // Main entry point for running the BART GFR sampler
  // If num_chains > 0, captures snapshots of the last num_chains GFR states for fork_chains()
  void run_gfr(BARTSamples& samples, int num_gfr, bool keep_gfr, int num_chains = 0);
 
  // Run a single chain of the BART MCMC sampler
  void run_mcmc(BARTSamples& samples, int num_burnin, int keep_every, int num_mcmc);
 
  // Run num_chains independent MCMC chains sequentually based on GFR snapshots captured by run_gfr() or re-initialized from root
  void run_mcmc_chains(BARTSamples& samples, int num_chains, int num_burnin, int keep_every, int num_mcmc);
 
  // Post-process samples by extracting test set predictions and running any necessary transformations
  void postprocess_samples(BARTSamples& samples, int start_sample = 0);
 
  // Serialize the internal RNG state to a string. std::mt19937 round-trips losslessly through
  // its stream operators, so this captures the exact position in the random stream -- used to
  // persist RNG state across a sample() / continue_sampling() boundary for bit-identical results.
  std::string GetRngState() const {
    std::ostringstream oss;
    oss << rng_;
    return oss.str();
  }
 
  // Restore the internal RNG state from a string produced by GetRngState(). Resumes the random
  // stream at exactly the captured position. Must be called after construction (InitializeState
  // unconditionally re-seeds rng_) and before any sampling draws.
  void SetRngState(const std::string& state) {
    std::istringstream iss(state);
    iss >> rng_;
  }
 
  // Regenerate the latent outcome for a continuation warm-start (probit or cloglog). The latent is not
  // persisted (it is re-drawn each MCMC iteration), so this draws a fresh latent to place the residual
  // / auxiliary state in a valid, stationary state before the first continued draw. No-op unless the
  // model uses a latent-augmentation link (probit or cloglog) and has a mean forest. MUST be called
  // after SetRngState so the draw comes from the resumed (or user-re-seeded) stream rather than the
  // pre-seed default RNG. For cloglog, the deterministic auxiliary state (eta / cutpoints / residual)
  // is set up in InitializeState's continuation cloglog block before this draw.
  void RegenerateLatentOutcome(BARTSamples& samples);
 
 private:
  void InitializeState(BARTSamples& samples, bool continuation = false);
  bool initialized_ = false;
 
  void RestoreStateFromGFRSnapshot(BARTSamples& samples, int snapshot_index);
 
  void WarmStartResetFromSample(BARTSamples& samples, BARTSamples& source, int idx);
 
  void RestoreStateDefault();
 
  void RunOneIteration(BARTSamples& samples, bool gfr, bool keep_sample, bool write_snapshot = false);
 
  struct MeanForestInitVisitor {
    BARTSampler& sampler;
    BARTSamples& samples;
    void operator()(GaussianConstantLeafModel& model) {
      sampler.mean_forest_ = std::make_unique<TreeEnsemble>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      samples.mean_forests = std::make_unique<ForestContainer>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      sampler.mean_forest_tracker_ = std::make_unique<ForestTracker>(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train);
      sampler.tree_prior_mean_ = std::make_unique<TreePrior>(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean);
      sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean);
      UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), !sampler.config_.leaf_constant_mean, std::minus<double>());
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
      sampler.has_mean_forest_ = true;
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      sampler.mean_forest_ = std::make_unique<TreeEnsemble>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      samples.mean_forests = std::make_unique<ForestContainer>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      sampler.mean_forest_tracker_ = std::make_unique<ForestTracker>(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train);
      sampler.tree_prior_mean_ = std::make_unique<TreePrior>(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean);
      sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean);
      UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), !sampler.config_.leaf_constant_mean, std::minus<double>());
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
      sampler.has_mean_forest_ = true;
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      sampler.mean_forest_ = std::make_unique<TreeEnsemble>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      samples.mean_forests = std::make_unique<ForestContainer>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      sampler.mean_forest_tracker_ = std::make_unique<ForestTracker>(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train);
      sampler.tree_prior_mean_ = std::make_unique<TreePrior>(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean);
      sampler.mean_forest_->SetLeafVector(sampler.init_val_mean_vec_);
      UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), true, std::minus<double>());
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
      sampler.has_mean_forest_ = true;
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      sampler.mean_forest_ = std::make_unique<TreeEnsemble>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      samples.mean_forests = std::make_unique<ForestContainer>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      sampler.mean_forest_tracker_ = std::make_unique<ForestTracker>(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train);
      sampler.tree_prior_mean_ = std::make_unique<TreePrior>(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean);
      sampler.mean_forest_->SetLeafValue(sampler.init_val_mean_ / sampler.config_.num_trees_mean);
      UpdateResidualEntireForest(*sampler.mean_forest_tracker_, *sampler.forest_dataset_, *sampler.residual_, sampler.mean_forest_.get(), false, std::minus<double>());
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
      sampler.has_mean_forest_ = true;
    }
  };
 
  struct MeanForestResetVisitor {
    BARTSampler& sampler;
    BARTSamples& samples;
    TreeEnsemble& forest;
    void operator()(GaussianConstantLeafModel& model) {
      sampler.mean_forest_->ReconstituteFromForest(forest);
      sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true);
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      sampler.mean_forest_->ReconstituteFromForest(forest);
      sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true);
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      sampler.mean_forest_->ReconstituteFromForest(forest);
      sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true);
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      sampler.mean_forest_->ReconstituteFromForest(forest);
      sampler.mean_forest_tracker_->ReconstituteFromForest(forest, *sampler.forest_dataset_, *sampler.residual_, true);
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
    }
  };
 
  struct MeanForestContinuationInitVisitor {
    BARTSampler& sampler;
    BARTSamples& samples;   // destination container (appended to for continuation)
    BARTSamples& source;    // seed source (== samples for continuation)
    int idx;                // seed index in source.mean_forests
    bool fresh_container;   // create a fresh destination mean_forests container (previous-model warm-start)
    void WarmStart() {
      // Reset active forest, tracker, and tree prior
      sampler.mean_forest_ = std::make_unique<TreeEnsemble>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      sampler.mean_forest_tracker_ = std::make_unique<ForestTracker>(sampler.forest_dataset_->GetCovariates(), sampler.config_.feature_types, sampler.config_.num_trees_mean, sampler.data_.n_train);
      sampler.tree_prior_mean_ = std::make_unique<TreePrior>(sampler.config_.alpha_mean, sampler.config_.beta_mean, sampler.config_.min_samples_leaf_mean, sampler.config_.max_depth_mean);
      if (fresh_container) {
        samples.mean_forests = std::make_unique<ForestContainer>(sampler.config_.num_trees_mean, sampler.config_.leaf_dim_mean, sampler.config_.leaf_constant_mean, sampler.config_.exponentiated_leaf_mean);
      }
      sampler.has_mean_forest_ = true;
      // Re-initialize the active forest and tracker from the seed sample
      TreeEnsemble& seed_forest = *source.mean_forests->GetEnsemble(idx);
      sampler.mean_forest_->ReconstituteFromForest(seed_forest);
      sampler.mean_forest_tracker_->ReconstituteFromForest(seed_forest, *sampler.forest_dataset_, *sampler.residual_, true);
      sampler.mean_forest_tracker_->UpdatePredictions(sampler.mean_forest_.get(), *sampler.forest_dataset_.get());
    }
    void operator()(GaussianConstantLeafModel& model) {
      WarmStart();
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      WarmStart();
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      WarmStart();
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      // Reconstitute the forest from the last sample exactly as the Gaussian case. The tracker
      // residual swap corrupts the cloglog residual (which holds raw y, not y - f), but
      // InitializeState's continuation cloglog block overwrites the residual with raw y and sets the
      // auxiliary state (eta / cutpoints) afterward; the latent is regenerated post-SetRngState.
      WarmStart();
    }
  };
 
  struct GFROneIterationVisitor {
    BARTSampler& sampler;
    BARTSamples& samples;
    bool keep_sample;
    void operator()(GaussianConstantLeafModel& model) {
      GFRSampleOneIter<GaussianConstantLeafModel, GaussianConstantSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types,
          sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads);
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      GFRSampleOneIter<GaussianUnivariateRegressionLeafModel, GaussianUnivariateRegressionSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types,
          sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads);
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      GFRSampleOneIter<GaussianMultivariateRegressionLeafModel, GaussianMultivariateRegressionSuffStat, int>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types,
          sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads,
          sampler.config_.leaf_dim_mean);
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      GFRSampleOneIter<CloglogOrdinalLeafModel, CloglogOrdinalSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, sampler.config_.feature_types,
          sampler.config_.cutpoint_grid_size, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/false,
          /*num_features_subsample=*/sampler.config_.num_features_subsample_mean, sampler.config_.num_threads);
    }
  };
 
  struct MCMCOneIterationVisitor {
    BARTSampler& sampler;
    BARTSamples& samples;
    bool keep_sample;
    void operator()(GaussianConstantLeafModel& model) {
      MCMCSampleOneIter<GaussianConstantLeafModel, GaussianConstantSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_threads=*/sampler.config_.num_threads);
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      MCMCSampleOneIter<GaussianUnivariateRegressionLeafModel, GaussianUnivariateRegressionSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_threads=*/sampler.config_.num_threads);
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      MCMCSampleOneIter<GaussianMultivariateRegressionLeafModel, GaussianMultivariateRegressionSuffStat, int>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/true,
          /*num_threads=*/sampler.config_.num_threads,
          sampler.config_.leaf_dim_mean);
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      MCMCSampleOneIter<CloglogOrdinalLeafModel, CloglogOrdinalSuffStat>(
          *sampler.mean_forest_, *sampler.mean_forest_tracker_, *samples.mean_forests, model,
          *sampler.forest_dataset_, *sampler.residual_, *sampler.tree_prior_mean_, sampler.rng_,
          sampler.config_.var_weights_mean, sampler.config_.sweep_update_indices_mean, sampler.global_variance_, /*keep_forest=*/keep_sample,
          /*pre_initialized=*/true, /*backfitting=*/false,
          /*num_threads=*/sampler.config_.num_threads);
    }
  };
 
  struct ScaleUpdateVisitor {
    BARTSampler& sampler;
    double leaf_scale;
    void operator()(GaussianConstantLeafModel& model) {
      model.SetScale(leaf_scale);
    }
    void operator()(GaussianUnivariateRegressionLeafModel& model) {
      model.SetScale(leaf_scale);
    }
    void operator()(GaussianMultivariateRegressionLeafModel& model) {
      // No-op for multivariate regression leaf model since scale is a vector
    }
    void operator()(CloglogOrdinalLeafModel& model) {
      // No-op for cloglog ordinal leaf model since scale is not a variance parameter
    }
  };
 
  BARTConfig& config_;
  BARTData& data_;
 
  BARTSamples* warmstart_source_ = nullptr;
  int warmstart_sample_num_ = 0;
 
  std::variant<GaussianConstantLeafModel, GaussianUnivariateRegressionLeafModel, GaussianMultivariateRegressionLeafModel, CloglogOrdinalLeafModel> mean_leaf_model_;
  LogLinearVarianceLeafModel variance_leaf_model_;
 
  std::unique_ptr<TreeEnsemble> mean_forest_;
  std::unique_ptr<ForestTracker> mean_forest_tracker_;
  std::unique_ptr<TreePrior> tree_prior_mean_;
  bool has_mean_forest_ = false;
  double init_val_mean_;
  std::vector<double> init_val_mean_vec_;
  std::unique_ptr<OrdinalSampler> ordinal_sampler_;
 
  std::unique_ptr<TreeEnsemble> variance_forest_;
  std::unique_ptr<ForestTracker> variance_forest_tracker_;
  std::unique_ptr<TreePrior> tree_prior_variance_;
  bool has_variance_forest_ = false;
  double init_val_variance_;
 
  std::unique_ptr<MultivariateRegressionRandomEffectsModel> random_effects_model_;
  std::unique_ptr<RandomEffectsTracker> random_effects_tracker_;
  std::unique_ptr<RandomEffectsDataset> random_effects_dataset_;
  bool has_random_effects_ = false;
 
  std::unique_ptr<ColumnVector> residual_;
  std::unique_ptr<ColumnVector> outcome_raw_;
  std::unique_ptr<ForestDataset> forest_dataset_;
  std::unique_ptr<ForestDataset> forest_dataset_test_;
  bool has_test_ = false;
 
  std::mt19937 rng_;
 
  double global_variance_;
  double leaf_scale_;
  std::vector<double> leaf_scale_multivariate_;
 
  // Global error scale model
  std::unique_ptr<GlobalHomoskedasticVarianceModel> var_model_;
  bool sample_sigma2_global_ = false;
 
  // Leaf scale model
  std::unique_ptr<LeafNodeHomoskedasticVarianceModel> leaf_scale_model_;
  bool sample_sigma2_leaf_ = false;
 
  struct GFRSnapshot {
    // Forest state
    std::unique_ptr<TreeEnsemble> mean_forest;      // null if no mean forest
    std::unique_ptr<TreeEnsemble> variance_forest;  // null if no variance forest
 
    // Global parameters
    double sigma2;
    double leaf_scale;
    std::vector<double> leaf_scale_multivariate;
 
    // Residual (incorporates forest + RFX contributions for a given sampler iteration)
    std::vector<double> residual;
 
    // Heteroskedastic variance model state
    std::vector<double> variance_weights;  // forest_dataset_ var_weights at snapshot time; only valid when variance_forest != null
 
    // Cloglog model state
    std::vector<double> cloglog_forest_preds;        // cached forest predictions; only populated with cloglog link
    std::vector<double> cloglog_latent_outcome;      // cached latent outcome state; only populated with cloglog link
    std::vector<double> cloglog_logscale_cutpoints;  // cached logscale cutpoints; only populated with cloglog link and ordinal outcome
 
    // RFX model state (only populated when has_random_effects_)
    Eigen::VectorXd rfx_working_parameter;
    Eigen::MatrixXd rfx_group_parameters;
    Eigen::MatrixXd rfx_group_parameter_covariance;
    Eigen::MatrixXd rfx_working_parameter_covariance;
    double rfx_variance_prior_shape;
    double rfx_variance_prior_scale;
  };
 
  std::vector<GFRSnapshot> gfr_snapshots_;
};
 
}  // namespace StochTree
 
#endif  // STOCHTREE_BART_SAMPLER_H_