Skip to main content

silvic_core/model/
predict.rs

1//! Prediction as LightGBM 3.3.5 predicts a dense row: the C API's row reader (`RowFunctionFromDenseMatric`) keeps a
2//! value only when `|v| > kZeroThreshold` or it is NaN, so tiny values read as 0; each tree then walks
3//! `Tree::GetLeaf` with `NumericalDecision` and `CategoricalDecision`, and `GBDT::PredictRaw` sums the leaf values in
4//! tree order in f64. `GBDT::Predict` applies the objective's `ConvertOutput` (the sigmoid for binary).
5
6use rayon::prelude::*;
7
8use crate::data::{DenseData, DenseMatrix, Layout};
9use crate::model::{ImportanceType, Model, ModelObjective};
10use crate::tree::{in_bitset, MissingType, Tree, CATEGORICAL_MASK, DEFAULT_LEFT_MASK};
11
12/// LightGBM `kZeroThreshold`, which is the float `1e-35f` widened to double.
13const ZERO_THRESHOLD: f64 = 1e-35_f32 as f64;
14
15impl Model {
16    /// Sum of the leaf values `row` reaches in the first `num_trees` trees (all when `None`).
17    ///
18    /// `row` holds at least `num_features()` values; NaN is a missing value.
19    pub fn predict_raw_row(&self, row: &[f64], num_trees: Option<usize>) -> f64 {
20        self.used_trees(num_trees)
21            .iter()
22            .fold(0.0, |score, tree| score + tree.leaf_value[leaf_index(tree, row)])
23    }
24
25    /// Predictions for every row, parallel over rows; `raw_score = false` applies the objective's transform.
26    ///
27    /// # Panics
28    /// When `matrix` has fewer columns than the model has features.
29    pub fn predict(&self, matrix: &DenseMatrix<'_>, raw_score: bool, num_trees: Option<usize>) -> Vec<f64> {
30        assert!(
31            matrix.cols >= self.num_features(),
32            "the matrix has {} columns, the model reads {} features",
33            matrix.cols,
34            self.num_features()
35        );
36        let convert = |raw: f64| if raw_score { raw } else { self.convert_output(raw) };
37        if let (DenseData::F64(values), Layout::RowMajor) = (matrix.data, matrix.layout) {
38            if matrix.cols > 0 {
39                return values[..matrix.rows * matrix.cols]
40                    .par_chunks_exact(matrix.cols)
41                    .map(|row| convert(self.predict_raw_row(row, num_trees)))
42                    .collect();
43            }
44        }
45        (0..matrix.rows)
46            .into_par_iter()
47            .map_init(
48                || vec![0.0; matrix.cols],
49                |row, index| {
50                    for (col, value) in row.iter_mut().enumerate() {
51                        *value = matrix.get(index, col);
52                    }
53                    convert(self.predict_raw_row(row, num_trees))
54                },
55            )
56            .collect()
57    }
58
59    /// LightGBM `GBDT::FeatureImportance`: per feature, the number of splits (`Split`) or their summed gain (`Gain`),
60    /// counting only splits with positive gain.
61    pub fn feature_importance(&self, kind: ImportanceType) -> Vec<f64> {
62        let mut importance = vec![0.0; self.num_features()];
63        for tree in &self.trees {
64            let nodes = tree.num_leaves.saturating_sub(1);
65            for (&feature, &gain) in tree.split_feature[..nodes].iter().zip(&tree.split_gain[..nodes]) {
66                if gain > 0.0 {
67                    importance[feature as usize] += match kind {
68                        ImportanceType::Split => 1.0,
69                        ImportanceType::Gain => f64::from(gain),
70                    };
71                }
72            }
73        }
74        importance
75    }
76
77    fn used_trees(&self, num_trees: Option<usize>) -> &[Tree] {
78        &self.trees[..num_trees.map_or(self.trees.len(), |n| n.min(self.trees.len()))]
79    }
80
81    /// The objective's `ConvertOutput`.
82    fn convert_output(&self, raw: f64) -> f64 {
83        match self.objective {
84            ModelObjective::Binary { sigmoid } => 1.0 / (1.0 + (-sigmoid * raw).exp()),
85            ModelObjective::Regression | ModelObjective::LambdaRank => raw,
86        }
87    }
88}
89
90impl Tree {
91    /// The leaf `row` reaches, as `Model::predict_raw_row` walks it: `row` holds raw feature values (NaN missing).
92    pub fn leaf_index(&self, row: &[f64]) -> usize {
93        leaf_index(self, row)
94    }
95}
96
97/// `Tree::Predict`'s leaf: a one-leaf tree has no nodes to walk. Like `GetLeaf`, a tree without categorical splits
98/// takes every node as numerical.
99fn leaf_index(tree: &Tree, row: &[f64]) -> usize {
100    if tree.num_leaves <= 1 {
101        return 0;
102    }
103    let mut node = 0;
104    while node >= 0 {
105        let index = node as usize;
106        let value = read_value(row[tree.split_feature[index] as usize]);
107        node = if tree.num_cat > 0 && tree.decision_type[index] & CATEGORICAL_MASK != 0 {
108            categorical_decision(tree, index, value)
109        } else {
110            numerical_decision(tree, index, value)
111        };
112    }
113    !node as usize
114}
115
116/// The dense-row reader's view of a value: magnitudes up to `kZeroThreshold` are absent, that is 0.
117#[inline]
118fn read_value(value: f64) -> f64 {
119    if value.abs() <= ZERO_THRESHOLD {
120        0.0
121    } else {
122        value
123    }
124}
125
126/// `Tree::NumericalDecision`: NaN is 0 unless the split has a NaN direction; a missing value (NaN, or `IsZero` for
127/// `MissingType::Zero`) takes the default direction.
128#[inline]
129fn numerical_decision(tree: &Tree, node: usize, value: f64) -> i32 {
130    let decision_type = tree.decision_type[node];
131    let missing_type = MissingType::from_decision_type(decision_type);
132    let value = if value.is_nan() && missing_type != MissingType::NaN {
133        0.0
134    } else {
135        value
136    };
137    let is_missing = match missing_type {
138        MissingType::None => false,
139        MissingType::Zero => (-ZERO_THRESHOLD..=ZERO_THRESHOLD).contains(&value),
140        MissingType::NaN => value.is_nan(),
141    };
142    let left = if is_missing {
143        decision_type & DEFAULT_LEFT_MASK != 0
144    } else {
145        value <= tree.threshold[node]
146    };
147    if left {
148        tree.left_child[node]
149    } else {
150        tree.right_child[node]
151    }
152}
153
154/// `Tree::CategoricalDecision`: NaN and negative values go right, others left when their bit is set. The cast
155/// truncates toward zero like `static_cast<int>`; out of `int` range C++ is undefined, Rust saturates, and the
156/// result goes right either way (negative, or past every bitset).
157#[inline]
158fn categorical_decision(tree: &Tree, node: usize, value: f64) -> i32 {
159    if value.is_nan() || (value as i32) < 0 {
160        return tree.right_child[node];
161    }
162    let cat_index = tree.threshold[node] as usize;
163    let start = tree.cat_boundaries[cat_index] as usize;
164    let end = tree.cat_boundaries[cat_index + 1] as usize;
165    if in_bitset(&tree.cat_threshold[start..end], value as i32) {
166        tree.left_child[node]
167    } else {
168        tree.right_child[node]
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175    use crate::tree::SplitChildren;
176
177    fn children(left_value: f64, right_value: f64) -> SplitChildren {
178        SplitChildren {
179            left_value,
180            right_value,
181            left_count: 1,
182            right_count: 1,
183            left_weight: 1.0,
184            right_weight: 1.0,
185            gain: 1.0,
186        }
187    }
188
189    fn model(tree: Tree) -> Model {
190        Model {
191            objective: ModelObjective::Binary { sigmoid: 2.0 },
192            feature_names: vec!["x".into()],
193            feature_infos: vec!["[-1:1]".into()],
194            trees: vec![tree],
195            parameters: Vec::new(),
196            pandas_categorical: None,
197        }
198    }
199
200    #[test]
201    fn missing_types_route_like_lightgbm() {
202        let cases = [
203            (MissingType::None, true, f64::NAN, -1.0),
204            (MissingType::None, false, f64::NAN, -1.0),
205            (MissingType::Zero, false, f64::NAN, 1.0),
206            (MissingType::Zero, true, 1e-36, -1.0),
207            (MissingType::Zero, false, 0.0, 1.0),
208            (MissingType::NaN, true, f64::NAN, -1.0),
209            (MissingType::NaN, false, f64::NAN, 1.0),
210            (MissingType::NaN, false, -1e-36, -1.0),
211        ];
212        for (missing, default_left, value, expected) in cases {
213            let mut tree = Tree::new();
214            tree.split_numerical(0, 0, 0.0, default_left, missing, &children(-1.0, 1.0));
215            let model = model(tree);
216            let case = format!("{missing:?} default_left={default_left} value={value}");
217            assert_eq!(model.predict_raw_row(&[value], None), expected, "{case}");
218        }
219    }
220
221    #[test]
222    fn categorical_values_outside_the_bitset_go_right() {
223        let mut tree = Tree::new();
224        tree.split_categorical(0, 0, &[0, 2, 40], MissingType::None, &children(-1.0, 1.0));
225        let model = model(tree);
226        for (value, expected) in [
227            (0.0, -1.0),
228            (-0.5, -1.0),
229            (2.9, -1.0),
230            (40.0, -1.0),
231            (1.0, 1.0),
232            (-1.0, 1.0),
233            (f64::NAN, 1.0),
234            (64.0, 1.0),
235            (1e10, 1.0),
236            (-1e10, 1.0),
237        ] {
238            assert_eq!(model.predict_raw_row(&[value], None), expected, "category {value}");
239        }
240    }
241
242    #[test]
243    fn predict_matches_rows_in_every_layout() {
244        let mut tree = Tree::new();
245        tree.split_numerical(0, 0, 0.5, false, MissingType::NaN, &children(-1.0, 1.0));
246        let mut model = model(tree);
247        model.feature_names.push("y".into());
248        model.feature_infos.push("none".into());
249        let row_major = [0.0, 9.0, 1.0, 9.0, f64::NAN, 9.0];
250        let column_major = [0.0f32, 1.0, f32::NAN, 9.0, 9.0, 9.0];
251        let raw = vec![-1.0, 1.0, 1.0];
252        let sigmoid: Vec<f64> = raw.iter().map(|r: &f64| 1.0 / (1.0 + (-2.0 * r).exp())).collect();
253        let rows = DenseMatrix {
254            data: DenseData::F64(&row_major),
255            rows: 3,
256            cols: 2,
257            layout: Layout::RowMajor,
258        };
259        assert_eq!(model.predict(&rows, true, None), raw);
260        assert_eq!(model.predict(&rows, false, None), sigmoid);
261        let cols = DenseMatrix {
262            data: DenseData::F32(&column_major),
263            rows: 3,
264            cols: 2,
265            layout: Layout::ColumnMajor,
266        };
267        assert_eq!(model.predict(&cols, true, None), raw);
268        assert_eq!(model.predict(&cols, true, Some(0)), vec![0.0; 3]);
269    }
270
271    #[test]
272    fn importance_counts_positive_gain_splits() {
273        let mut tree = Tree::new();
274        let mut split = children(0.0, 0.0);
275        split.gain = 2.5;
276        tree.split_numerical(0, 0, 0.0, false, MissingType::None, &split);
277        split.gain = 0.0;
278        tree.split_numerical(0, 0, 0.0, false, MissingType::None, &split);
279        let model = model(tree);
280        assert_eq!(model.feature_importance(ImportanceType::Split), vec![1.0]);
281        assert_eq!(model.feature_importance(ImportanceType::Gain), vec![2.5]);
282    }
283}