dtrees_rs/algorithms/common/heuristics/
mod.rs

1//! Heuristics that rank candidate features, best first.
2
3pub mod helpers;
4
5use crate::algorithms::common::heuristics::helpers::{
6    entropy, gini_index, information_gain, weighted_entropy,
7};
8use crate::algorithms::common::utils::deduce_sibling_error_with_buffer;
9use crate::cover::Cover;
10use crate::globals::item;
11
12/// Scores a split from the class counts of the parent and both children.
13type ScoreFn = Box<dyn Fn(&[usize], &[usize], &[usize], f64) -> f64>;
14
15/// Ranks candidate features at a node.
16pub trait Heuristic: Send + Sync {
17    /// Sorts `candidates` best first and returns their scores in the same
18    /// order (empty when the heuristic does not score).
19    fn compute(&self, cover: &mut Cover, candidates: &mut Vec<usize>) -> Vec<f64>;
20
21    /// Scores each candidate with `scorer` and sorts them, ascending when
22    /// `lower_is_better` and descending otherwise.
23    fn compute_with_scorer(
24        &self,
25        parent_entropy: f64,
26        cover: &mut Cover,
27        candidates: &mut Vec<usize>,
28        scorer: ScoreFn,
29        lower_is_better: bool,
30    ) -> Vec<f64> {
31        if candidates.is_empty() {
32            return vec![];
33        }
34
35        let root_distribution = cover.labels_count();
36        let mut left_distribution = vec![0; cover.num_labels];
37        let mut right_distribution = vec![0; cover.num_labels];
38
39        let mut scores: Vec<_> = candidates
40            .iter()
41            .map(|&attr| {
42                cover.branch_on(item(attr, 0));
43                cover.labels_count_with_buffer(&mut left_distribution);
44                cover.backtrack();
45
46                deduce_sibling_error_with_buffer(
47                    &root_distribution,
48                    &left_distribution,
49                    &mut right_distribution,
50                );
51
52                let score = scorer(
53                    &root_distribution,
54                    &left_distribution,
55                    &right_distribution,
56                    parent_entropy,
57                );
58
59                (attr, score)
60            })
61            .collect();
62
63        if lower_is_better {
64            scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
65        } else {
66            scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
67        }
68
69        candidates.clear();
70        let (sorted_candidates, scores): (Vec<usize>, Vec<f64>) = scores.into_iter().unzip();
71        candidates.extend(sorted_candidates);
72        scores
73    }
74}
75
76/// Keeps the candidates in their original order.
77#[derive(Default)]
78pub struct NoHeuristic;
79
80impl Heuristic for NoHeuristic {
81    fn compute(&self, _cover: &mut Cover, _candidates: &mut Vec<usize>) -> Vec<f64> {
82        vec![]
83    }
84}
85
86/// Lowest weighted Gini impurity of the children first.
87#[derive(Default)]
88pub struct GiniIndex;
89
90impl Heuristic for GiniIndex {
91    fn compute(&self, cover: &mut Cover, candidates: &mut Vec<usize>) -> Vec<f64> {
92        self.compute_with_scorer(0.0, cover, candidates, Box::new(gini_index), true)
93    }
94}
95
96/// Highest information gain first.
97#[derive(Default)]
98pub struct InformationGain;
99
100impl Heuristic for InformationGain {
101    fn compute(&self, cover: &mut Cover, candidates: &mut Vec<usize>) -> Vec<f64> {
102        if candidates.is_empty() {
103            return vec![];
104        }
105
106        let parent_distribution = cover.labels_count();
107        let parent_entropy = entropy(&parent_distribution);
108
109        self.compute_with_scorer(
110            parent_entropy,
111            cover,
112            candidates,
113            Box::new(information_gain),
114            false,
115        )
116    }
117}
118
119/// Lowest weighted entropy of the children first.
120#[derive(Default)]
121pub struct WeightedEntropy;
122
123impl Heuristic for WeightedEntropy {
124    fn compute(&self, cover: &mut Cover, candidates: &mut Vec<usize>) -> Vec<f64> {
125        self.compute_with_scorer(0.0, cover, candidates, Box::new(weighted_entropy), true)
126    }
127}
dtrees_rs/algorithms/common/heuristics/mod.rs

dtrees_rs/algorithms/common/heuristics/
mod.rs
