forust 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of forust might be problematic. Click here for more details.

Files changed (43) hide show
  1. {forust-0.2.0 → forust-0.2.2}/Cargo.toml +2 -2
  2. {forust-0.2.0 → forust-0.2.2}/PKG-INFO +16 -2
  3. {forust-0.2.0 → forust-0.2.2}/README.md +15 -1
  4. {forust-0.2.0 → forust-0.2.2}/forust/__init__.py +40 -0
  5. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/.github/workflows/CI.yml +3 -3
  6. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/Cargo.toml +3 -3
  7. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/README.md +15 -1
  8. forust-0.2.2/local_dependencies/forust-ml/dist/forust-0.2.2-cp311-cp311-macosx_10_7_x86_64.whl +0 -0
  9. forust-0.2.2/local_dependencies/forust-ml/dist/forust-0.2.2.tar.gz +0 -0
  10. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/rs-example.md +1 -1
  11. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/data.rs +42 -1
  12. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/gradientbooster.rs +81 -0
  13. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/node.rs +55 -52
  14. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/partial_dependence.rs +30 -40
  15. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/splitter.rs +6 -28
  16. forust-0.2.2/local_dependencies/forust-ml/src/tree.rs +395 -0
  17. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/utils.rs +11 -0
  18. {forust-0.2.0 → forust-0.2.2}/src/lib.rs +16 -0
  19. {forust-0.2.0 → forust-0.2.2}/tests/test_booster.py +46 -0
  20. forust-0.2.0/local_dependencies/forust-ml/dist/forust-0.2.0-cp310-cp310-macosx_10_7_x86_64.whl +0 -0
  21. forust-0.2.0/local_dependencies/forust-ml/dist/forust-0.2.0.tar.gz +0 -0
  22. forust-0.2.0/local_dependencies/forust-ml/src/tree.rs +0 -299
  23. {forust-0.2.0 → forust-0.2.2}/.gitignore +0 -0
  24. {forust-0.2.0 → forust-0.2.2}/LICENSE +0 -0
  25. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/.gitignore +0 -0
  26. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/.pre-commit-config.yaml +0 -0
  27. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/LICENSE +0 -0
  28. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/benches/forust_benchmarks.rs +0 -0
  29. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/resources/pdp_plot_age.png +0 -0
  30. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/resources/pdp_plot_age_mono.png +0 -0
  31. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/resources/titanic.csv +0 -0
  32. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/resources/tree-image-crop.png +0 -0
  33. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/scripts/make_resources.py +0 -0
  34. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/scripts/run-python-tests.ps1 +0 -0
  35. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/scripts/run-python-tests.sh +0 -0
  36. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/binning.rs +0 -0
  37. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/constraints.rs +0 -0
  38. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/errors.rs +0 -0
  39. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/histogram.rs +0 -0
  40. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/lib.rs +0 -0
  41. {forust-0.2.0 → forust-0.2.2}/local_dependencies/forust-ml/src/objective.rs +0 -0
  42. {forust-0.2.0 → forust-0.2.2}/pyproject.toml +0 -0
  43. {forust-0.2.0 → forust-0.2.2}/scratch.py +0 -0
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "py-forust"
3
- version = "0.2.0"
3
+ version = "0.2.2"
4
4
  edition = "2021"
5
5
 
6
6
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -10,6 +10,6 @@ crate-type = ["cdylib"]
10
10
 
11
11
  [dependencies]
12
12
  pyo3 = { version = "0.17", features = ["extension-module"] }
13
- forust-ml = { version="0.2.0", path= "local_dependencies/forust-ml" }
13
+ forust-ml = { version="0.2.2", path= "local_dependencies/forust-ml" }
14
14
  numpy = "0.17.2"
15
15
  ndarray = "0.15.1"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: forust
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Classifier: Programming Language :: Rust
5
5
  Classifier: Programming Language :: Python :: Implementation :: CPython
6
6
  Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -47,7 +47,7 @@ pip install forust
47
47
 
