BART.h

#ifndef BART_H
#define BART_H
#include <vector>
#include <random>
#include <unordered_map>
#include <boost/random.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/distributions/chi_squared.hpp>
 
using namespace std;
 
class bart_node {
public:
    int feature_number; 
    int split_index; 
    float node_value; 
    bart_node *parent; 
    bart_node *childs[2]; 
    vector<int> observation_indexes; 
    bool mark_change; 
 
    int num_feature_options; 
    int num_split_options; 
 
    bart_node() {
        feature_number = -1;
        split_index = -1;
        node_value = 0;
        parent = NULL; //root
                       //Childrens
        childs[0] = NULL; //smaller than or equal
        childs[1] = NULL; //biger than
        mark_change = true;
        num_feature_options = 0;
        num_split_options = 0;
    }
 
    bart_node(const bart_node &cp) {
        feature_number = cp.feature_number;
        split_index = cp.split_index;
        node_value = cp.node_value;
        parent = cp.parent;
        childs[0] = cp.childs[0];
        childs[1] = cp.childs[1];
        observation_indexes = cp.observation_indexes;
        mark_change = cp.mark_change;
        num_feature_options = cp.num_feature_options;
        num_split_options = cp.num_split_options;
    }
 
    void list_all_nodes(vector<bart_node *> &all_nodes);
    int depth();
    float variance(const vector<float> &y);
    void deep_clone(bart_node *&target);
    void print_tree(const vector<vector<float>> &feature_sorted_values) const; //for debug
    void validate_tree(const vector<vector<float>> &feature_sorted_values, const vector<float> &x, int nftrs) const; //for debug
private:
};
 
class tree_change_details {
public:
    vector<bart_node *>changed_nodes_before; 
    vector<bart_node *>changed_nodes_after; 
 
    int action = -1; 
    int num_node_selection = 0; 
};
 
enum bart_data_prior_type {
    regression_mean_shift = 0,
    classification = 1
};
 
class bart_tree;
 
class bart_params {
public:
    int min_obs_in_node; 
 
    float alpha; 
    float beta; 
 
                //general:
    bart_data_prior_type data_prior_type;
 
    //params for prior:
    float k; 
    float nu; 
    float lambda; 
 
    bart_params() {
        min_obs_in_node = 0;
        alpha = 1;
        beta = 1;
        sigsq_mu = 0;
        mean_mu = 0;
        nu = 3;
        k = 2;
        lambda = 1;
        data_prior_type = regression_mean_shift;
    }
 
    void set_classification(int num_trees) {
        mean_mu = 0;
        //k = 2;
        //nu = 3.0;
        data_prior_type = bart_data_prior_type::classification;
 
        sigsq_mu = pow(3 / (k * sqrt(num_trees)), 2);
    }
    void set_regression(int num_trees, float sample_var_y_in_data) {
        mean_mu = 0;
        //nu = 3.0;
        //k = 2;
        data_prior_type = bart_data_prior_type::regression_mean_shift;
 
        sigsq_mu = pow(1 / (2 * k * sqrt(num_trees)), 2);
 
        boost::math::chi_squared_distribution<> chi_dist(nu);
        double ten_pctile_chisq_df_hyper_nu = boost::math::cdf(chi_dist, 1 - 0.9);
        lambda = sample_var_y_in_data * ten_pctile_chisq_df_hyper_nu / nu;
        //lambda = 2.7e-3;
        //printf("ten_pctile_chisq_df_hyper_nu=%2.5f, var=%2.5f\n", ten_pctile_chisq_df_hyper_nu, sample_var_y_in_data);
    }
private:
    float mean_mu; 
    float sigsq_mu; 
 
    friend class bart_tree;
};
 
class BART;
 
class bart_tree {
public:
    bart_node *root; 
    double tree_loglikelihood; 
    bart_params params; 
 
    void next_gen_tree(const vector<float> &x, const vector<float> &y);
 
    void clone_tree(bart_tree &tree);
 
    void set_sigma(double sig) {
        sigma = sig;
    }
 
    bart_tree(const bart_tree &other) {
        root = other.root;
        tree_loglikelihood = other.tree_loglikelihood;
        _rnd_gen = other._rnd_gen;
        action_priors = other.action_priors;
        feature_to_sorted_vals = other.feature_to_sorted_vals;
        feature_to_val_index = other.feature_to_val_index;
        params = other.params;
    }
 
    bart_tree() {
        root = NULL;
        tree_loglikelihood = -1;
        _rnd_gen = mt19937(_rd());
 
        //Default Values:
        action_priors = { (float)0.25, (float)0.25, (float)0.4, (float)0.1 };
        params.min_obs_in_node = 5;
        params.alpha = (float)0.95;
        params.beta = 1.0;
 
        params.data_prior_type = bart_data_prior_type::regression_mean_shift;
        params.mean_mu = 0;
        params.nu = 3;
        params.sigsq_mu = 0;
        params.lambda = 1.0;
    }
private:
    //general variables
    random_device _rd;
    mt19937 _rnd_gen;
    vector<vector<float>> feature_to_sorted_vals;
    vector<unordered_map<float, int>> feature_to_val_index;
    double sigma;
 
