Skip to main content

radiate_gp/collections/trees/
eval.rs

1use super::Tree;
2use crate::{Eval, EvalMut, TreeNode, eval::EvalInto, node::Node};
3
4/// Implements the [Eval] trait for [`Tree<T>`] where `T` is `Eval<[V], V>`. All this really does is
5/// call the `eval` method on the root node of the [Tree]. The real work is
6/// done in the [TreeNode] implementation below.
7impl<T, V> Eval<[V], V> for Tree<T>
8where
9    T: Eval<[V], V>,
10    V: Clone,
11{
12    #[inline]
13    fn eval(&self, input: &[V]) -> V {
14        self.root()
15            .map(|root| root.eval(input))
16            .unwrap_or_else(|| panic!("Tree has no root node."))
17    }
18}
19
20/// Implements the [Eval] trait for `Vec<Tree<T>>`. This is a wrapper around a `Vec<Tree<T>>`
21/// and allows for the evaluation of each [Tree] in the `Vec` with a single input.
22/// This is useful for things like `Ensemble` models where multiple models are used to make a prediction.
23///
24/// This is a simple implementation that just maps over the `Vec` and calls [Eval] on each [Tree].
25impl<T, V> Eval<[V], Vec<V>> for Vec<Tree<T>>
26where
27    T: Eval<[V], V>,
28    V: Clone,
29{
30    #[inline]
31    fn eval(&self, inputs: &[V]) -> Vec<V> {
32        self.iter().map(|tree| tree.eval(inputs)).collect()
33    }
34}
35
36impl<T, V> EvalMut<[V], Vec<V>> for Tree<T>
37where
38    T: Eval<[V], V>,
39    V: Clone + Default,
40{
41    #[inline]
42    fn eval_mut(&mut self, input: &[V]) -> Vec<V> {
43        vec![self.eval(input)]
44    }
45}
46
47impl<T, V> EvalInto<[V], [V]> for Tree<T>
48where
49    T: Eval<[V], V>,
50    V: Clone,
51{
52    #[inline]
53    fn eval_into(&self, input: &[V], buffer: &mut [V]) {
54        buffer[0] = self
55            .root()
56            .map(|root| root.eval(input))
57            .unwrap_or_else(|| panic!("Tree has no root node."));
58    }
59}
60
61impl<T, V> EvalInto<[V], [V]> for Vec<Tree<T>>
62where
63    T: Eval<[V], V>,
64    V: Clone,
65{
66    #[inline]
67    fn eval_into(&self, input: &[V], buffer: &mut [V]) {
68        for i in 0..self.len() {
69            buffer[i] = self[i]
70                .root()
71                .map(|root| root.eval(input))
72                .unwrap_or_else(|| panic!("Tree has no root node."));
73        }
74    }
75}
76
77/// Implements the [Eval] trait for `Vec<&TreeNode<T>>`. This is a wrapper around a `Vec<&TreeNode<T>>`
78/// and allows for the evaluation of each [TreeNode] in the `Vec` with a single input.
79/// The len of the input slice must equal the number of nodes in the `Vec`.
80impl<T, V> Eval<[V], Vec<V>> for &[TreeNode<T>]
81where
82    T: Eval<[V], V>,
83    V: Clone,
84{
85    #[inline]
86    fn eval(&self, inputs: &[V]) -> Vec<V> {
87        self.iter().map(|node| node.eval(inputs)).collect()
88    }
89}
90
91impl<T, V> EvalInto<[V], [V]> for &[&TreeNode<T>]
92where
93    T: Eval<[V], V>,
94    V: Clone,
95{
96    #[inline]
97    fn eval_into(&self, input: &[V], buffer: &mut [V]) {
98        for (i, node) in self.iter().enumerate() {
99            buffer[i] = node.eval(input);
100        }
101    }
102}
103
104impl<T, V> Eval<[V], V> for TreeNode<T>
105where
106    T: Eval<[V], V>,
107    V: Clone,
108{
109    #[inline]
110    fn eval(&self, input: &[V]) -> V {
111        (&self).eval(input)
112    }
113}
114
115/// Implements the [Eval] trait for `TreeNode<T>` where `T` is `Eval<[V], V>`. This is where the real work is done.
116/// It recursively evaluates the [TreeNode] and its children until it reaches a leaf node,
117/// at which point it applies the `T`'s eval fn to the input.
118///
119/// Because a [Tree] has only a single root node, this can only be used to return a single value.
120/// We assume here that each leaf can eval the incoming input - this is a safe and the
121/// only real logical assumption we can make.
122impl<T, V> Eval<[V], V> for &TreeNode<T>
123where
124    T: Eval<[V], V>,
125    V: Clone,
126{
127    #[inline]
128    fn eval(&self, input: &[V]) -> V {
129        if let Some(children) = self.children() {
130            let mut inputs = Vec::with_capacity(children.len());
131
132            for child in children {
133                inputs.push(child.eval(input));
134            }
135
136            return self.value().eval(&inputs);
137        }
138
139        self.value().eval(input)
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::{Op, TreeNode};
147
148    #[test]
149    fn test_tree_reduce_simple() {
150        let mut root = TreeNode::new(Op::add());
151
152        root.add_child(TreeNode::new(Op::constant(1.0)));
153        root.add_child(TreeNode::new(Op::constant(2.0)));
154
155        let result = root.eval(&[]);
156
157        assert_eq!(result, 3.0);
158    }
159
160    #[test]
161    fn test_tree_reduce_complex() {
162        let node = TreeNode::new(Op::add())
163            .attach(
164                TreeNode::new(Op::mul())
165                    .attach(TreeNode::new(Op::constant(2.0)))
166                    .attach(TreeNode::new(Op::constant(3.0))),
167            )
168            .attach(
169                TreeNode::new(Op::add())
170                    .attach(TreeNode::new(Op::constant(2.0)))
171                    .attach(TreeNode::new(Op::var(0))),
172            );
173
174        let nine = node.eval(&[1_f32]);
175        let ten = node.eval(&[2_f32]);
176        let eleven = node.eval(&[3_f32]);
177
178        assert_eq!(nine, 9.0);
179        assert_eq!(ten, 10.0);
180        assert_eq!(eleven, 11.0);
181    }
182}