radiate_gp/collections/graphs/
iter.rs

1use crate::{Node, NodeType, collections::GraphNode};
2use radiate_core::Valid;
3use std::collections::VecDeque;
4
5/// [GraphIterator] is a trait that provides an iterator over any &[[`GraphNode<T>`]]. The iterator is used to
6/// traverse the said nodes in a pseudo-topological order.
7pub trait GraphIterator<'a, T> {
8    fn iter_topological(&'a self) -> GraphTopologicalIterator<'a, T>;
9
10    fn get_nodes_of_type(
11        &'a self,
12        node_type: NodeType,
13    ) -> impl Iterator<Item = &'a GraphNode<T>> + 'a
14    where
15        Self: AsRef<[GraphNode<T>]>,
16        T: 'a,
17    {
18        self.as_ref()
19            .iter()
20            .filter(move |node| node.node_type() == node_type)
21    }
22}
23
24impl<'a, G: AsRef<[GraphNode<T>]>, T> GraphIterator<'a, T> for G {
25    fn iter_topological(&'a self) -> GraphTopologicalIterator<'a, T> {
26        GraphTopologicalIterator::new(self.as_ref())
27    }
28}
29
30/// [GraphIterator] is an iterator that traverses a &[[`GraphNode<T>`]] in sudo-topological order. I say
31/// "sudo-topological" because it is not a true topological order, but rather a topological order
32/// that allows for recurrent connections.
33pub struct GraphTopologicalIterator<'a, T> {
34    graph: &'a [GraphNode<T>],
35    completed: Vec<bool>,
36    index_queue: VecDeque<usize>,
37    pending_index: usize,
38    remaining: usize,
39}
40
41impl<'a, T> GraphTopologicalIterator<'a, T> {
42    /// Create a new `GraphIterator` from a reference to a [`GraphNode<T>`].
43    ///
44    /// # Arguments
45    /// - `graph`: A reference to the `Graph` to iterate over.
46    pub fn new(graph: &'a [GraphNode<T>]) -> Self {
47        let is_valid = !graph.iter().any(|node| !node.is_valid());
48        GraphTopologicalIterator {
49            graph,
50            completed: vec![false; graph.len()],
51            index_queue: VecDeque::with_capacity(graph.len()),
52            pending_index: if is_valid { 0 } else { graph.len() },
53            remaining: if is_valid { graph.len() } else { 0 },
54        }
55    }
56}
57
58/// Implement the `Iterator` trait for [GraphTopologicalIterator].
59/// The `Item` type is a reference to a [GraphNode].
60///
61/// This implementation is a bit more complex than the typical iterator implementation. The iterator
62/// must traverse the graph in a pseudo-topological order. This means that it must iterate over the
63/// nodes in the graph in an order that respects the dependencies between the nodes. We
64/// do this by keeping track of which nodes have been completed and which nodes are pending, it
65/// then iterates over the nodes in the graph, checking the dependencies of each node to determine
66/// if it can be completed. If a node can be completed, it is added to the index queue, which is
67/// used to determine the order in which the nodes are returned by the iterator.
68/// It is a 'pseudo' topological order because it allows for recurrent connections in the graph.
69impl<'a, T> Iterator for GraphTopologicalIterator<'a, T> {
70    type Item = &'a GraphNode<T>;
71
72    #[inline]
73    fn next(&mut self) -> Option<Self::Item> {
74        let mut min_pending_index = self.graph.len();
75        for index in self.pending_index..self.graph.len() {
76            if self.completed[index] {
77                continue;
78            }
79
80            let node = &self.graph[index];
81
82            let mut degree = node.incoming().len();
83            for incoming_index in node.incoming() {
84                let incoming_node = &self.graph[*incoming_index];
85                if self.completed[incoming_node.index()] || incoming_node.is_recurrent() {
86                    degree -= 1;
87                }
88            }
89
90            if degree == 0 {
91                self.completed[node.index()] = true;
92                self.index_queue.push_back(node.index());
93            } else {
94                min_pending_index = std::cmp::min(min_pending_index, node.index());
95            }
96        }
97
98        self.pending_index = min_pending_index;
99        self.index_queue.pop_front().map(|idx| {
100            self.remaining -= 1;
101            &self.graph[idx]
102        })
103    }
104
105    /// Okay this is actually a pretty straightforward implementation,
106    /// but it _does_ have a performance impact that makes it useful to have.
107    fn size_hint(&self) -> (usize, Option<usize>) {
108        (self.remaining, Some(self.remaining))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::collections::{Graph, GraphNode, NodeType};
116    use crate::ops::Op;
117
118    #[test]
119    fn test_graph_iterator() {
120        let graph = Graph::<Op<f64>>::new(vec![
121            GraphNode::from((0, NodeType::Input, Op::var(0))).with_outgoing([2]),
122            GraphNode::from((1, NodeType::Input, Op::var(1))).with_outgoing([2]),
123            GraphNode::from((2, NodeType::Vertex, Op::add()))
124                .with_incoming([0, 1])
125                .with_outgoing([3]),
126            GraphNode::from((3, NodeType::Output, Op::linear())).with_incoming([2]),
127        ]);
128
129        let mut iter = graph.iter_topological();
130
131        assert_eq!(iter.next().unwrap().index(), 0);
132        assert_eq!(iter.next().unwrap().index(), 1);
133        assert_eq!(iter.next().unwrap().index(), 2);
134        assert_eq!(iter.next().unwrap().index(), 3);
135        assert!(iter.next().is_none());
136    }
137
138    #[test]
139    fn test_graph_iterator_recurrent() {
140        let nodes = vec![
141            GraphNode::from((0, NodeType::Input, Op::<f64>::var(0), vec![], vec![2])),
142            GraphNode::from((1, NodeType::Input, Op::<f64>::var(1), vec![], vec![2])),
143            GraphNode::from((2, NodeType::Vertex, Op::<f64>::add(), vec![0, 1], vec![3])),
144            GraphNode::from((3, NodeType::Vertex, Op::<f64>::mul(), vec![2], vec![2])),
145            GraphNode::from((4, NodeType::Output, Op::<f64>::linear(), vec![3], vec![])),
146        ];
147
148        let graph = Graph::new(nodes);
149        let mut iter = graph.iter_topological();
150
151        assert_eq!(iter.next().unwrap().index(), 0);
152        assert_eq!(iter.next().unwrap().index(), 1);
153        assert_eq!(iter.next().unwrap().index(), 2);
154        assert_eq!(iter.next().unwrap().index(), 3);
155        assert_eq!(iter.next().unwrap().index(), 4);
156        assert!(iter.next().is_none());
157    }
158
159    #[test]
160    fn test_graph_iterator_disconnected() {
161        let nodes = vec![
162            GraphNode::from((0, NodeType::Input, Op::<f64>::var(0))).with_outgoing([2]),
163            GraphNode::from((1, NodeType::Input, Op::<f64>::var(1))),
164            GraphNode::from((2, NodeType::Vertex, Op::<f64>::add())).with_incoming([0]),
165            GraphNode::from((3, NodeType::Output, Op::<f64>::linear())).with_incoming([2]),
166        ];
167
168        let results = Graph::new(nodes)
169            .iter_topological()
170            .map(|node| node.index())
171            .collect::<Vec<usize>>();
172
173        assert!(results.is_empty());
174    }
175
176    #[test]
177    fn test_graph_deep_cycles() {
178        let mut graph = Graph::<Op<f32>>::default();
179
180        graph.insert(NodeType::Input, Op::var(0));
181        graph.insert(NodeType::Vertex, Op::diff());
182        graph.insert(NodeType::Output, Op::sigmoid());
183        graph.insert(NodeType::Vertex, Op::div());
184        graph.insert(NodeType::Vertex, Op::pow());
185        graph.insert(NodeType::Edge, Op::weight());
186        graph.insert(NodeType::Edge, Op::identity());
187        graph.insert(NodeType::Vertex, Op::exp());
188        graph.insert(NodeType::Vertex, Op::cos());
189        graph.insert(NodeType::Edge, Op::weight());
190
191        graph.attach(0, 1);
192        graph.attach(1, 1);
193        graph.attach(4, 1);
194        graph.attach(7, 1);
195        graph.attach(1, 2);
196        graph.attach(3, 2);
197        graph.attach(9, 2);
198        graph.attach(0, 3);
199        graph.attach(5, 3);
200        graph.attach(0, 4);
201        graph.attach(8, 4);
202        graph.attach(1, 5);
203        graph.attach(3, 6);
204        graph.attach(4, 7);
205        graph.attach(6, 8);
206        graph.attach(7, 9);
207
208        graph.set_cycles(vec![]);
209
210        let results = graph
211            .iter_topological()
212            .map(|node| node.index())
213            .collect::<Vec<usize>>();
214
215        assert_eq!(results, vec![0, 1, 3, 4, 5, 6, 7, 8, 9, 2]);
216    }
217
218    #[test]
219    fn test_size_hint_tracks_actual_remaining() {
220        let graph = Graph::<Op<f64>>::new(vec![
221            GraphNode::from((0, NodeType::Input, Op::var(0))).with_outgoing([2]),
222            GraphNode::from((1, NodeType::Input, Op::var(1))).with_outgoing([2]),
223            GraphNode::from((2, NodeType::Vertex, Op::add()))
224                .with_incoming([0, 1])
225                .with_outgoing([3]),
226            GraphNode::from((3, NodeType::Output, Op::linear())).with_incoming([2]),
227        ]);
228
229        let mut iter = graph.iter_topological();
230        let mut actual_remaining = 4;
231
232        assert_eq!(iter.size_hint(), (actual_remaining, Some(actual_remaining)));
233
234        iter.next();
235        actual_remaining -= 1;
236        assert_eq!(iter.size_hint(), (actual_remaining, Some(actual_remaining)));
237
238        while iter.next().is_some() {
239            actual_remaining -= 1;
240            assert_eq!(iter.size_hint(), (actual_remaining, Some(actual_remaining)));
241        }
242
243        assert_eq!(actual_remaining, 0);
244        assert_eq!(iter.size_hint(), (0, Some(0)));
245    }
246
247    #[test]
248    fn test_invalid_graph_size_hint_is_zero() {
249        let nodes = vec![
250            GraphNode::from((0, NodeType::Input, Op::<f64>::var(0))).with_outgoing([2]),
251            GraphNode::from((1, NodeType::Input, Op::<f64>::var(1))),
252            GraphNode::from((2, NodeType::Vertex, Op::<f64>::add())).with_incoming([0]),
253            GraphNode::from((3, NodeType::Output, Op::<f64>::linear())).with_incoming([2]),
254        ];
255
256        let graph = Graph::new(nodes);
257        let iter = graph.iter_topological();
258
259        assert_eq!(iter.size_hint(), (0, Some(0)));
260    }
261
262    #[test]
263    fn test_eval_order_collect_allocates_exactly_once() {
264        let mut graph = Graph::<Op<f32>>::default();
265
266        graph.insert(NodeType::Input, Op::var(0));
267        graph.insert(NodeType::Vertex, Op::diff());
268        graph.insert(NodeType::Output, Op::sigmoid());
269        graph.insert(NodeType::Vertex, Op::div());
270        graph.insert(NodeType::Vertex, Op::pow());
271        graph.insert(NodeType::Edge, Op::weight());
272        graph.insert(NodeType::Edge, Op::identity());
273        graph.insert(NodeType::Vertex, Op::exp());
274        graph.insert(NodeType::Vertex, Op::cos());
275        graph.insert(NodeType::Edge, Op::weight());
276
277        graph.attach(0, 1);
278        graph.attach(1, 1);
279        graph.attach(4, 1);
280        graph.attach(7, 1);
281        graph.attach(1, 2);
282        graph.attach(3, 2);
283        graph.attach(9, 2);
284        graph.attach(0, 3);
285        graph.attach(5, 3);
286        graph.attach(0, 4);
287        graph.attach(8, 4);
288        graph.attach(1, 5);
289        graph.attach(3, 6);
290        graph.attach(4, 7);
291        graph.attach(6, 8);
292        graph.attach(7, 9);
293
294        graph.set_cycles(vec![]);
295
296        let eval_order = graph
297            .iter_topological()
298            .map(|n| n.index())
299            .collect::<Vec<usize>>();
300
301        assert_eq!(eval_order.len(), graph.len());
302        assert_eq!(eval_order.capacity(), eval_order.len());
303    }
304}
radiate_gp/collections/graphs/iter.rs

radiate_gp/collections/graphs/
iter.rs