48
48
  To use in a rust project add the following to your Cargo.toml file.
49
49
  ```toml
50
- forust-ml = "0.2.0"
50
+ forust-ml = "0.2.2"
51
51
  ```
52
52
 
53
53
  ## Usage
@@ -105,6 +105,13 @@ model.fit(X, y)
105
105
  # Predict on data
106
106
  model.predict(X.head())
107
107
  # array([-1.94919663, 2.25863229, 0.32963671, 2.48732194, -3.00371813])
108
+
109
+ # predict contributions
110
+ model.predict_contributions(X.head())
111
+ # array([[-0.63014213, 0.33880048, -0.16520798, -0.07798772, -0.85083578,
112
+ # -1.07720813],
113
+ # [ 1.05406709, 0.08825999, 0.21662544, -0.12083538, 0.35209258,
114
+ # -1.07720813],
108
115
  ```
109
116
 
110
117
  The `fit` method accepts the following arguments.
@@ -123,6 +130,13 @@ The predict method accepts the following arguments.
123
130
  passed, the `parallel` attribute of the booster will be used.
124
131
  Defaults to `None`.
125
132
 
133
+ The `predict_contributions` method will predict with the fitted booster on new data, returning the feature contribution matrix. The last column is the bias term.
134
+ - `X` ***(FrameLike)***: Either a pandas DataFrame, or a 2 dimensional numpy array, with numeric data.
135
+ - `parallel` ***(Optional[bool], optional)***: Optionally specify if the predict
136
+ function should run in parallel on multiple threads. If `None` is
137
+ passed, the `parallel` attribute of the booster will be used.
138
+ Defaults to `None`.
139
+
126
140
  ### Inspecting the Model
127
141
 
128
142
  Once the booster has been fit, each individual tree structure can be retrieved in text form, using the `text_dump` method. This method returns a list, the same length as the number of trees in the model.
@@ -26,7 +26,7 @@ pip install forust
26
26
 
27
27
  To use in a rust project add the following to your Cargo.toml file.
28
28
  ```toml
29
- forust-ml = "0.2.0"
29
+ forust-ml = "0.2.2"
30
30
  ```
31
31
 
32
32
  ## Usage
@@ -84,6 +84,13 @@ model.fit(X, y)
84
84
  # Predict on data
85
85
  model.predict(X.head())
86
86
  # array([-1.94919663, 2.25863229, 0.32963671, 2.48732194, -3.00371813])
87
+
88
+ # predict contributions
89
+ model.predict_contributions(X.head())
90
+ # array([[-0.63014213, 0.33880048, -0.16520798, -0.07798772, -0.85083578,
91
+ # -1.07720813],
92
+ # [ 1.05406709, 0.08825999, 0.21662544, -0.12083538, 0.35209258,
93
+ # -1.07720813],
87
94
  ```
88
95
 
89
96
  The `fit` method accepts the following arguments.
@@ -102,6 +109,13 @@ The predict method accepts the following arguments.
102
109
  passed, the `parallel` attribute of the booster will be used.
103
110
  Defaults to `None`.
104
111
 
112
+ The `predict_contributions` method will predict with the fitted booster on new data, returning the feature contribution matrix. The last column is the bias term.
113
+ - `X` ***(FrameLike)***: Either a pandas DataFrame, or a 2 dimensional numpy array, with numeric data.
114
+ - `parallel` ***(Optional[bool], optional)***: Optionally specify if the predict
115
+ function should run in parallel on multiple threads. If `None` is
116
+ passed, the `parallel` attribute of the booster will be used.
117
+ Defaults to `None`.
118
+
105
119
  ### Inspecting the Model
106
120
 
107
121
  Once the booster has been fit, each individual tree structure can be retrieved in text form, using the `text_dump` method. This method returns a list, the same length as the number of trees in the model.
@@ -35,6 +35,15 @@ class BoosterType:
35
35
  ) -> np.ndarray:
36
36
  raise NotImplementedError()
37
37
 