    vector<float> action_priors;
    //actions:
    tree_change_details do_grow(const vector<float> &x, const vector<float> &y);
    tree_change_details do_prune(const vector<float> &x, const vector<float> &y);
    tree_change_details do_change(const vector<float> &x, const vector<float> &y);
    tree_change_details do_swap(const vector<float> &x, const vector<float> &y);
 
    float score_leaf(const vector<float> &y, const vector<int> &obs_indexes);
 
    double node_data_likelihood(const vector<bart_node *> &leaf_node, const vector<float> &x, int nftrs, const vector<float> &y);
    void calc_likelihood(const vector<float> &x, int nftrs, const vector<float> &y); //will calc before mean in all leaves
 
                                                                                     //helper for change:
    void get_avalible_change(const vector<float> &x, int nftrs, vector<bart_node *> &good_idx);
    void get_avalible_feats_change(const bart_node *selected_node, const vector<float> &x, int nftrs, vector<int> &good_idx);
    void propogate_change_down(bart_node *current_node, const vector<float> &x, int nftrs, vector<bart_node *> &list_nodes_after);
 
    //helper for grow:
    void get_avalible_grow(const vector<float> &x, int nftrs, vector<bart_node *> &good_idx) const;
    void get_avalible_feats_grow(const bart_node *selected_node, const vector<float> &x, int nftrs, vector<int> &good_idx);
 
    //helper for prune;
    void get_avalible_prune(const vector<float> &x, int nftrs, vector<bart_node *> &good_idx);
 
    //helper for swap:
    void get_avalible_swap(const vector<float> &x, int nftrs, vector<bart_node *> &good_idx);
 
    bool has_split(const bart_node *current_node, const vector<float> &x, int nftrs, int feature_num) const;
    bool select_split(bart_node *current_node, const vector<float> &x, int nftrs,
        vector<vector<int>> &split_obx_indexes);
 
    void commit_change(const tree_change_details &change);
    void rollback_change(const tree_change_details &change);
 
    int clear_tree_mem(bart_node *node); //erase count
protected:
    void predict(const vector<float> &x, int nSamples, vector<float> &scores) const;
    void predict_on_train(const vector<float> &x, int nSamples, vector<float> &scores) const; //faster
 
    friend class BART;
};
 
class BART {
public:
    int ntrees; 
    int iter_count; 
    int burn_count; 
    int restart_count; 
    bart_params tree_params; 
 
    void learn(const vector<float> &x, const vector<float> &y);
 
    void predict(const vector<float> &x, int nSamples, vector<float> &scores) const;
 
    BART(int ntrees, int iterations, int burn_cnt, int restart_cnt, bart_params &tree_pr) {
        nftrs = 0;
        //default:
        this->ntrees = ntrees;
        this->iter_count = iterations;
        this->burn_count = burn_cnt;
        this->tree_params = tree_pr;
        this->restart_count = restart_cnt;
 
        _trees.resize(ntrees);
        for (size_t i = 0; i < ntrees; ++i)
            _trees[i].params = this->tree_params;
        trans_y_b = 0;
        trans_y_max = 0;
    }
 
    void operator=(const BART &other) {
        ntrees = other.ntrees; iter_count = other.iter_count;
        burn_count = other.burn_count; restart_count = other.restart_count;
        tree_params = other.tree_params;
        nftrs = other.nftrs;
        _trees.resize(other._trees.size());
        for (size_t i = 0; i < other._trees.size(); ++i) {
            _trees[i].action_priors = other._trees[i].action_priors;
            _trees[i].root = other._trees[i].root;
            _trees[i].tree_loglikelihood = other._trees[i].tree_loglikelihood;
            _trees[i]._rnd_gen = other._trees[i]._rnd_gen;
            _trees[i].feature_to_sorted_vals = other._trees[i].feature_to_sorted_vals;
            _trees[i].feature_to_val_index = other._trees[i].feature_to_val_index;
            _trees[i].params = other._trees[i].params;
        }
    }
 
    ~BART() {
        for (size_t i = 0; i < _trees.size(); ++i)
            _trees[i].clear_tree_mem(_trees[i].root);
    }
 
private:
    int nftrs;
    vector<bart_tree> _trees;
    float trans_y_b; //movement of y values
    float trans_y_max;
 
    void transform_y(vector<float> &y);
    void untransform_y(vector<float> &y) const;
 
    void init_hyper_parameters(const vector<float> &residuals);
    void update_sigma_param(boost::mt19937 &rng, const vector<float> &residuals, double &sigma); //for regression
    void update_latent_z_params(boost::random::random_number_generator<boost::mt19937> &rng_gen,
        const vector<float> &x, const vector<float> &y, const vector<bart_tree> &forest_trees,
        vector<float> &residuals); //for classification
};
 
#endif // !BART_H