Skip to main content

radiate_gp/collections/graphs/
eval.rs

1use super::{Graph, GraphNode, iter::GraphIterator};
2use crate::{
3    Eval, EvalMut, NodeType,
4    eval::{EvalInto, EvalIntoMut},
5    node::Node,
6};
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9use std::ops::Range;
10
11/// A cache for storing intermediate results during graph evaluation.
12///
13/// This cache is used to store the inputs and outputs of each node in the graph
14/// during evaluation, allowing for more efficient re-evaluation of nodes when
15/// their inputs change. If we want to save a graph's evaluation between different evals,
16/// we need to keep track of the inputs and outputs from previous runs incase of recurrent
17/// structures. This cache is the answer to that.
18#[derive(Clone, Debug, PartialEq)]
19#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
20pub struct GraphEvalCache<V> {
21    eval_order: Vec<usize>,
22    outputs: Vec<V>,
23    inputs: Vec<V>,
24    input_ranges: Vec<Range<usize>>,
25    output_indices: Vec<usize>,
26}
27
28/// [GraphEvaluator] is a struct that is used to evaluate a [Graph] of [GraphNode]'s. It uses the [GraphIterator]
29/// to traverse the [Graph] in a pseudo-topological order and evaluate the nodes in the correct order.
30pub struct GraphEvaluator<'a, T, V> {
31    nodes: &'a [GraphNode<T>],
32    inner: GraphEvalCache<V>,
33}
34
35impl<'a, T, V> GraphEvaluator<'a, T, V>
36where
37    T: Eval<[V], V>,
38    V: Default + Clone,
39{
40    /// Creates a new [GraphEvaluator] with the given [Graph]. We pre-allocate the necessary
41    /// storage for inputs and outputs based on the structure of the graph on creation.
42    /// This way, we can reuse the same evaluator for multiple evaluations of the same graph
43    /// without needing to reallocate memory each time.
44    ///
45    /// # Arguments
46    /// * graph - The [Graph] to reduce.
47    #[inline]
48    pub fn new<N>(graph: &'a N) -> Self
49    where
50        N: AsRef<[GraphNode<T>]>,
51    {
52        let nodes = graph.as_ref();
53
54        let mut total_inputs = 0;
55        let mut input_ranges = Vec::with_capacity(nodes.len());
56
57        for node in nodes {
58            let k = node.incoming().len();
59            input_ranges.push(total_inputs..total_inputs + k);
60            total_inputs += k;
61        }
62
63        let output_indices = graph
64            .get_nodes_of_type(NodeType::Output)
65            .map(|n| n.index())
66            .collect::<Vec<usize>>();
67
68        GraphEvaluator {
69            nodes,
70            inner: GraphEvalCache {
71                inputs: vec![V::default(); total_inputs],
72                outputs: vec![V::default(); nodes.len()],
73                eval_order: nodes.iter_topological().map(|n| n.index()).collect(),
74                input_ranges,
75                output_indices,
76            },
77        }
78    }
79
80    pub fn take_cache(self) -> GraphEvalCache<V> {
81        self.inner
82    }
83}
84
85impl<T, V> EvalMut<[V], Vec<V>> for GraphEvaluator<'_, T, V>
86where
87    T: Eval<[V], V>,
88    V: Copy + Default,
89{
90    #[inline]
91    fn eval_mut(&mut self, input: &[V]) -> Vec<V> {
92        let out_len = self.inner.output_indices.len();
93        let mut buffer = vec![V::default(); out_len];
94        self.eval_into_mut(input, &mut buffer[..]);
95        buffer
96    }
97}
98
99impl<T, V> EvalIntoMut<[V], [V]> for GraphEvaluator<'_, T, V>
100where
101    T: Eval<[V], V>,
102    V: Copy + Default,
103{
104    #[inline]
105    fn eval_into_mut(&mut self, input: &[V], buffer: &mut [V]) {
106        for &index in self.inner.eval_order.iter() {
107            let node = &self.nodes[index];
108            let incoming = node.incoming();
109
110            if incoming.is_empty() {
111                self.inner.outputs[index] = node.eval(input);
112            } else {
113                let range = &self.inner.input_ranges[index];
114                let buf = &mut self.inner.inputs[range.clone()];
115
116                for (dst, &src_idx) in buf.iter_mut().zip(incoming.iter()) {
117                    *dst = self.inner.outputs[src_idx];
118                }
119
120                self.inner.outputs[index] = node.eval(buf);
121            }
122        }
123
124        let mut count = 0;
125        for &idx in self.inner.output_indices.iter() {
126            buffer[count] = self.inner.outputs[idx];
127            count += 1;
128        }
129    }
130}
131
132impl<T, V> EvalInto<[Vec<V>], Vec<Vec<V>>> for Graph<T>
133where
134    T: Eval<[V], V>,
135    V: Copy + Default,
136{
137    /// Evaluates the [Graph] with the given input `Vec<Vec<T>>`. Returns the output of the [Graph] as `Vec<Vec<T>>`.
138    ///
139    /// # Arguments
140    /// * `input` - A `Vec<Vec<T>>` to evaluate the [Graph] with.
141    ///
142    /// # Returns
143    /// * A `Vec<Vec<T>>` which is the output of the [Graph].
144    #[inline]
145    fn eval_into(&self, input: &[Vec<V>], buffer: &mut Vec<Vec<V>>) {
146        let mut evaluator = GraphEvaluator::new(self);
147        for i in 0..input.len() {
148            evaluator.eval_into_mut(&input[i], &mut buffer[i]);
149        }
150    }
151}
152
153impl<T, V> Eval<[Vec<V>], Vec<Vec<V>>> for Graph<T>
154where
155    T: Eval<[V], V>,
156    V: Copy + Default,
157{
158    /// Evaluates the [Graph] with the given input `Vec<Vec<T>>`. Returns the output of the [Graph] as `Vec<Vec<T>>`.
159    /// This is intended to be used when evaluating a batch of inputs.
160    ///
161    /// # Arguments
162    /// * `input` - A `Vec<Vec<T>>` to evaluate the [Graph] with.
163    ///
164    /// # Returns
165    /// * A `Vec<Vec<T>>` which is the output of the [Graph].
166    #[inline]
167    fn eval(&self, input: &[Vec<V>]) -> Vec<Vec<V>> {
168        let mut evaluator = GraphEvaluator::new(self);
169        input
170            .iter()
171            .map(|input| evaluator.eval_mut(input))
172            .collect()
173    }
174}
175
176impl<T, V> Eval<[V], V> for GraphNode<T>
177where
178    T: Eval<[V], V>,
179    V: Copy,
180{
181    /// Evaluates the [GraphNode] with the given input. Returns the output of the [GraphNode].
182    /// # Arguments
183    /// * `inputs` - A `Vec` of `V` to evaluate the [GraphNode] with.
184    ///
185    /// # Returns
186    /// * A `V` which is the output of the [GraphNode].
187    #[inline]
188    fn eval(&self, inputs: &[V]) -> V {
189        self.value().eval(inputs)
190    }
191}
192
193impl<'a, G, T, V> From<(&'a G, GraphEvalCache<V>)> for GraphEvaluator<'a, T, V>
194where
195    G: AsRef<[GraphNode<T>]>,
196    T: Eval<[V], V>,
197    V: Default + Clone,
198{
199    fn from((graph, cache): (&'a G, GraphEvalCache<V>)) -> Self {
200        if cache.eval_order.is_empty() || graph.as_ref().len() != cache.eval_order.len() {
201            return GraphEvaluator::new(graph);
202        }
203
204        GraphEvaluator {
205            nodes: graph.as_ref(),
206            inner: cache,
207        }
208    }
209}
210
211impl<'a, T, V> From<&'a Graph<T>> for GraphEvaluator<'a, T, V>
212where
213    T: Eval<[V], V>,
214    V: Default + Clone,
215{
216    fn from(graph: &'a Graph<T>) -> Self {
217        GraphEvaluator::new(graph)
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::{Graph, Op};
225
226    fn round(value: f32, places: u32) -> f32 {
227        let factor = 10_f32.powi(places as i32);
228        (value * factor).round() / factor
229    }
230
231    #[test]
232    fn test_graph_eval_simple() {
233        let mut graph = Graph::<Op<f32>>::default();
234
235        let idx_one = graph.insert(NodeType::Input, Op::var(0));
236        let idx_two = graph.insert(NodeType::Input, Op::constant(5_f32));
237        let idx_three = graph.insert(NodeType::Vertex, Op::add());
238        let idx_four = graph.insert(NodeType::Output, Op::linear());
239
240        graph
241            .attach(idx_one, idx_three)
242            .attach(idx_two, idx_three)
243            .attach(idx_three, idx_four);
244
245        let six = graph.eval(&[vec![1_f32]]);
246        let seven = graph.eval(&[vec![2_f32]]);
247        let eight = graph.eval(&[vec![3_f32]]);
248
249        assert_eq!(six, vec![vec![6_f32]]);
250        assert_eq!(seven, vec![vec![7_f32]]);
251        assert_eq!(eight, vec![vec![8_f32]]);
252        assert_eq!(graph.len(), 4);
253    }
254
255    #[test]
256    fn test_graph_eval_recurrent() {
257        let mut graph = Graph::<Op<f32>>::default();
258
259        graph.insert(NodeType::Input, Op::var(0));
260        graph.insert(NodeType::Vertex, Op::diff());
261        graph.insert(NodeType::Output, Op::sigmoid());
262        graph.insert(NodeType::Edge, Op::weight_with(-1.41));
263        graph.insert(NodeType::Vertex, Op::sigmoid());
264        graph.insert(NodeType::Vertex, Op::exp());
265        graph.insert(NodeType::Edge, Op::weight_with(-1.10));
266        graph.insert(NodeType::Vertex, Op::exp());
267        graph.insert(NodeType::Vertex, Op::exp());
268        graph.insert(NodeType::Vertex, Op::div());
269
270        graph.attach(0, 1);
271        graph.attach(1, 1);
272        graph.attach(4, 1);
273        graph.attach(7, 1);
274        graph.attach(8, 1);
275        graph.attach(1, 2);
276        graph.attach(3, 2);
277        graph.attach(6, 2);
278        graph.attach(5, 3);
279        graph.attach(1, 4);
280        graph.attach(0, 5);
281        graph.attach(9, 6);
282        graph.attach(4, 7);
283        graph.attach(7, 8);
284        graph.attach(0, 9);
285        graph.attach(9, 9);
286
287        graph.set_cycles(vec![]);
288
289        let mut evaluator = GraphEvaluator::new(&graph);
290
291        let out1 = evaluator.eval_mut(&vec![0.0])[0];
292        let out2 = evaluator.eval_mut(&vec![0.0])[0];
293        let out3 = evaluator.eval_mut(&vec![0.0])[0];
294        let out4 = evaluator.eval_mut(&vec![1.0])[0];
295        let out5 = evaluator.eval_mut(&vec![0.0])[0];
296        let out6 = evaluator.eval_mut(&vec![0.0])[0];
297        let out7 = evaluator.eval_mut(&vec![0.0])[0];
298
299        assert_eq!(round(out1, 3), 0.196);
300        assert_eq!(round(out2, 3), 0.000);
301        assert_eq!(round(out3, 3), 0.902);
302        assert_eq!(round(out4, 3), 0.000);
303        assert_eq!(round(out5, 3), 1.000);
304        assert_eq!(round(out6, 3), 0.000);
305        assert_eq!(round(out7, 3), 1.000);
306    }
307}