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        for (count, &idx) in self.inner.output_indices.iter().enumerate() {
125            buffer[count] = self.inner.outputs[idx];
126        }
127    }
128}
129
130impl<T, V> EvalInto<[Vec<V>], Vec<Vec<V>>> for Graph<T>
131where
132    T: Eval<[V], V>,
133    V: Copy + Default,
134{
135    /// Evaluates the [Graph] with the given input `Vec<Vec<T>>`. Returns the output of the [Graph] as `Vec<Vec<T>>`.
136    ///
137    /// # Arguments
138    /// * `input` - A `Vec<Vec<T>>` to evaluate the [Graph] with.
139    ///
140    /// # Returns
141    /// * A `Vec<Vec<T>>` which is the output of the [Graph].
142    #[inline]
143    fn eval_into(&self, input: &[Vec<V>], buffer: &mut Vec<Vec<V>>) {
144        let mut evaluator = GraphEvaluator::new(self);
145        for i in 0..input.len() {
146            evaluator.eval_into_mut(&input[i], &mut buffer[i]);
147        }
148    }
149}
150
151impl<T, V> Eval<[Vec<V>], Vec<Vec<V>>> for Graph<T>
152where
153    T: Eval<[V], V>,
154    V: Copy + Default,
155{
156    /// Evaluates the [Graph] with the given input `Vec<Vec<T>>`. Returns the output of the [Graph] as `Vec<Vec<T>>`.
157    /// This is intended to be used when evaluating a batch of inputs.
158    ///
159    /// # Arguments
160    /// * `input` - A `Vec<Vec<T>>` to evaluate the [Graph] with.
161    ///
162    /// # Returns
163    /// * A `Vec<Vec<T>>` which is the output of the [Graph].
164    #[inline]
165    fn eval(&self, input: &[Vec<V>]) -> Vec<Vec<V>> {
166        let mut evaluator = GraphEvaluator::new(self);
167        input
168            .iter()
169            .map(|input| evaluator.eval_mut(input))
170            .collect()
171    }
172}
173
174impl<T, V> Eval<[V], V> for GraphNode<T>
175where
176    T: Eval<[V], V>,
177    V: Copy,
178{
179    /// Evaluates the [GraphNode] with the given input. Returns the output of the [GraphNode].
180    /// # Arguments
181    /// * `inputs` - A `Vec` of `V` to evaluate the [GraphNode] with.
182    ///
183    /// # Returns
184    /// * A `V` which is the output of the [GraphNode].
185    #[inline]
186    fn eval(&self, inputs: &[V]) -> V {
187        self.value().eval(inputs)
188    }
189}
190
191impl<'a, G, T, V> From<(&'a G, GraphEvalCache<V>)> for GraphEvaluator<'a, T, V>
192where
193    G: AsRef<[GraphNode<T>]>,
194    T: Eval<[V], V>,
195    V: Default + Clone,
196{
197    fn from((graph, cache): (&'a G, GraphEvalCache<V>)) -> Self {
198        if cache.eval_order.is_empty() || graph.as_ref().len() != cache.eval_order.len() {
199            return GraphEvaluator::new(graph);
200        }
201
202        GraphEvaluator {
203            nodes: graph.as_ref(),
204            inner: cache,
205        }
206    }
207}
208
209impl<'a, T, V> From<&'a Graph<T>> for GraphEvaluator<'a, T, V>
210where
211    T: Eval<[V], V>,
212    V: Default + Clone,
213{
214    fn from(graph: &'a Graph<T>) -> Self {
215        GraphEvaluator::new(graph)
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use crate::{Graph, Op};
223
224    fn round(value: f32, places: u32) -> f32 {
225        let factor = 10_f32.powi(places as i32);
226        (value * factor).round() / factor
227    }
228
229    #[test]
230    fn test_graph_eval_simple() {
231        let mut graph = Graph::<Op<f32>>::default();
232
233        let idx_one = graph.insert(NodeType::Input, Op::var(0));
234        let idx_two = graph.insert(NodeType::Input, Op::constant(5_f32));
235        let idx_three = graph.insert(NodeType::Vertex, Op::add());
236        let idx_four = graph.insert(NodeType::Output, Op::linear());
237
238        graph
239            .attach(idx_one, idx_three)
240            .attach(idx_two, idx_three)
241            .attach(idx_three, idx_four);
242
243        let six = graph.eval(&[vec![1_f32]]);
244        let seven = graph.eval(&[vec![2_f32]]);
245        let eight = graph.eval(&[vec![3_f32]]);
246
247        assert_eq!(six, vec![vec![6_f32]]);
248        assert_eq!(seven, vec![vec![7_f32]]);
249        assert_eq!(eight, vec![vec![8_f32]]);
250        assert_eq!(graph.len(), 4);
251    }
252
253    #[test]
254    fn test_graph_eval_recurrent() {
255        let mut graph = Graph::<Op<f32>>::default();
256
257        graph.insert(NodeType::Input, Op::var(0));
258        graph.insert(NodeType::Vertex, Op::diff());
259        graph.insert(NodeType::Output, Op::sigmoid());
260        graph.insert(NodeType::Edge, Op::weight_with(-1.41));
261        graph.insert(NodeType::Vertex, Op::sigmoid());
262        graph.insert(NodeType::Vertex, Op::exp());
263        graph.insert(NodeType::Edge, Op::weight_with(-1.10));
264        graph.insert(NodeType::Vertex, Op::exp());
265        graph.insert(NodeType::Vertex, Op::exp());
266        graph.insert(NodeType::Vertex, Op::div());
267
268        graph.attach(0, 1);
269        graph.attach(1, 1);
270        graph.attach(4, 1);
271        graph.attach(7, 1);
272        graph.attach(8, 1);
273        graph.attach(1, 2);
274        graph.attach(3, 2);
275        graph.attach(6, 2);
276        graph.attach(5, 3);
277        graph.attach(1, 4);
278        graph.attach(0, 5);
279        graph.attach(9, 6);
280        graph.attach(4, 7);
281        graph.attach(7, 8);
282        graph.attach(0, 9);
283        graph.attach(9, 9);
284
285        graph.set_cycles(vec![]);
286
287        let mut evaluator = GraphEvaluator::new(&graph);
288
289        let out1 = evaluator.eval_mut(&vec![0.0])[0];
290        let out2 = evaluator.eval_mut(&vec![0.0])[0];
291        let out3 = evaluator.eval_mut(&vec![0.0])[0];
292        let out4 = evaluator.eval_mut(&vec![1.0])[0];
293        let out5 = evaluator.eval_mut(&vec![0.0])[0];
294        let out6 = evaluator.eval_mut(&vec![0.0])[0];
295        let out7 = evaluator.eval_mut(&vec![0.0])[0];
296
297        assert_eq!(round(out1, 3), 0.196);
298        assert_eq!(round(out2, 3), 0.000);
299        assert_eq!(round(out3, 3), 0.902);
300        assert_eq!(round(out4, 3), 0.000);
301        assert_eq!(round(out5, 3), 1.000);
302        assert_eq!(round(out6, 3), 0.000);
303        assert_eq!(round(out7, 3), 1.000);
304    }
305}