38
+ def predict_contributions(
39
+ self,
40
+ flat_data: np.ndarray,
41
+ rows: int,
42
+ cols: int,
43
+ parallel: bool = True,
44
+ ) -> np.ndarray:
45
+ raise NotImplementedError
46
+
38
47
  def value_partial_dependence(
39
48
  self,
40
49
  feature: int,
@@ -238,6 +247,37 @@ class GradientBooster:
238
247
  parallel=parallel_,
239
248
  )
240
249
 
250
+ def predict_contributions(
251
+ self, X: FrameLike, parallel: Union[bool, None] = None
252
+ ) -> np.ndarray:
253
+ """Predict with the fitted booster on new data, returning the feature
254
+ contribution matrix. The last column is the bias term.
255
+
256
+ Args:
257
+ X (FrameLike): Either a pandas DataFrame, or a 2 dimensional numpy array.
258
+ parallel (Union[bool, None], optional): Optionally specify if the predict
259
+ function should run in parallel on multiple threads. If `None` is
260
+ passed, the `parallel` attribute of the booster will be used.
261
+ Defaults to `None`.
262
+
263
+ Returns:
264
+ np.ndarray: Returns a numpy array of the predictions.
265
+ """
266
+ X_ = X.to_numpy() if isinstance(X, pd.DataFrame) else X
267
+ if not np.issubdtype(X_.dtype, "float64"):
268
+ X_ = X_.astype(dtype="float64", copy=False)
269
+
270
+ parallel_ = self.parallel if parallel is None else parallel
271
+ flat_data = X_.ravel(order="F")
272
+ rows, cols = X_.shape
273
+ contributions = self.booster.predict_contributions(
274
+ flat_data=flat_data,
275
+ rows=rows,
276
+ cols=cols,
277
+ parallel=parallel_,
278
+ )
279
+ return np.reshape(contributions, (X_.shape[0], X_.shape[1] + 1))
280
+
241
281
  def partial_dependence(self, X: FrameLike, feature: Union[str, int]) -> np.ndarray:
