1use 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
12const ZERO_THRESHOLD: f64 = 1e-35_f32 as f64;
14
15impl Model {
16 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 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 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 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 pub fn leaf_index(&self, row: &[f64]) -> usize {
93 leaf_index(self, row)
94 }
95}
96
97fn 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#[inline]
118fn read_value(value: f64) -> f64 {
119 if value.abs() <= ZERO_THRESHOLD {
120 0.0
121 } else {
122 value
123 }
124}
125
126#[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#[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}