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