242
282
  """Calculate the partial dependence values of a feature. For each unique
243
283
  value of the feature, this gives the estimate of the predicted value for that
@@ -5,7 +5,7 @@ jobs:
5
5
  windows-build-tests:
6
6
  strategy:
7
7
  matrix:
8
- pyversion: ["3.8", "3.9", "3.10"]
8
+ pyversion: ["3.8", "3.9", "3.10", "3.11"]
9
9
  runs-on: "windows-latest"
10
10
  steps:
11
11
  - uses: actions/checkout@v2
@@ -58,7 +58,7 @@ jobs:
58
58
  macos-build-test:
59
59
  strategy:
60
60
  matrix:
61
- pyversion: ["3.8", "3.9", "3.10"]
61
+ pyversion: ["3.8", "3.9", "3.10", "3.11"]
62
62
  runs-on: "macos-latest"
63
63
  steps:
64
64
  - uses: actions/checkout@v2
@@ -113,7 +113,7 @@ jobs:
113
113
  runs-on: ubuntu-latest
114
114
  strategy:
115
115
  matrix:
116
- pyversion: ["3.8", "3.9", "3.10"]
116
+ pyversion: ["3.8", "3.9", "3.10", "3.11"]
117
117
  # target: [x86_64, i686]
118
118
  steps:
119
119
  - uses: actions/checkout@v2
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "forust-ml"
3
- version = "0.2.0"
3
+ version = "0.2.2"
4
4
  edition = "2021"
5
5
  authors = ["James Inlow <james.d.inlow@gmail.com>"]
6
6
  homepage = "https://github.com/jinlow/forust"
@@ -12,8 +12,8 @@ repository = "https://github.com/jinlow/forust"
12
12
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
13
13
 
14
14
  [dependencies]
15
- rayon = "1.5"
16
- thiserror = "1.0.31"
15
+ rayon = "1.7"
16
+ thiserror = "1.0"
17
17
  serde_json = {version = "1.0", features = ["float_roundtrip"] }
18
18
  serde = { version = "1.0", features = ["derive"] }
19
19
 
@@ -26,7 +26,7 @@ pip install forust
26
26
 
27
27
  To use in a rust project add the following to your Cargo.toml file.
28
28
  ```toml
29
- forust-ml = "0.2.0"
29
+ forust-ml = "0.2.2"
30
30
  ```
31
31
 
32
32
  ## Usage
@@ -84,6 +84,13 @@ model.fit(X, y)
84
84
  # Predict on data
85
85
  model.predict(X.head())
86
86
  # array([-1.94919663, 2.25863229, 0.32963671, 2.48732194, -3.00371813])
87
+
88
+ # predict contributions
89
+ model.predict_contributions(X.head())
90
+ # array([[-0.63014213, 0.33880048, -0.16520798, -0.07798772, -0.85083578,
91
+ # -1.07720813],
92
+ # [ 1.05406709, 0.08825999, 0.21662544, -0.12083538, 0.35209258,
93
+ # -1.07720813],
87
94
  ```
88
95
 
89
96
  The `fit` method accepts the following arguments.
@@ -102,6 +109,13 @@ The predict method accepts the following arguments.
102
109
  passed, the `parallel` attribute of the booster will be used.
103
110
  Defaults to `None`.
104
111
 
112
+ The `predict_contributions` method will predict with the fitted booster on new data, returning the feature contribution matrix. The last column is the bias term.
113
+ - `X` ***(FrameLike)***: Either a pandas DataFrame, or a 2 dimensional numpy array, with numeric data.
114
+ - `parallel` ***(Optional[bool], optional)***: Optionally specify if the predict
115
+ function should run in parallel on multiple threads. If `None` is
116
+ passed, the `parallel` attribute of the booster will be used.
117
+ Defaults to `None`.
118
+
105
119
  ### Inspecting the Model
106
120
 
107
121
  Once the booster has been fit, each individual tree structure can be retrieved in text form, using the `text_dump` method. This method returns a list, the same length as the number of trees in the model.
@@ -3,7 +3,7 @@
3
3
  To run this example, add the following code to your `Cargo.toml` file.
4
4
  ```toml
5
5
  [dependencies]
6
- forust-ml = "0.2.0"
6
+ forust-ml = "0.2.2"
7
7
  polars = "0.24"
8
8
  reqwest = { version = "0.11", features = ["blocking"] }
9
9
  ```
@@ -93,6 +93,7 @@ pub struct Matrix<'a, T> {
93
93
  }
94
94
 
95
95
  impl<'a, T> Matrix<'a, T> {
96
+ // Defaults to column major
96
97
  pub fn new(data: &'a [T], rows: usize, cols: usize) -> Self {
97
98
  Matrix {
98
99
  data,
@@ -118,6 +119,14 @@ impl<'a, T> Matrix<'a, T> {
118
119
  idx
119
120
  }
120
121
 
122
+ /// Get access to a row of the data, as an iterator.
123
+ pub fn get_row_iter(
124
+ &self,
125
+ row: usize,
126
+ ) -> std::iter::StepBy<std::iter::Skip<std::slice::Iter<T>>> {
127
+ self.data.iter().skip(row).step_by(self.rows)
128
+ }
129
+
121
130
  /// Get a slice of a column in the matrix.
122
131
  ///
123
132
  /// * `col` - The index of the column to select.
@@ -137,6 +146,30 @@ impl<'a, T> Matrix<'a, T> {
137
146
  }
138
147
  }
139
148
 
149
+ impl<'a, T> Matrix<'a, T>
150
+ where
151
+ T: Copy,
152
+ {
153
+ /// Get a row of the data as a vector.
154
+ pub fn get_row(&self, row: usize) -> Vec<T> {
155
+ self.get_row_iter(row).copied().collect()
156
+ }
157
+ }
158
+
159
+ /// A lightweight row major matrix, this is primarily
160
+ /// for returning data to the user, it is especially
161
+ /// suited for appending rows to, such as when building
162
+ /// up a matrix of contributions to return to the
163
+ /// user, the added benefit is it will be even
164
+ /// faster to return to numpy.
165
+ // pub struct RowMajorMatrix<T> {
166
+ // pub data: Vec<T>,
167
+ // pub rows: usize,
168
+ // pub cols: usize,
169
+ // stride1: usize,
170
+ // stride2: usize,
171
+ // }
172
+
140
173
  impl<'a, T> fmt::Display for Matrix<'a, T>
141
174
  where
142
175
  T: FromStr + std::fmt::Display,
@@ -212,7 +245,6 @@ impl<T> JaggedMatrix<T> {
212
245
  n_records: 0,
213
246
  }
214
247
  }
215
-
216
248
  /// Get the column of a jagged array.
217
249
  pub fn get_col(&self, col: usize) -> &[T] {
218
250
  assert!(col < self.ends.len());
@@ -270,6 +302,15 @@ mod tests {
270
302
  assert_eq!(m.get_col(1), &vec![5, 6, 7]);
271
303
  }
272
304
 
305
+ #[test]
306
+ fn test_matrix_row() {
307
+ let v = vec![1, 2, 3, 5, 6, 7];
308
+ let m = Matrix::new(&v, 3, 2);
309
+ assert_eq!(m.get_row(2), vec![3, 7]);
310
+ assert_eq!(m.get_row(0), vec![1, 5]);
311
+ assert_eq!(m.get_row(1), vec![2, 6]);
312
+ }
313
+
273
314
  #[test]
274
315
  fn test_jaggedmatrix_get_col() {
275
316
  let vecs = vec![vec![0], vec![5, 4, 3, 2], vec![4, 5]];
@@ -5,6 +5,7 @@ use crate::errors::ForustError;
5
5
  use crate::objective::{gradient_hessian_callables, ObjectiveType};
6
6
  use crate::splitter::MissingImputerSplitter;
7
7
  use crate::tree::Tree;
8
+ use rayon::prelude::*;
8
9
  use serde::{Deserialize, Serialize};
9
10
  use std::fs;
10
11
 
@@ -215,6 +216,84 @@ impl GradientBooster {
215
216
  init_preds
216
217
  }
217
218
 
219
+ // pub fn predict(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
220
+ // // After we disconvered it's faster materializing the row once, and then
221
+ // // Passing that to each tree, let's see if we can do the same with the booster
222
+ // // prediction...
223
+ // // Clean this up..
224
+ // let mut init_preds = vec![self.base_score; data.rows];
225
+ // if parallel {
226
+ // init_preds.par_iter_mut().enumerate().for_each(|(i, p)| {
227
+ // let pred_row = data.get_row(i);
228
+ // for t in &self.trees {
229
+ // *p += t.predict_row_from_row_slice(&pred_row);
230
+ // }
231
+ // });
232
+ // } else {
233
+ // init_preds.iter_mut().enumerate().for_each(|(i, p)| {
234
+ // let pred_row = data.get_row(i);
235
+ // for t in &self.trees {
236
+ // *p += t.predict_row_from_row_slice(&pred_row);
237
+ // }
238
+ // });
239
+ // }
240
+ // init_preds
241
+ // }
242
+
243
+ /// Generate predictions on data using the gradient booster.
244
+ ///
245
+ /// * `data` - Either a pandas DataFrame, or a 2 dimensional numpy array.
246
+ pub fn predict_contributions(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
247
+ let weights: Vec<Vec<f64>> = if parallel {
248
+ self.trees
249
+ .par_iter()
250
+ .map(|t| t.distribute_leaf_weights())
251
+ .collect()
252
+ } else {
253
+ self.trees
254
+ .iter()
255
+ .map(|t| t.distribute_leaf_weights())
256
+ .collect()
257
+ };
258
+ let mut contribs = vec![0.; (data.cols + 1) * data.rows];
259
+
260
+ // Add the bias term to every bias value...
261
+ let bias_idx = data.cols + 1;
262
+ contribs
263
+ .iter_mut()
264
+ .skip(bias_idx - 1)
265
+ .step_by(bias_idx)
266
+ .for_each(|v| *v += self.base_score);
267
+
268
+ // Clean this up..
269
+ // materializing a row, and then passing that to all of the
270
+ // trees seems to be the fastest approach (5X faster), we should test
271
+ // something like this for normal predictions.
272
+ if parallel {
273
+ data.index
274
+ .par_iter()
275
+ .zip(contribs.par_chunks_mut(data.cols + 1))
276
+ .for_each(|(row, c)| {
277
+ let r_ = data.get_row(*row);
278
+ self.trees.iter().zip(weights.iter()).for_each(|(t, w)| {
279
+ t.predict_contributions_row(&r_, c, w);
280
+ });
281
+ });
282
+ } else {
283
+ data.index
284
+ .iter()
285
+ .zip(contribs.chunks_mut(data.cols + 1))
286
+ .for_each(|(row, c)| {
287
+ let r_ = data.get_row(*row);
288
+ self.trees.iter().zip(weights.iter()).for_each(|(t, w)| {
289
+ t.predict_contributions_row(&r_, c, w);
290
+ });
291
+ });
292
+ }
293
+
294
+ contribs
295
+ }
296
+
218
297
  /// Given a value, return the partial dependence value of that value for that
219
298
  /// feature in the model.
220
299
  ///
@@ -390,6 +469,8 @@ mod tests {
390
469
  let sample_weight = vec![1.; y.len()];
391
470
  booster.fit(&data, &y, &sample_weight).unwrap();
392
471
  let preds = booster.predict(&data, false);
472
+ let contribs = booster.predict_contributions(&data, false);
473
+ assert_eq!(contribs.len(), (data.cols + 1) * data.rows);
393
474
  println!("{}", booster.trees[0]);
394
475
  println!("{}", booster.trees[0].nodes.len());
395
476
  println!("{}", booster.trees.last().unwrap().nodes.len());
@@ -23,10 +23,11 @@ pub struct SplittableNode {
23
23
  pub stop_idx: usize,
24
24
  pub lower_bound: f32,
25
25
  pub upper_bound: f32,
26
+ pub is_leaf: bool,
26
27
  }
27
28
 
28
29
  #[derive(Deserialize, Serialize)]
29
- pub struct ParentNode {
30
+ pub struct Node {
30
31
  pub num: usize,
31
32
  pub weight_value: f32,
32
33
  pub hessian_sum: f32,
@@ -37,6 +38,32 @@ pub struct ParentNode {
37
38
  pub missing_node: usize,
38
39
  pub left_child: usize,
39
40
  pub right_child: usize,
41
+ pub is_leaf: bool,
42
+ }
43
+
44
+ impl Node {
45
+ /// Update all the info that is needed if this node is a
46
+ /// parent node, this consumes the SplitableNode.
47
+ pub fn make_parent_node(&mut self, split_node: SplittableNode) {
48
+ self.is_leaf = false;
49
+ self.missing_node = split_node.missing_node;
50
+ self.split_value = split_node.split_value;
51
+ self.split_feature = split_node.split_feature;
52
+ self.split_gain = split_node.split_gain;
53
+ self.left_child = split_node.left_child;
54
+ self.right_child = split_node.right_child;
55
+ }
56
+ /// Get the path that should be traveled down, given a value.
57
+ pub fn get_child_idx(&self, v: &f64) -> usize {
58
+ if v.is_nan() {
59
+ self.missing_node
60
+ } else if v < &self.split_value {
61
+ self.left_child
62
+ } else {
63
+ // if v >= &self.split_value
64
+ self.right_child
65
+ }
66
+ }
40
67
  }
41
68
 
42
69
  #[derive(Deserialize, Serialize)]
@@ -47,13 +74,6 @@ pub struct LeafNode {
47
74
  pub depth: usize,
48
75
  }
49
76
 
50
- #[derive(Deserialize, Serialize)]
51
- pub enum TreeNode {
52
- Parent(ParentNode),
53
- Leaf(LeafNode),
54
- Splittable(SplittableNode),
55
- }
56
-
57
77
  impl SplittableNode {
58
78
  pub fn from_node_info(
59
79
  num: usize,
@@ -81,9 +101,12 @@ impl SplittableNode {
81
101
  stop_idx,
82
102
  lower_bound: node_info.bounds.0,
83
103
  upper_bound: node_info.bounds.1,
104
+ is_leaf: true,
84
105
  }
85
106
  }
86
107
 
108
+ /// Create a default splitable node,
109
+ /// we default to the node being a leaf.
87
110
  #[allow(clippy::too_many_arguments)]
88
111
  pub fn new(
89
112
  num: usize,
@@ -116,6 +139,7 @@ impl SplittableNode {
116
139
  stop_idx,
117
140
  lower_bound,
118
141
  upper_bound,
142
+ is_leaf: true,
119
143
  }
120
144
  }
121
145
 
@@ -136,17 +160,10 @@ impl SplittableNode {
136
160
  MissingInfo::Branch(_) => todo!(),
137
161
  MissingInfo::EmptyBranch => todo!(),
138
162
  };
163
+ self.is_leaf = false;
139
164
  }
140
- pub fn as_leaf_node(&self) -> TreeNode {
141
- TreeNode::Leaf(LeafNode {
142
- num: self.num,
143
- weight_value: self.weight_value,
144
- hessian_sum: self.hessian_sum,
145
- depth: self.depth,
146
- })
147
- }
148
- pub fn as_parent_node(&self) -> TreeNode {
149
- TreeNode::Parent(ParentNode {
165
+ pub fn as_node(&self) -> Node {
166
+ Node {
150
167
  num: self.num,
151
168
  weight_value: self.weight_value,
152
169
  hessian_sum: self.hessian_sum,
@@ -157,47 +174,33 @@ impl SplittableNode {
157
174
  split_gain: self.split_gain,
158
175
  left_child: self.left_child,
159
176
  right_child: self.right_child,
160
- })
177
+ is_leaf: self.is_leaf,
178
+ }
161
179
  }
162
180
  }
163
181
 
164
- impl fmt::Display for TreeNode {
182
+ impl fmt::Display for Node {
165
183
  // This trait requires `fmt` with this exact signature.
166
184
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
167
- match self {
168
- TreeNode::Leaf(leaf) => write!(
185
+ if self.is_leaf {
186
+ write!(
169
187
  f,
170
188
  "{}:leaf={},cover={}",
171
- leaf.num, leaf.weight_value, leaf.hessian_sum
172
- ),
173
- TreeNode::Parent(parent) => {
174
- write!(
175
- f,
176
- "{}:[{} < {}] yes={},no={},missing={},gain={},cover={}",
177
- parent.num,
178
- parent.split_feature,
179
- parent.split_value,
180
- parent.left_child,
181
- parent.right_child,
182
- parent.missing_node,
183
- parent.split_gain,
184
- parent.hessian_sum
185
- )
186
- }
187
- TreeNode::Splittable(node) => {
188
- write!(
189
- f,
190
- "SPLITTABLE - {}:[{} < {}] yes={},no={},missing={},gain={},cover={}",
191
- node.num,
192
- node.split_feature,
193
- node.split_value,
194
- node.missing_node,
195
- node.left_child,
196
- node.right_child,
197
- node.split_gain,
198
- node.hessian_sum
199
- )
200
- }
189
+ self.num, self.weight_value, self.hessian_sum
190
+ )
191
+ } else {
192
+ write!(
193
+ f,
194
+ "{}:[{} < {}] yes={},no={},missing={},gain={},cover={}",
195
+ self.num,
196
+ self.split_feature,
197
+ self.split_value,
198
+ self.left_child,
199
+ self.right_child,
200
+ self.missing_node,
201
+ self.split_gain,
202
+ self.hessian_sum
203
+ )
201
204
  }
202
205
  }
203
206
  }