Growing a (Decision) Tree

#[derive(Debug)]
pub struct Rule {
    dimension: String,
    cutoff: f64,
    metric: f64,
}

#[derive(Debug)]
pub struct Decision {
    rule: Option<Rule>,
    confidence: f64,
    prediction: String,
}

pub fn predict_majority_dataframe<'a>(
    data: & 'a DataFrame, target: & str
) -> PolarsResult<Decision> {
    // extract the categorical target column
    let labels = data.column(target)?.categorical()?;

    let total = labels.len() as f64;

    // count all categories and sort them
    let result_count = labels.value_counts()?;

    // get the most frequent category
    let result_cat = result_count.head(Some(1));

    // transform the series into a categorical vector
    let actual_cat = result_cat
        .column(target)?
        .categorical()?;

    // collect all categories as strings
    let string_cat: Vec<String> = actual_cat
        .iter_str()
        .flatten()
        .map(|name| (*name).into())
        .collect();

    let probability: Vec<f64>= result_cat
        .column("counts")?
        .u32()?
    .iter()
    .flatten()
    .map(|c| (c as f64)/total)
        .collect();
    // return the most common category as a string
    return Ok(
        Decision{
            rule: None,
            prediction: string_cat
                .get(0)
                .unwrap()
                .to_string(),
            confidence: probability
                .get(0)
                .unwrap()
                .to_owned()
        }
    );
}

pub fn evaluate_best_split<'a>(
    data: & DataFrame,
    features: & HashSet <& str>,
    target: & str,
) -> PolarsResult<Rule> {

    // iteratively evaluate the metric on all features
    let metrics: PolarsResult<Vec<LazyFrame>> = features
        .iter()
        .map(|feature| {
            Ok(evaluate_metric(& data, feature, target)?
                .lazy()
                .with_column(feature.lit().alias("feature")))
        })
        .collect();

    // join all results in a single dataframe
    let concat_rules = UnionArgs {
        parallel: true,
        rechunk: true,
        to_supertypes: true,
    };
    let concat_metrics: DataFrame = concat(metrics?, concat_rules)?.collect()?;

    // search for the best split
    let expr: Expr = col("metrics").lt_eq(col("metrics").min());
    let best_split: DataFrame = concat_metrics
        .clone()
        .lazy()
        .filter(expr)
        .select([col("feature"), col("split"), col("metrics")])
        .collect()?;

    let chosen_features: Vec<String> = best_split
        .column("feature")?
        .str()?
        .iter()
        .flatten()
        .map(|name| <& str as Into<String>>::into(name))
        .collect();

    let chosen_split_point: f64 = best_split.column("split")?.f64()?.get(0).unwrap();

    let split_metric: f64 = best_split.column("metrics")?.f64()?.get(0).unwrap();
    Ok(Rule {
        dimension: chosen_features
            .get(0)
            .unwrap()
            .to_string(),
        cutoff: chosen_split_point,
        metric: split_metric,
    })
}

        if (!current_features.is_empty()) && // exhausted features
            (confidence < 1.0) && // all elements belong to one category
            (data.shape().0 > self.min_size) && // size is below minimum threshold
            (level <= self.max_level){ // maximum depth reached

#[derive(Debug)]
pub struct DTreeBuilder<'a>{
    max_level: usize,
    min_size: usize,
    features: HashSet<& 'a str>,
    target: & 'a str,
    reuse_features: bool
}

// uses a struct to define trees constraints
impl <'a>DTreeBuilder<'a> {
    pub fn new(features: HashSet<& 'a str>, target : & 'a str) -> DTreeBuilder<'a>{
        DTreeBuilder{
            max_level: 3,
            min_size: 1,
            features,
            target,
            reuse_features: true
        }
    }

    pub fn set_max_level(mut self, max_level: usize) -> DTreeBuilder<'a>{
        self.max_level = max_level;
        self
    }

    pub fn set_min_size(mut self, min_size: usize) -> DTreeBuilder<'a>{
        self.min_size = min_size;
        self
    }

    pub fn set_reuse_features(mut self, reuse_features : bool) -> DTreeBuilder<'a>{
        self.reuse_features = reuse_features;
        self
    }
}

impl <'a>DTreeBuilder<'a> {
    // ...
    pub fn build(
        & self,
        data: & DataFrame,
    ) -> PolarsResult<btree::Tree<Decision>> {
        let current_features = if !self.reuse_features {
            let feats = self.features.clone();
            Some(feats)
        }else{
            None
        };
        println!("{1:->0$}{2:?}{1:-<0$}", 20, "\n", self);
        let root = self.build_node(data, 1, & current_features)?;
        Ok(btree::Tree::from_node(root))
    }
    // ...
}

impl <'a>DTreeBuilder<'a> {
    // ...
    fn build_node(
        & self,
        data: & DataFrame,
        level: usize, // tree depth
        features: & Option<HashSet<& str>>, // optionally used to remove features
    ) -> PolarsResult<btree::Node<Decision>> {
        let prediction = predict_majority_dataframe(data, self.target)?;
        let confidence = prediction.confidence;
        let mut node = btree::Node::new(prediction);
        let current_features = features.clone().unwrap_or(self.features.clone());
        // check stop conditions
        if (!current_features.is_empty()) && // exhausted features
            (confidence < 1.0) && // all elements belong to one category
            (data.shape().0 > self.min_size) && // size is below minimum threshold
            (level <= self.max_level){ // maximum depth reached
                let rule = evaluate_best_split(data, & current_features, self.target)?;
                let higher: DataFrame = data
                    .clone()
                    .lazy()
                    .filter(col(& rule.dimension).gt(rule.cutoff))
                    .collect()?;
                let lower: DataFrame = data
                    .clone()
                    .lazy()
                    .filter(col(& rule.dimension).lt_eq(rule.cutoff))
                    .collect()?;
                // remove features only if requested by the user
                let next_features = match features {
                    None => None,
                    Some(feats) => {
                        let mut reduced_features =
                            feats.clone();
                        reduced_features.remove(rule.dimension.as_str());
                        let feats_vec: Vec<String> = reduced_features
                            .iter()
                            .map(|s| s.to_string())
                            .collect();
                        Some(reduced_features)
                    }
                };
                node.value.rule = Some(rule);
                // creates leaves
                node.left = self
                    .build_node(& higher, level + 1, & next_features)?
                    .into();
                node.right = self
                    .build_node(& lower, level + 1, & next_features)?
                    .into();
            }
        Ok(node)
    }
    // ...
}