1use ndarray::{Array1, ArrayView1, ArrayView2};
18use rand::seq::index::sample;
19use rand::{Rng, SeedableRng};
20use rand_chacha::ChaCha8Rng;
21
22use super::{Classifier, Regressor};
23use crate::error::{DmlError, Result};
24
25#[derive(Debug, Clone, Copy, PartialEq)]
27pub enum MaxFeatures {
28 All,
30 Sqrt,
32 Log2,
34 Fraction(f64),
36 Count(usize),
38}
39
40impl MaxFeatures {
41 pub(super) fn resolve(self, p: usize) -> usize {
42 let k = match self {
43 MaxFeatures::All => p,
44 MaxFeatures::Sqrt => (p as f64).sqrt() as usize,
45 MaxFeatures::Log2 => (p as f64).log2() as usize,
46 MaxFeatures::Fraction(f) => (f * p as f64) as usize,
47 MaxFeatures::Count(k) => k,
48 };
49 k.clamp(1, p.max(1))
50 }
51}
52
53#[derive(Debug, Clone, Copy)]
55pub(crate) struct TreeParams {
56 pub max_depth: Option<usize>,
57 pub min_samples_split: usize,
58 pub min_samples_leaf: usize,
59 pub max_features: MaxFeatures,
60}
61
62#[derive(Debug, Clone, Copy)]
63pub(crate) enum Node {
64 Leaf {
65 value: f64,
66 },
67 Split {
68 feature: usize,
69 threshold: f64,
70 left: u32,
71 right: u32,
72 },
73}
74
75#[derive(Debug, Clone)]
76pub(crate) struct Tree {
77 pub nodes: Vec<Node>,
78}
79
80impl Tree {
81 pub(super) fn predict_row(&self, row: ArrayView1<'_, f64>) -> f64 {
82 match self.nodes[self.leaf_index(row)] {
83 Node::Leaf { value } => value,
84 Node::Split { .. } => unreachable!("leaf_index returns a leaf"),
85 }
86 }
87
88 pub(crate) fn leaf_index(&self, row: ArrayView1<'_, f64>) -> usize {
90 let mut i = 0;
91 loop {
92 match self.nodes[i] {
93 Node::Leaf { .. } => return i,
94 Node::Split {
95 feature,
96 threshold,
97 left,
98 right,
99 } => {
100 i = if row[feature] <= threshold {
101 left
102 } else {
103 right
104 } as usize
105 }
106 }
107 }
108 }
109}
110
111pub(crate) struct TreeBuilder<'x, 'b> {
117 pub x: ArrayView2<'x, f64>,
118 pub y: &'b [f64],
119 pub w: &'b [f64],
120 pub params: TreeParams,
121 pub n_try: usize,
122 pub rng: ChaCha8Rng,
123 pub importance: Vec<f64>,
125}
126
127struct Best {
128 feature: usize,
129 threshold: f64,
130 gain: f64,
131}
132
133impl TreeBuilder<'_, '_> {
134 pub(crate) fn grow(mut self, rows: Vec<usize>) -> (Tree, Vec<f64>) {
135 let mut nodes = vec![Node::Leaf { value: 0.0 }];
136 let mut stack = vec![(0usize, rows, 0usize)];
138 while let Some((id, rows, depth)) = stack.pop() {
139 let (sw, swy, swy2) = rows.iter().fold((0.0, 0.0, 0.0), |(a, b, c), &i| {
140 let (wi, yi) = (self.w[i], self.y[i]);
141 (a + wi, b + wi * yi, c + wi * yi * yi)
142 });
143 let value = swy / sw;
144 let sse = swy2 - swy * swy / sw;
145 let splittable = rows.len() >= self.params.min_samples_split
146 && rows.len() >= 2 * self.params.min_samples_leaf
147 && self.params.max_depth.map_or(true, |d| depth < d)
148 && sse > 1e-12 * swy2.max(f64::MIN_POSITIVE);
149 let best = if splittable {
150 self.best_split(&rows, sw, swy)
151 } else {
152 None
153 };
154 let Some(best) = best else {
155 nodes[id] = Node::Leaf { value };
156 continue;
157 };
158 self.importance[best.feature] += best.gain;
159 let (left_rows, right_rows): (Vec<usize>, Vec<usize>) = rows
160 .into_iter()
161 .partition(|&i| self.x[[i, best.feature]] <= best.threshold);
162 let left = nodes.len();
163 nodes.push(Node::Leaf { value: 0.0 });
164 nodes.push(Node::Leaf { value: 0.0 });
165 nodes[id] = Node::Split {
166 feature: best.feature,
167 threshold: best.threshold,
168 left: left as u32,
169 right: (left + 1) as u32,
170 };
171 stack.push((left + 1, right_rows, depth + 1));
172 stack.push((left, left_rows, depth + 1));
173 }
174 (Tree { nodes }, self.importance)
175 }
176
177 fn best_split(&mut self, rows: &[usize], sw: f64, swy: f64) -> Option<Best> {
180 let p = self.x.ncols();
181 let parent = swy * swy / sw;
182 let min_leaf = self.params.min_samples_leaf;
183 let mut best: Option<Best> = None;
184 let mut sorted: Vec<(f64, usize)> = Vec::with_capacity(rows.len());
185 for feature in sample(&mut self.rng, p, self.n_try).into_iter() {
186 sorted.clear();
187 sorted.extend(rows.iter().map(|&i| (self.x[[i, feature]], i)));
188 sorted.sort_unstable_by(|a, b| a.0.total_cmp(&b.0));
189 if sorted[0].0 == sorted[sorted.len() - 1].0 {
190 continue; }
192 let (mut wl, mut sl) = (0.0, 0.0);
193 for k in 0..sorted.len() - 1 {
194 let i = sorted[k].1;
195 wl += self.w[i];
196 sl += self.w[i] * self.y[i];
197 let (v, v_next) = (sorted[k].0, sorted[k + 1].0);
198 let n_left = k + 1;
199 if v == v_next || n_left < min_leaf || sorted.len() - n_left < min_leaf {
200 continue;
201 }
202 let (wr, sr) = (sw - wl, swy - sl);
203 if wl <= 0.0 || wr <= 0.0 {
204 continue;
205 }
206 let gain = sl * sl / wl + sr * sr / wr - parent;
207 if best.as_ref().map_or(true, |b| gain > b.gain) {
208 let mid = v + (v_next - v) / 2.0;
210 let threshold = if mid < v_next { mid } else { v };
211 best = Some(Best {
212 feature,
213 threshold,
214 gain,
215 });
216 }
217 }
218 }
219 best.filter(|b| b.gain > 1e-12 * parent.abs().max(f64::MIN_POSITIVE))
220 }
221}
222
223#[derive(Debug, Clone)]
225struct Forest {
226 n_estimators: usize,
227 params: TreeParams,
228 bootstrap: bool,
229 max_samples: Option<f64>,
230 seed: u64,
231 trees: Vec<Tree>,
232 y_offset: f64,
233 n_features: usize,
234 importances: Option<Array1<f64>>,
235}
236
237impl Forest {
238 fn new(max_features: MaxFeatures) -> Self {
239 Self {
240 n_estimators: 100,
241 params: TreeParams {
242 max_depth: None,
243 min_samples_split: 2,
244 min_samples_leaf: 1,
245 max_features,
246 },
247 bootstrap: true,
248 max_samples: None,
249 seed: 0,
250 trees: Vec::new(),
251 y_offset: 0.0,
252 n_features: 0,
253 importances: None,
254 }
255 }
256
257 fn unfitted(&self) -> Self {
258 Self {
259 trees: Vec::new(),
260 importances: None,
261 ..self.clone()
262 }
263 }
264
265 fn validate(&self) -> Result<()> {
266 if self.n_estimators == 0 {
267 return Err(DmlError::InvalidArgument(
268 "n_estimators must be at least 1".into(),
269 ));
270 }
271 if self.params.min_samples_leaf == 0 || self.params.min_samples_split < 2 {
272 return Err(DmlError::InvalidArgument(
273 "min_samples_leaf must be ≥ 1 and min_samples_split ≥ 2".into(),
274 ));
275 }
276 if let MaxFeatures::Fraction(f) = self.params.max_features {
277 if !(f > 0.0 && f <= 1.0) {
278 return Err(DmlError::InvalidArgument(
279 "max_features fraction must be in (0, 1]".into(),
280 ));
281 }
282 }
283 if let Some(m) = self.max_samples {
284 if !(m > 0.0 && m <= 1.0) {
285 return Err(DmlError::InvalidArgument(
286 "max_samples must be in (0, 1]".into(),
287 ));
288 }
289 }
290 Ok(())
291 }
292
293 fn fit(
294 &mut self,
295 x: ArrayView2<'_, f64>,
296 y: ArrayView1<'_, f64>,
297 w: Option<ArrayView1<'_, f64>>,
298 ) -> Result<()> {
299 self.validate()?;
300 let (n, p) = x.dim();
301 if y.len() != n {
302 return Err(DmlError::DimensionMismatch("forest fit x/y rows".into()));
303 }
304 if n == 0 {
305 return Err(DmlError::InvalidArgument("empty training set".into()));
306 }
307 let weights: Vec<f64> = match w {
308 Some(w) => {
309 if w.len() != n {
310 return Err(DmlError::DimensionMismatch("sample weight length".into()));
311 }
312 if !w.iter().all(|&v| v >= 0.0 && v.is_finite()) || w.sum() <= 0.0 {
313 return Err(DmlError::InvalidArgument(
314 "sample weights must be non-negative, finite and not all zero".into(),
315 ));
316 }
317 w.to_vec()
318 }
319 None => vec![1.0; n],
320 };
321 let total_w: f64 = weights.iter().sum();
323 let offset = y.iter().zip(&weights).map(|(&a, &b)| a * b).sum::<f64>() / total_w;
324 let yc: Vec<f64> = y.iter().map(|&v| v - offset).collect();
325 let n_draw = self
326 .max_samples
327 .map_or(n, |m| ((m * n as f64).round() as usize).max(1));
328 let n_try = self.params.max_features.resolve(p);
329
330 let grow_tree = |t: usize| -> (Tree, Vec<f64>) {
331 let tree_seed = self
332 .seed
333 .wrapping_add((t as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
334 let mut rng = ChaCha8Rng::seed_from_u64(tree_seed);
335 let mut eff = vec![0.0; n];
336 if self.bootstrap {
337 for _ in 0..n_draw {
338 eff[rng.gen_range(0..n)] += 1.0;
339 }
340 for (e, &wi) in eff.iter_mut().zip(&weights) {
341 *e *= wi;
342 }
343 } else {
344 eff.clone_from(&weights);
345 }
346 let rows: Vec<usize> = (0..n).filter(|&i| eff[i] > 0.0).collect();
347 let rows = if rows.is_empty() {
348 (0..n).collect()
349 } else {
350 rows
351 };
352 TreeBuilder {
353 x,
354 y: &yc,
355 w: &eff,
356 params: self.params,
357 n_try,
358 rng,
359 importance: vec![0.0; p],
360 }
361 .grow(rows)
362 };
363
364 #[cfg(feature = "parallel")]
365 let grown: Vec<(Tree, Vec<f64>)> = {
366 use rayon::prelude::*;
367 (0..self.n_estimators)
368 .into_par_iter()
369 .map(grow_tree)
370 .collect()
371 };
372 #[cfg(not(feature = "parallel"))]
373 let grown: Vec<(Tree, Vec<f64>)> = (0..self.n_estimators).map(grow_tree).collect();
374
375 let mut importances = Array1::<f64>::zeros(p);
377 let mut counted = 0.0;
378 for (_, imp) in &grown {
379 let total: f64 = imp.iter().sum();
380 if total > 0.0 {
381 importances += &(Array1::from(imp.clone()) / total);
382 counted += 1.0;
383 }
384 }
385 if counted > 0.0 {
386 importances /= counted;
387 }
388 self.trees = grown.into_iter().map(|(t, _)| t).collect();
389 self.y_offset = offset;
390 self.n_features = p;
391 self.importances = Some(importances);
392 Ok(())
393 }
394
395 fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
396 if self.trees.is_empty() {
397 return Err(DmlError::NotFitted);
398 }
399 if x.ncols() != self.n_features {
400 return Err(DmlError::DimensionMismatch(format!(
401 "x has {} columns, the forest was fitted with {}",
402 x.ncols(),
403 self.n_features
404 )));
405 }
406 let per_row = |row: ArrayView1<'_, f64>| -> f64 {
407 self.trees.iter().map(|t| t.predict_row(row)).sum::<f64>() / self.trees.len() as f64
408 + self.y_offset
409 };
410 Ok(x.rows().into_iter().map(per_row).collect())
411 }
412}
413
414macro_rules! forest_builders {
416 () => {
417 pub fn n_estimators(mut self, n: usize) -> Self {
419 self.forest.n_estimators = n;
420 self
421 }
422
423 pub fn max_depth(mut self, depth: Option<usize>) -> Self {
425 self.forest.params.max_depth = depth;
426 self
427 }
428
429 pub fn min_samples_split(mut self, n: usize) -> Self {
431 self.forest.params.min_samples_split = n;
432 self
433 }
434
435 pub fn min_samples_leaf(mut self, n: usize) -> Self {
438 self.forest.params.min_samples_leaf = n;
439 self
440 }
441
442 pub fn max_features(mut self, m: MaxFeatures) -> Self {
444 self.forest.params.max_features = m;
445 self
446 }
447
448 pub fn bootstrap(mut self, bootstrap: bool) -> Self {
450 self.forest.bootstrap = bootstrap;
451 self
452 }
453
454 pub fn max_samples(mut self, fraction: Option<f64>) -> Self {
456 self.forest.max_samples = fraction;
457 self
458 }
459
460 pub fn seed(mut self, seed: u64) -> Self {
462 self.forest.seed = seed;
463 self
464 }
465
466 pub fn feature_importances(&self) -> Option<&Array1<f64>> {
468 self.forest.importances.as_ref()
469 }
470 };
471}
472
473#[derive(Debug, Clone)]
488pub struct RandomForestRegressor {
489 forest: Forest,
490}
491
492impl Default for RandomForestRegressor {
493 fn default() -> Self {
494 Self::new()
495 }
496}
497
498impl RandomForestRegressor {
499 pub fn new() -> Self {
500 Self {
501 forest: Forest::new(MaxFeatures::All),
502 }
503 }
504
505 forest_builders!();
506}
507
508impl Regressor for RandomForestRegressor {
509 fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
510 self.forest.fit(x, y, None)
511 }
512
513 fn fit_weighted(
514 &mut self,
515 x: ArrayView2<'_, f64>,
516 y: ArrayView1<'_, f64>,
517 w: ArrayView1<'_, f64>,
518 ) -> Result<()> {
519 self.forest.fit(x, y, Some(w))
520 }
521
522 fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
523 self.forest.predict(x)
524 }
525
526 fn fresh(&self) -> Box<dyn Regressor> {
527 Box::new(Self {
528 forest: self.forest.unfitted(),
529 })
530 }
531}
532
533#[derive(Debug, Clone)]
539pub struct RandomForestClassifier {
540 forest: Forest,
541}
542
543impl Default for RandomForestClassifier {
544 fn default() -> Self {
545 Self::new()
546 }
547}
548
549impl RandomForestClassifier {
550 pub fn new() -> Self {
551 Self {
552 forest: Forest::new(MaxFeatures::Sqrt),
553 }
554 }
555
556 forest_builders!();
557}
558
559impl Classifier for RandomForestClassifier {
560 fn fit(&mut self, x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
561 if !y.iter().all(|&v| v == 0.0 || v == 1.0) {
562 return Err(DmlError::InvalidArgument(
563 "RandomForestClassifier labels must be 0 or 1".into(),
564 ));
565 }
566 self.forest.fit(x, y, None)
567 }
568
569 fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
570 Ok(self.forest.predict(x)?.mapv(|p| p.clamp(0.0, 1.0)))
572 }
573
574 fn fresh(&self) -> Box<dyn Classifier> {
575 Box::new(Self {
576 forest: self.forest.unfitted(),
577 })
578 }
579}
580
581#[cfg(test)]
582mod tests {
583 use super::*;
584 use ndarray::{array, Array2};
585
586 fn nonlinear(n: usize, seed: u64) -> (Array2<f64>, Array1<f64>) {
587 let mut rng = ChaCha8Rng::seed_from_u64(seed);
588 let x = Array2::from_shape_simple_fn((n, 3), || rng.gen::<f64>() * 4.0 - 2.0);
589 let y = x
590 .rows()
591 .into_iter()
592 .map(|r| r[0].sin() * 2.0 + (r[1] > 0.0) as u8 as f64)
593 .collect();
594 (x, y)
595 }
596
597 #[test]
598 fn single_tree_without_bootstrap_fits_a_step_exactly() {
599 let x = array![[0.0], [1.0], [2.0], [3.0], [4.0], [5.0]];
600 let y = array![1.0, 1.0, 1.0, 7.0, 7.0, 7.0];
601 let mut rf = RandomForestRegressor::new()
602 .n_estimators(1)
603 .bootstrap(false);
604 rf.fit(x.view(), y.view()).unwrap();
605 let pred = rf
606 .predict(array![[2.4], [2.6], [-10.0], [10.0]].view())
607 .unwrap();
608 assert_eq!(pred, array![1.0, 7.0, 1.0, 7.0]);
609 assert!(
611 matches!(rf.forest.trees[0].nodes[0], Node::Split { threshold, .. } if threshold == 2.5)
612 );
613 }
614
615 #[test]
616 fn forest_learns_a_nonlinear_function() {
617 let (x, y) = nonlinear(2000, 1);
618 let (x_test, y_test) = nonlinear(500, 2);
619 let mut rf = RandomForestRegressor::new()
620 .n_estimators(100)
621 .min_samples_leaf(3)
622 .seed(7);
623 rf.fit(x.view(), y.view()).unwrap();
624 let resid = &y_test - &rf.predict(x_test.view()).unwrap();
625 let r2 = 1.0
626 - resid.mapv(|e| e * e).sum()
627 / y_test.mapv(|v| (v - y_test.mean().unwrap()).powi(2)).sum();
628 assert!(r2 > 0.9, "out-of-sample R² = {r2}");
629 let imp = rf.feature_importances().unwrap();
631 assert!((imp.sum() - 1.0).abs() < 1e-9);
632 assert!(imp[2] < imp[0] && imp[2] < imp[1], "{imp}");
633 }
634
635 #[test]
636 fn same_seed_same_forest_and_seeds_matter() {
637 let (x, y) = nonlinear(300, 3);
638 let fit = |seed| {
639 let mut rf = RandomForestRegressor::new()
640 .n_estimators(20)
641 .max_features(MaxFeatures::Count(2))
642 .seed(seed);
643 rf.fit(x.view(), y.view()).unwrap();
644 rf.predict(x.view()).unwrap()
645 };
646 assert_eq!(fit(1), fit(1));
647 assert_ne!(fit(1), fit(2));
648 }
649
650 #[test]
651 fn weights_equal_row_duplication_without_bootstrap() {
652 let x = array![[0.0], [1.0], [2.0], [3.0]];
653 let y = array![0.0, 1.0, 5.0, 6.0];
654 let mut weighted = RandomForestRegressor::new()
655 .n_estimators(1)
656 .bootstrap(false)
657 .max_depth(Some(1));
658 weighted
659 .fit_weighted(x.view(), y.view(), array![3.0, 1.0, 1.0, 1.0].view())
660 .unwrap();
661 let mut dup = RandomForestRegressor::new()
662 .n_estimators(1)
663 .bootstrap(false)
664 .max_depth(Some(1));
665 let xd = array![[0.0], [0.0], [0.0], [1.0], [2.0], [3.0]];
666 let yd = array![0.0, 0.0, 0.0, 1.0, 5.0, 6.0];
667 dup.fit(xd.view(), yd.view()).unwrap();
668 let grid = array![[0.0], [1.5], [3.0]];
669 let diff = weighted.predict(grid.view()).unwrap() - dup.predict(grid.view()).unwrap();
670 assert!(diff.iter().all(|d| d.abs() < 1e-12));
671 }
672
673 #[test]
674 fn min_samples_leaf_and_depth_are_respected() {
675 let (x, y) = nonlinear(200, 4);
676 let mut stump = RandomForestRegressor::new()
677 .n_estimators(3)
678 .max_depth(Some(1));
679 stump.fit(x.view(), y.view()).unwrap();
680 assert!(stump.forest.trees.iter().all(|t| t.nodes.len() <= 3));
681 let mut big_leaves = RandomForestRegressor::new()
682 .n_estimators(1)
683 .bootstrap(false)
684 .min_samples_leaf(50);
685 big_leaves.fit(x.view(), y.view()).unwrap();
686 assert!(big_leaves.forest.trees[0].nodes.len() <= 7); }
688
689 #[test]
690 fn classifier_probabilities_track_the_truth() {
691 let mut rng = ChaCha8Rng::seed_from_u64(5);
692 let n = 3000;
693 let x = Array2::from_shape_simple_fn((n, 2), || rng.gen::<f64>() * 4.0 - 2.0);
694 let truth: Array1<f64> = x
695 .rows()
696 .into_iter()
697 .map(|r| 1.0 / (1.0 + (-2.0 * r[0]).exp()))
698 .collect();
699 let labels = truth.mapv(|p| f64::from(rng.gen::<f64>() < p));
700 let mut rf = RandomForestClassifier::new().min_samples_leaf(20).seed(1);
701 rf.fit(x.view(), labels.view()).unwrap();
702 let p = rf.predict_proba(x.view()).unwrap();
703 assert!(p.iter().all(|&v| (0.0..=1.0).contains(&v)));
704 let mae = (&p - &truth).mapv(f64::abs).mean().unwrap();
705 assert!(mae < 0.08, "mean |p̂ − p| = {mae}");
706 assert!(rf.fit(x.view(), (&labels * 2.0).view()).is_err());
707 }
708
709 #[test]
710 fn invalid_settings_and_unfitted_use_error() {
711 let x = array![[0.0], [1.0]];
712 let y = array![0.0, 1.0];
713 assert!(RandomForestRegressor::new()
714 .n_estimators(0)
715 .fit(x.view(), y.view())
716 .is_err());
717 assert!(RandomForestRegressor::new()
718 .max_features(MaxFeatures::Fraction(1.5))
719 .fit(x.view(), y.view())
720 .is_err());
721 assert_eq!(
722 RandomForestRegressor::new().predict(x.view()).unwrap_err(),
723 DmlError::NotFitted
724 );
725 assert_eq!(MaxFeatures::Sqrt.resolve(10), 3);
726 assert_eq!(MaxFeatures::Log2.resolve(1), 1);
727 assert_eq!(MaxFeatures::Count(50).resolve(4), 4);
728 }
729}