dtrees_rs/tree/
mod.rs

1//! Binary decision trees stored as an arena of nodes.
2
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6/// The content of a tree node.
7#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
8pub struct NodeInfos {
9    /// The feature tested by an internal node; `None` for a leaf.
10    pub test: Option<usize>,
11    /// Error of the subtree rooted here.
12    pub error: f64,
13    /// Score of the node for searches that optimise another metric.
14    pub metric: Option<f64>,
15    /// The prediction of a leaf.
16    pub out: Option<f64>,
17}
18
19impl Default for NodeInfos {
20    fn default() -> Self {
21        NodeInfos::new()
22    }
23}
24
25impl NodeInfos {
26    pub fn new() -> NodeInfos {
27        NodeInfos {
28            test: None,
29            error: <f64>::INFINITY,
30            metric: None,
31            out: None,
32        }
33    }
34}
35
36/// A node of a [`Tree`]. Child indices of `0` mean "no child".
37#[derive(Copy, Clone, Serialize, Deserialize, Debug, Default)]
38pub struct TreeNode {
39    /// What the node tests or predicts.
40    pub value: NodeInfos,
41    /// The node's own index in the arena.
42    pub index: usize,
43    /// Index of the child where the tested feature is 0.
44    pub left: usize,
45    /// Index of the child where the tested feature is 1.
46    pub right: usize,
47}
48
49impl TreeNode {
50    pub fn new(value: NodeInfos) -> TreeNode {
51        TreeNode {
52            value,
53            index: 0,
54            left: 0,
55            right: 0,
56        }
57    }
58}
59
60/// A binary decision tree stored as an arena of nodes, the root at index 0.
61#[derive(Clone, Serialize, Deserialize, Debug)]
62pub struct Tree {
63    tree: Vec<TreeNode>,
64}
65
66impl Default for Tree {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl Tree {
73    pub fn new() -> Self {
74        Tree { tree: Vec::new() }
75    }
76
77    pub fn with_capacity(capacity: usize) -> Self {
78        Tree {
79            tree: Vec::with_capacity(capacity),
80        }
81    }
82
83    pub fn is_empty(&self) -> bool {
84        self.tree.is_empty()
85    }
86
87    pub fn len(&self) -> usize {
88        self.tree.len()
89    }
90
91    /// Number of nodes reachable from the root.
92    pub fn actual_len(&self) -> usize {
93        self.count_node_recursion(self.get_root_index())
94    }
95
96    fn count_node_recursion(&self, node_index: usize) -> usize {
97        let mut left_index = 0;
98        let mut right_index = 0;
99        if let Some(node) = self.get_node(node_index) {
100            if node.left == node.right {
101                return 1;
102            } else {
103                left_index = node.left;
104                right_index = node.right;
105            }
106        }
107
108        let mut count = 0;
109
110        if left_index != 0 {
111            count += self.count_node_recursion(left_index);
112        }
113        if right_index != 0 {
114            count += self.count_node_recursion(right_index);
115        }
116
117        count + 1
118    }
119
120    /// Appends `node` as the left or right child of `parent` and returns its
121    /// index. The first node added becomes the root and `parent` is ignored.
122    pub fn add_node(&mut self, parent: usize, is_left: bool, mut node: TreeNode) -> usize {
123        node.index = self.tree.len();
124        self.tree.push(node);
125        let position = self.tree.len() - 1;
126        if position == 0 {
127            return position;
128        }
129        if let Some(parent_node) = self.tree.get_mut(parent) {
130            if is_left {
131                parent_node.left = position
132            } else {
133                parent_node.right = position
134            }
135        };
136        position
137    }
138
139    pub fn add_root(&mut self, root: TreeNode) -> usize {
140        self.add_node(0, false, root)
141    }
142
143    pub fn add_default_root(&mut self) -> usize {
144        self.add_root(TreeNode::default())
145    }
146
147    pub fn add_left_node(&mut self, parent: usize, node: TreeNode) -> usize {
148        self.add_node(parent, true, node)
149    }
150    pub fn add_right_node(&mut self, parent: usize, node: TreeNode) -> usize {
151        self.add_node(parent, false, node)
152    }
153
154    pub fn create_child(&mut self, parent: usize, left: bool) -> usize {
155        self.add_node(parent, left, TreeNode::default())
156    }
157
158    pub fn get_root_index(&self) -> usize {
159        0
160    }
161
162    pub fn get_node(&self, index: usize) -> Option<&TreeNode> {
163        self.tree.get(index)
164    }
165
166    pub fn get_node_mut(&mut self, index: usize) -> Option<&mut TreeNode> {
167        self.tree.get_mut(index)
168    }
169
170    pub fn get_left_child(&self, node: &TreeNode) -> Option<&TreeNode> {
171        if node.left == 0 {
172            None
173        } else {
174            self.tree.get(node.left)
175        }
176    }
177    pub fn get_left_child_mut(&mut self, node: &TreeNode) -> Option<&mut TreeNode> {
178        if node.left == 0 {
179            None
180        } else {
181            self.tree.get_mut(node.left)
182        }
183    }
184
185    pub fn get_right_child(&self, node: &TreeNode) -> Option<&TreeNode> {
186        if node.right == 0 {
187            None
188        } else {
189            self.tree.get(node.right)
190        }
191    }
192    pub fn get_right_child_mut(&mut self, node: &TreeNode) -> Option<&mut TreeNode> {
193        if node.right == 0 {
194            None
195        } else {
196            self.tree.get_mut(node.right)
197        }
198    }
199
200    /// A complete tree skeleton of the given depth, with empty nodes.
201    pub fn empty_tree(depth: usize) -> Tree {
202        let mut tree = Tree::new();
203        let value = NodeInfos::new();
204        let node = TreeNode::new(value);
205        let root = tree.add_root(node);
206        Self::build_tree_recurse(&mut tree, root, depth);
207        tree
208    }
209
210    fn build_tree_recurse(tree: &mut Tree, parent: usize, depth: usize) {
211        if depth == 0 {
212            if let Some(parent_node) = tree.get_node_mut(parent) {
213                parent_node.left = 0;
214                parent_node.right = 0;
215            }
216        } else {
217            let value = NodeInfos::new();
218            let node = TreeNode::new(value);
219            let left = tree.add_node(parent, true, node);
220            Self::build_tree_recurse(tree, left, depth - 1);
221            let node = TreeNode::new(value);
222            let right = tree.add_node(parent, false, node);
223            Self::build_tree_recurse(tree, right, depth - 1);
224        }
225    }
226
227    pub fn root_details(&self) -> NodeInfos {
228        self.get_node(self.get_root_index())
229            .map(|node| node.value)
230            .unwrap_or_default()
231    }
232
233    pub fn node_details(&self, index: usize) -> NodeInfos {
234        self.get_node(index)
235            .map(|node| node.value)
236            .unwrap_or_default()
237    }
238
239    pub fn root_error(&self) -> f64 {
240        self.get_node(self.get_root_index())
241            .map(|node| node.value.error)
242            .unwrap_or(f64::MAX)
243    }
244
245    pub fn node_error(&self, index: usize) -> f64 {
246        self.get_node(index)
247            .map(|node| node.value.error)
248            .unwrap_or(f64::MAX)
249    }
250
251    pub fn node_metric(&self, index: usize) -> Option<f64> {
252        self.get_node(index)
253            .map(|node| node.value.metric)
254            .unwrap_or(Some(0.0))
255    }
256
257    pub fn root_output(&self) -> Option<f64> {
258        self.get_node(self.get_root_index())
259            .and_then(|node| node.value.out)
260    }
261
262    pub fn node_output(&self, index: usize) -> Option<f64> {
263        self.get_node(index).and_then(|node| node.value.out)
264    }
265
266    pub fn root_test(&self) -> Option<usize> {
267        self.get_node(self.get_root_index())
268            .and_then(|node| node.value.test)
269    }
270
271    pub fn node_test(&self, index: usize) -> Option<usize> {
272        self.get_node(index).and_then(|node| node.value.test)
273    }
274
275    /// A builder that edits the node at `index`.
276    pub fn update_node(&mut self, index: usize) -> Option<NodeUpdater<'_>> {
277        self.get_node_mut(index).map(NodeUpdater::new)
278    }
279
280    pub fn update_root(&mut self) -> Option<NodeUpdater<'_>> {
281        self.get_node_mut(0).map(NodeUpdater::new)
282    }
283
284    /// `(left, right)` child indices of the node at `index`; `0` means none.
285    pub fn node_children(&self, index: usize) -> (usize, usize) {
286        self.get_node(index)
287            .map_or((0, 0), |node| (node.left, node.right))
288    }
289
290    /// Sets the `(error, prediction)` of the node at `index`.
291    pub fn update_leaf_node(&mut self, index: usize, error: (f64, f64)) -> &mut Self {
292        if let Some(updater) = self.update_node(index) {
293            updater.error(error.0).output(error.1);
294        }
295        self
296    }
297
298    /// Copies the subtree of `origin` rooted at `origin_index` onto the node at
299    /// `index`, creating children as needed.
300    pub fn update_subtree(&mut self, index: usize, origin: &Tree, origin_index: usize) {
301        let (left_index, right_index) = self.update_node(index).map_or((0, 0), |updater| {
302            updater
303                .value(origin.node_details(origin_index))
304                .get_children()
305        });
306
307        let (origin_left_index, origin_right_index) = origin.node_children(origin_index);
308
309        for (branch_value, (&source, mut dest)) in [origin_left_index, origin_right_index]
310            .iter()
311            .zip([left_index, right_index])
312            .enumerate()
313        {
314            if source > 0 {
315                if dest == 0 {
316                    dest = self.create_child(index, branch_value == 0);
317                }
318                self.update_subtree(dest, origin, source);
319            }
320        }
321    }
322
323    /// Merges pairs of sibling leaves that predict the same class into their
324    /// parent.
325    pub fn clean_orphaned_nodes(&mut self) {
326        if self.is_empty() {
327            return;
328        }
329        self.cleanup_leaves(self.get_root_index())
330    }
331
332    fn cleanup_leaves(&mut self, index: usize) {
333        let (left, right) = self.node_children(index);
334        if left != 0 {
335            self.cleanup_leaves(left);
336        }
337        if right != 0 {
338            self.cleanup_leaves(right)
339        }
340
341        let has_children = left != 0 || right != 0;
342        if has_children && self.can_be_leaf(index) {
343            self.update_node(index).map(|updater| updater.leaf());
344            return;
345        }
346
347        if has_children && self.is_leaf(left) && self.is_leaf(right) {
348            // Two leaves that predict the same class are one leaf. Both
349            // outputs have to be set: two unset ones are equal too.
350            if let (Some(output), Some(other)) = (self.node_output(left), self.node_output(right)) {
351                if output == other {
352                    self.update_node(index)
353                        .map(|updater| updater.output(output).clean_test().leaf());
354                }
355            }
356        }
357    }
358
359    fn is_leaf(&self, index: usize) -> bool {
360        let (left, right) = self.node_children(index);
361        (left == 0) && (right == 0)
362    }
363
364    fn can_be_leaf(&self, index: usize) -> bool {
365        self.node_test(index).is_none() && self.node_output(index).is_some()
366    }
367}
368
369/// One node per line, indented by depth, children after their parent.
370impl fmt::Display for Tree {
371    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
372        let mut stack: Vec<(usize, Option<&TreeNode>)> = Vec::new();
373        stack.push((0, self.get_node(self.get_root_index())));
374        while let Some((depth, node)) = stack.pop() {
375            if let Some(node) = node {
376                writeln!(f, "{}----{:?}", "    ".repeat(depth), node.value)?;
377                stack.push((depth + 1, self.get_right_child(node)));
378                stack.push((depth + 1, self.get_left_child(node)));
379            }
380        }
381        Ok(())
382    }
383}
384
385/// Chained setters for one node of a [`Tree`].
386pub struct NodeUpdater<'a> {
387    node: &'a mut TreeNode,
388}
389
390impl<'a> NodeUpdater<'a> {
391    pub fn new(node: &'a mut TreeNode) -> Self {
392        Self { node }
393    }
394
395    pub fn value(self, value: NodeInfos) -> Self {
396        self.node.value = value;
397        self
398    }
399
400    pub fn error(self, error: f64) -> Self {
401        self.node.value.error = error;
402        self
403    }
404
405    pub fn metric(self, metric: f64) -> Self {
406        self.node.value.metric = Some(metric);
407        self
408    }
409
410    pub fn output(self, output: f64) -> Self {
411        self.node.value.out = Some(output);
412        self
413    }
414
415    pub fn test(self, test: usize) -> Self {
416        self.node.value.test = Some(test);
417        self
418    }
419
420    pub fn left_child(self, index: usize) -> Self {
421        self.node.left = index;
422        self
423    }
424
425    pub fn right_child(self, index: usize) -> Self {
426        self.node.right = index;
427        self
428    }
429
430    /// Detaches the node's children. The tested feature is kept; use
431    /// [`Self::clean_test`] to clear it.
432    pub fn leaf(self) -> Self {
433        self.node.left = 0;
434        self.node.right = 0;
435        self
436    }
437
438    pub fn clean_test(self) -> Self {
439        self.node.value.test = None;
440        self
441    }
442
443    pub fn get_children(&self) -> (usize, usize) {
444        (self.node.left, self.node.right)
445    }
446}
447
448#[cfg(test)]
449mod binary_tree_test {
450    use crate::tree::{NodeInfos, Tree, TreeNode};
451
452    #[test]
453    fn create_node_data() {
454        let data = NodeInfos::new();
455        assert_eq!(data.error, <f64>::INFINITY);
456        assert!(data.test.is_none());
457        assert_eq!(data.out, None);
458    }
459
460    #[test]
461    fn create_tree_node() {
462        let data = NodeInfos::new();
463        let node = TreeNode::new(data);
464        assert_eq!(node.right, 0);
465        assert_eq!(node.left, 0);
466        assert_eq!(node.index, 0);
467    }
468
469    #[test]
470    fn tree_default() {
471        let tree = Tree::default();
472        assert_eq!(tree.len(), 0);
473    }
474
475    #[test]
476    fn tree_new() {
477        let tree = Tree::default();
478        assert_eq!(tree.len(), 0);
479    }
480
481    #[test]
482    fn tree_is_empty() {
483        let mut tree = Tree::new();
484        assert!(tree.is_empty());
485
486        let root = TreeNode::new(NodeInfos::default());
487        tree.add_root(root);
488        assert!(!tree.is_empty());
489    }
490
491    #[test]
492    fn binarytree_add_root() {
493        let mut tree: Tree = Tree::new();
494        let root = TreeNode::new(NodeInfos::default());
495        let root_index = tree.add_root(root);
496        assert_eq!(0, root_index);
497    }
498
499    #[test]
500    fn binarytree_get_root_index() {
501        let mut tree = Tree::new();
502        let root = TreeNode::new(NodeInfos::default());
503        let _ = tree.add_root(root);
504        let root_index = tree.get_root_index();
505        assert_eq!(0, root_index);
506    }
507
508    #[test]
509    fn binarytree_get_left_child() {
510        let mut tree = Tree::new();
511        let root = TreeNode::new(NodeInfos::default());
512        let root_index = tree.add_root(root);
513        let node_infos = NodeInfos {
514            test: Some(15),
515            error: 0.0,
516            metric: None,
517            out: None,
518        };
519        let left_node = TreeNode::new(node_infos);
520        let _ = tree.add_left_node(root_index, left_node);
521        let root = tree.get_node(root_index).unwrap();
522        let left_node = tree.get_left_child(root).unwrap();
523        assert_eq!(left_node.value.test, Some(15));
524    }
525
526    #[test]
527    fn binarytree_get_right_child() {
528        let mut tree = Tree::new();
529        let root = TreeNode::new(NodeInfos::default());
530        let root_index = tree.add_root(root);
531        let node_infos = NodeInfos {
532            test: Some(15),
533            error: 0.0,
534            metric: None,
535            out: None,
536        };
537        let right_node = TreeNode::new(node_infos);
538        let _ = tree.add_right_node(root_index, right_node);
539        let root = tree.get_node(root_index).unwrap();
540        let right_node = tree.get_right_child(root).unwrap();
541        assert_eq!(right_node.value.test, Some(15));
542    }
543
544    #[test]
545    fn test_get_node() {
546        let mut tree = Tree::new();
547        let node_infos = NodeInfos {
548            test: Some(33),
549            error: 0.0,
550            metric: None,
551            out: None,
552        };
553        let root = TreeNode::new(node_infos);
554        let _ = tree.add_root(root);
555        let root_index = tree.get_root_index();
556        let root = tree.get_node(root_index).unwrap();
557        assert_eq!(Some(33), root.value.test)
558    }
559
560    #[test]
561    fn test_get_node_mut() {
562        let mut tree = Tree::new();
563        let node_infos = NodeInfos {
564            test: Some(55),
565            error: 0.0,
566            metric: None,
567            out: None,
568        };
569        let root = TreeNode::new(node_infos);
570        let _ = tree.add_root(root);
571        let root_index = tree.get_root_index();
572        let root = tree.get_node_mut(root_index).unwrap();
573        root.value.test = Some(12);
574        assert_eq!(Some(12), root.value.test);
575    }
576
577    #[test]
578    fn test_add_left_node() {
579        let mut tree = Tree::new();
580        let node_infos = NodeInfos {
581            test: Some(15),
582            error: 0.0,
583            metric: None,
584            out: None,
585        };
586        let root = TreeNode::new(node_infos);
587        let root_index = tree.add_root(root);
588        let node_infos = NodeInfos {
589            test: Some(11),
590            error: 0.0,
591            metric: None,
592            out: None,
593        };
594        let left_node = TreeNode::new(node_infos);
595        let _ = tree.add_left_node(root_index, left_node);
596        let root = tree.get_node(root_index).unwrap();
597        let left_node = tree.get_left_child(root).unwrap();
598        assert_eq!(left_node.value.test, Some(11));
599    }
600
601    #[test]
602    fn test_add_right_node() {
603        let mut tree = Tree::new();
604        let node_infos = NodeInfos {
605            test: Some(15),
606            error: 0.0,
607            metric: None,
608            out: None,
609        };
610        let root = TreeNode::new(node_infos);
611        let root_index = tree.add_root(root);
612        let node_infos = NodeInfos {
613            test: Some(22),
614            error: 0.0,
615            metric: None,
616            out: None,
617        };
618        let right_node = TreeNode::new(node_infos);
619        let _ = tree.add_right_node(root_index, right_node);
620        let root = tree.get_node(root_index).unwrap();
621        let right_node = tree.get_right_child(root).unwrap();
622        assert_eq!(right_node.value.test, Some(22));
623    }
624}
dtrees_rs/tree/mod.rs

dtrees_rs/tree/
mod.rs
