silvic_core/
tree.rs

1//! A regression tree in LightGBM 3.3.5's own array layout (`include/LightGBM/tree.h`), so `Tree::ToString` output
2//! is a field-by-field dump.
3//!
4//! Splits mirror `Tree::Split`: the split leaf keeps its index as the left child, the right child becomes leaf
5//! `num_leaves`, and the new internal node is `num_leaves - 1`. A child index `>= 0` is an internal node and `~leaf`
6//! a leaf. The root starts with value and weight 0, which is why LightGBM writes `internal_weight=0` for node 0.
7
8/// `decision_type` bit 0: categorical split.
9pub const CATEGORICAL_MASK: i8 = 1;
10/// `decision_type` bit 1: missing values go left.
11pub const DEFAULT_LEFT_MASK: i8 = 2;
12/// LightGBM `kZeroThreshold` (`1e-35f`, a float widened to double): smaller magnitudes are written as 0.
13const ZERO_THRESHOLD: f64 = 1e-35_f32 as f64;
14
15/// LightGBM `MissingType`, stored in `decision_type` bits 2-3.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum MissingType {
18    None = 0,
19    Zero = 1,
20    NaN = 2,
21}
22
23impl MissingType {
24    pub fn from_decision_type(decision_type: i8) -> MissingType {
25        match (decision_type >> 2) & 3 {
26            1 => MissingType::Zero,
27            2 => MissingType::NaN,
28            _ => MissingType::None,
29        }
30    }
31}
32
33/// The statistics of the two children a split creates.
34#[derive(Debug, Clone, Copy, PartialEq)]
35pub struct SplitChildren {
36    pub left_value: f64,
37    pub right_value: f64,
38    pub left_count: i64,
39    pub right_count: i64,
40    /// Hessian sums.
41    pub left_weight: f64,
42    pub right_weight: f64,
43    pub gain: f32,
44}
45
46#[derive(Debug, Clone, PartialEq)]
47pub struct Tree {
48    pub num_leaves: usize,
49    pub num_cat: usize,
50    /// Input (real) feature index of each internal node.
51    pub split_feature: Vec<i32>,
52    pub split_gain: Vec<f32>,
53    /// Numerical: the value `<=` which goes left. Categorical: index into `cat_boundaries`.
54    pub threshold: Vec<f64>,
55    pub decision_type: Vec<i8>,
56    pub left_child: Vec<i32>,
57    pub right_child: Vec<i32>,
58    pub leaf_value: Vec<f64>,
59    pub leaf_weight: Vec<f64>,
60    pub leaf_count: Vec<i64>,
61    pub internal_value: Vec<f64>,
62    pub internal_weight: Vec<f64>,
63    pub internal_count: Vec<i64>,
64    /// `num_cat + 1` offsets into `cat_threshold`.
65    pub cat_boundaries: Vec<i32>,
66    /// Bitsets of category values going left, 32 per word.
67    pub cat_threshold: Vec<u32>,
68    pub shrinkage: f64,
69    /// Training only: each leaf's parent node (`-1` for the root).
70    pub leaf_parent: Vec<i32>,
71    /// Training only: each leaf's depth (root 0).
72    pub leaf_depth: Vec<i32>,
73}
74
75impl Default for Tree {
76    fn default() -> Self {
77        Tree::new()
78    }
79}
80
81impl Tree {
82    /// A single leaf of value 0.
83    pub fn new() -> Tree {
84        Tree {
85            num_leaves: 1,
86            num_cat: 0,
87            split_feature: Vec::new(),
88            split_gain: Vec::new(),
89            threshold: Vec::new(),
90            decision_type: Vec::new(),
91            left_child: Vec::new(),
92            right_child: Vec::new(),
93            leaf_value: vec![0.0],
94            leaf_weight: vec![0.0],
95            leaf_count: vec![0],
96            internal_value: Vec::new(),
97            internal_weight: Vec::new(),
98            internal_count: Vec::new(),
99            cat_boundaries: vec![0],
100            cat_threshold: Vec::new(),
101            shrinkage: 1.0,
102            leaf_parent: vec![-1],
103            leaf_depth: vec![0],
104        }
105    }
106
107    /// `Tree::Split` for a numerical feature; returns the new (right) leaf's index.
108    pub fn split_numerical(
109        &mut self,
110        leaf: usize,
111        feature: usize,
112        threshold: f64,
113        default_left: bool,
114        missing: MissingType,
115        children: &SplitChildren,
116    ) -> usize {
117        let mut decision_type = (missing as i8) << 2;
118        if default_left {
119            decision_type |= DEFAULT_LEFT_MASK;
120        }
121        self.split(leaf, feature, threshold, decision_type, children)
122    }
123
124    /// `Tree::SplitCategorical`: `categories` are the category values going left.
125    pub fn split_categorical(
126        &mut self,
127        leaf: usize,
128        feature: usize,
129        categories: &[i32],
130        missing: MissingType,
131        children: &SplitChildren,
132    ) -> usize {
133        let bitset = to_bitset(categories);
134        let cat_index = self.num_cat as f64;
135        let last = *self.cat_boundaries.last().expect("cat_boundaries starts with 0");
136        self.cat_boundaries.push(last + bitset.len() as i32);
137        self.cat_threshold.extend_from_slice(&bitset);
138        self.num_cat += 1;
139        let decision_type = CATEGORICAL_MASK | ((missing as i8) << 2);
140        self.split(leaf, feature, cat_index, decision_type, children)
141    }
142
143    fn split(&mut self, leaf: usize, feature: usize, threshold: f64, decision_type: i8, c: &SplitChildren) -> usize {
144        let node = self.num_leaves - 1;
145        let new_leaf = self.num_leaves;
146        let parent = self.leaf_parent[leaf];
147        if parent >= 0 {
148            let parent = parent as usize;
149            if self.left_child[parent] == !(leaf as i32) {
150                self.left_child[parent] = node as i32;
151            } else {
152                self.right_child[parent] = node as i32;
153            }
154        }
155        self.split_feature.push(feature as i32);
156        self.split_gain.push(c.gain);
157        self.threshold.push(threshold);
158        self.decision_type.push(decision_type);
159        self.left_child.push(!(leaf as i32));
160        self.right_child.push(!(new_leaf as i32));
161        self.internal_weight.push(self.leaf_weight[leaf]);
162        self.internal_value.push(self.leaf_value[leaf]);
163        self.internal_count.push(c.left_count + c.right_count);
164        self.leaf_value[leaf] = nan_to_zero(c.left_value);
165        self.leaf_weight[leaf] = c.left_weight;
166        self.leaf_count[leaf] = c.left_count;
167        self.leaf_value.push(nan_to_zero(c.right_value));
168        self.leaf_weight.push(c.right_weight);
169        self.leaf_count.push(c.right_count);
170        self.leaf_parent[leaf] = node as i32;
171        self.leaf_parent.push(node as i32);
172        let depth = self.leaf_depth[leaf] + 1;
173        self.leaf_depth[leaf] = depth;
174        self.leaf_depth.push(depth);
175        self.num_leaves += 1;
176        new_leaf
177    }
178
179    /// `Tree::Shrinkage`.
180    pub fn shrink(&mut self, rate: f64) {
181        for value in self.leaf_value.iter_mut().chain(self.internal_value.iter_mut()) {
182            *value = round_to_zero(*value * rate);
183        }
184        self.shrinkage *= rate;
185    }
186
187    /// `Tree::AddBias`: folds the initial score into the first tree, whose shrinkage LightGBM then writes as 1.
188    pub fn add_bias(&mut self, bias: f64) {
189        for value in self.leaf_value.iter_mut().chain(self.internal_value.iter_mut()) {
190            *value = round_to_zero(*value + bias);
191        }
192        self.shrinkage = 1.0;
193    }
194
195    /// `Tree::AsConstantTree`.
196    pub fn as_constant(&mut self, value: f64) {
197        *self = Tree::new();
198        self.leaf_value[0] = value;
199    }
200
201    /// `Tree::SetLeafOutput`.
202    pub fn set_leaf_output(&mut self, leaf: usize, value: f64) {
203        self.leaf_value[leaf] = round_to_zero(value);
204    }
205
206    /// Depth of the deepest leaf.
207    pub fn max_depth(&self) -> i32 {
208        self.leaf_depth.iter().copied().max().unwrap_or(0)
209    }
210}
211
212fn nan_to_zero(value: f64) -> f64 {
213    if value.is_nan() {
214        0.0
215    } else {
216        value
217    }
218}
219
220/// LightGBM `MaybeRoundToZero`.
221fn round_to_zero(value: f64) -> f64 {
222    if value.abs() > ZERO_THRESHOLD {
223        value
224    } else {
225        0.0
226    }
227}
228
229/// LightGBM `Common::ConstructBitset`: one bit per non-negative value, as few words as the largest value needs.
230pub fn to_bitset(values: &[i32]) -> Vec<u32> {
231    let max = values.iter().copied().filter(|&v| v >= 0).max().unwrap_or(0) as usize;
232    let mut words = vec![0u32; max / 32 + 1];
233    for &value in values.iter().filter(|&&v| v >= 0) {
234        words[value as usize / 32] |= 1 << (value as usize % 32);
235    }
236    words
237}
238
239/// LightGBM `Common::FindInBitset`.
240pub fn in_bitset(words: &[u32], value: i32) -> bool {
241    if value < 0 {
242        return false;
243    }
244    let word = value as usize / 32;
245    word < words.len() && (words[word] >> (value as usize % 32)) & 1 == 1
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    fn children(gain: f32) -> SplitChildren {
253        SplitChildren {
254            left_value: -1.0,
255            right_value: 1.0,
256            left_count: 10,
257            right_count: 20,
258            left_weight: 10.0,
259            right_weight: 20.0,
260            gain,
261        }
262    }
263
264    #[test]
265    fn splits_follow_lightgbm_node_numbering() {
266        let mut tree = Tree::new();
267        let right = tree.split_numerical(0, 3, 0.5, true, MissingType::None, &children(5.0));
268        assert_eq!(right, 1);
269        let right2 = tree.split_numerical(1, 1, 2.0, false, MissingType::NaN, &children(2.0));
270        assert_eq!(right2, 2);
271        tree.split_numerical(0, 0, -1.0, true, MissingType::None, &children(1.0));
272        // Node 0 splits leaf 0 (left) / leaf 1 (right); node 1 replaces leaf 1; node 2 replaces leaf 0.
273        assert_eq!(tree.left_child, vec![2, -2, -1]);
274        assert_eq!(tree.right_child, vec![1, -3, -4]);
275        assert_eq!(tree.decision_type, vec![2, 8, 2]);
276        assert_eq!(tree.internal_weight[0], 0.0);
277        assert_eq!(tree.leaf_depth, vec![2, 2, 2, 2]);
278    }
279
280    #[test]
281    fn categorical_split_records_bitset() {
282        let mut tree = Tree::new();
283        tree.split_categorical(0, 2, &[1, 3, 33], MissingType::None, &children(1.0));
284        assert_eq!(tree.cat_boundaries, vec![0, 2]);
285        assert_eq!(tree.cat_threshold, vec![0b1010, 0b10]);
286        assert_eq!(tree.threshold, vec![0.0]);
287        assert_eq!(tree.decision_type, vec![1]);
288        assert!(in_bitset(&tree.cat_threshold, 33));
289        assert!(!in_bitset(&tree.cat_threshold, 2));
290        assert!(!in_bitset(&tree.cat_threshold, -1));
291    }
292
293    #[test]
294    fn bias_resets_shrinkage() {
295        let mut tree = Tree::new();
296        tree.split_numerical(0, 0, 0.0, true, MissingType::None, &children(1.0));
297        tree.shrink(0.1);
298        tree.add_bias(2.0);
299        assert_eq!(tree.shrinkage, 1.0);
300        assert_eq!(tree.leaf_value, vec![-0.1 + 2.0, 0.1 + 2.0]);
301        assert_eq!(tree.internal_value, vec![2.0]);
302    }
303}
silvic_core/tree.rs

silvic_core/
tree.rs
