radiate_gp/collections/graphs/
builder.rs

1use super::aggregate::GraphAggregate;
2use crate::{
3    Arity, Factory, NodeStore,
4    collections::{Graph, GraphNode, NodeType},
5};
6
7impl<T: Clone + Default> Graph<T> {
8    /// Creates a directed graph with the given input, output sizes and values.
9    /// The values are used to initialize the nodes in the graph with the given values.
10    ///
11    /// # Example
12    /// ```
13    /// use radiate_gp::*;
14    ///
15    /// let values = vec![
16    ///     (NodeType::Input, vec![Op::var(0), Op::var(1), Op::var(2)]),
17    ///     (NodeType::Output, vec![Op::sigmoid()]),
18    /// ];
19    ///
20    /// let graph = Graph::directed(3, 3, values);
21    ///
22    /// assert_eq!(graph.len(), 6);
23    /// ```
24    ///
25    /// The graph will have 6 nodes, 3 input nodes and 3 output nodes where each input node is
26    /// connected to each output node. Such as:
27    /// ``` text
28    /// [0, 1, 2] -> [3, 4, 5]
29    /// ```
30    ///
31    /// # Arguments
32    /// * `input_size` - The number of input nodes.
33    /// * `output_size` - The number of output nodes.
34    /// * `values` - The values to initialize the nodes with.
35    ///
36    /// # Returns
37    /// A new directed graph.
38    pub fn directed(
39        input_size: usize,
40        output_size: usize,
41        values: impl Into<NodeStore<T>>,
42    ) -> Graph<T> {
43        let builder = NodeBuilder::new(values);
44
45        let input_nodes = builder.input(input_size);
46        let output_nodes = builder.output(output_size);
47
48        GraphAggregate::new()
49            .all_to_all(&input_nodes, &output_nodes)
50            .build()
51    }
52
53    /// Creates a recurrent graph with the given input and output sizes.
54    /// The values are used to initialize the nodes in the graph with the given values.
55    /// The graph will have a recurrent connection from each hidden vertex to itself.
56    /// The graph will have a one-to-one connection from each input node to each hidden vertex.
57    /// The graph will have an all-to-all connection from each hidden vertex to each output node.
58    ///
59    /// # Example
60    /// ```
61    /// use radiate_gp::*;
62    ///
63    /// let values = vec![
64    ///   (NodeType::Input, vec![Op::var(0), Op::var(1), Op::var(2)]),
65    ///   (NodeType::Vertex, vec![Op::linear()]),
66    ///   (NodeType::Output, vec![Op::sigmoid()]),
67    /// ];
68    ///
69    /// let graph = Graph::recurrent(3, 3, values);
70    ///
71    /// assert_eq!(graph.len(), 9);
72    /// ```
73    ///
74    /// The graph will have 9 nodes, 3 input nodes, 3 hidden nodes with recurrent connections to themselves,
75    /// and 3 output nodes. Such as:
76    /// ``` text
77    /// [0, 1, 2] -> [3, 4, 5]
78    ///     [3, 4, 5] -> [6, 7, 8]
79    ///         [6, 7, 8] -> [3, 4, 5]
80    /// [3, 4, 5] -> [9, 10, 11]
81    /// ```
82    ///
83    /// # Arguments
84    /// * `input_size` - The number of input nodes.
85    /// * `output_size` - The number of output nodes.
86    /// * `values` - The values to initialize the nodes with.
87    ///
88    /// # Returns
89    /// A new recurrent graph.
90    pub fn recurrent(
91        input_size: usize,
92        output_size: usize,
93        values: impl Into<NodeStore<T>>,
94    ) -> Graph<T> {
95        let builder = NodeBuilder::new(values);
96
97        let input = builder.input(input_size);
98        let aggregate = builder.vertices(input_size);
99        let output = builder.output(output_size);
100
101        GraphAggregate::new()
102            .one_to_one(&input, &aggregate)
103            .cycle(&aggregate)
104            .all_to_all(&aggregate, &output)
105            .build()
106    }
107
108    /// Creates a weighted directed graph with the given input and output sizes.
109    ///
110    /// This will result in the same graph as `Graph::directed` but with an additional edge
111    /// connecting each input node to each output node.
112    ///
113    /// # Arguments
114    /// * `input_size` - The number of input nodes.
115    /// * `output_size` - The number of output nodes.
116    ///
117    /// # Returns
118    /// A new weighted directed graph.
119    pub fn weighted_directed(
120        input_size: usize,
121        output_size: usize,
122        values: impl Into<NodeStore<T>>,
123    ) -> Graph<T> {
124        let builder = NodeBuilder::new(values);
125
126        let input = builder.input(input_size);
127        let output = builder.output(output_size);
128        let weights = builder.edge(input_size * output_size);
129
130        GraphAggregate::new()
131            .one_to_many(&input, &weights)
132            .many_to_one(&weights, &output)
133            .build()
134    }
135
136    /// Creates a weighted recurrent graph with the given input and output sizes.
137    /// This will result in the same graph as `Graph::recurrent` but with an additional edge
138    /// connecting each hidden vertex to each output node.
139    ///
140    /// # Arguments
141    /// * `input_size` - The number of input nodes.
142    /// * `output_size` - The number of output nodes.
143    ///
144    /// # Returns
145    /// A new weighted recurrent graph.
146    pub fn weighted_recurrent(
147        input_size: usize,
148        output_size: usize,
149        values: impl Into<NodeStore<T>>,
150    ) -> Graph<T> {
151        let builder = NodeBuilder::new(values);
152
153        let input = builder.input(input_size);
154        let aggregate = builder.vertices(input_size);
155        let output = builder.output(output_size);
156        let weights = builder.edge(input_size * output_size);
157
158        GraphAggregate::new()
159            .one_to_one(&input, &aggregate)
160            .cycle(&aggregate)
161            .one_to_many(&aggregate, &weights)
162            .many_to_one(&weights, &output)
163            .build()
164    }
165
166    /// Creates a Long Short-Term Memory (LSTM) graph with the given input and output sizes.
167    /// The graph will have the following structure:
168    /// - Input nodes connected to forget, input, candidate, and output gates.
169    /// - Hidden state connected to forget, input, candidate, and output gates.
170    /// - Forget gate connected to cell state.
171    /// - Input gate connected to candidate and cell state.
172    /// - Candidate connected to cell state.
173    /// - Cell state connected to hidden state.
174    /// - Output gate connected to hidden state.
175    /// - Hidden state connected to output nodes.
176    ///
177    /// # Arguments
178    /// * `input_size` - The number of input nodes.
179    /// * `output_size` - The number of output nodes.
180    /// * `store` - The node store.
181    ///
182    /// # Returns
183    /// A new LSTM graph.
184    pub fn lstm(input_size: usize, output_size: usize, store: impl Into<NodeStore<T>>) -> Graph<T> {
185        let builder = NodeBuilder::new(store);
186
187        let input = builder.input(input_size);
188        let output = builder.output(output_size);
189
190        let cell_state = builder.vertices_with_arity(1, Arity::Any);
191        let hidden_state = builder.vertices_with_arity(1, Arity::Any);
192
193        let forget_gate = builder.vertices_with_arity(1, Arity::Any);
194        let input_gate = builder.vertices_with_arity(1, Arity::Any);
195        let output_gate = builder.vertices_with_arity(1, Arity::Any);
196        let candidate = builder.vertices_with_arity(1, Arity::Any);
197
198        GraphAggregate::new()
199            .all_to_all(&input, &forget_gate)
200            .all_to_all(&input, &input_gate)
201            .all_to_all(&input, &output_gate)
202            .all_to_all(&input, &candidate)
203            .one_to_one(&hidden_state, &forget_gate)
204            .one_to_one(&hidden_state, &input_gate)
205            .one_to_one(&hidden_state, &output_gate)
206            .one_to_one(&hidden_state, &candidate)
207            .one_to_one(&forget_gate, &cell_state)
208            .one_to_one(&input_gate, &candidate)
209            .one_to_one(&candidate, &cell_state)
210            .one_to_one(&cell_state, &hidden_state)
211            .one_to_one(&output_gate, &hidden_state)
212            .all_to_all(&hidden_state, &output)
213            .build()
214    }
215
216    /// Creates a Gated Recurrent Unit (GRU) graph with the given input and output sizes.
217    /// The graph will have the following structure:
218    /// - Input nodes connected to reset, update, and candidate gates.
219    /// - Hidden state connected to reset, update, and candidate gates.
220    /// - Reset gate connected to hidden state.
221    /// - Update gate connected to blend and gate flip.
222    /// - Candidate connected to blend.
223    /// - Blend connected to hidden state.
224    /// - Gate flip connected to hidden state.
225    /// - Hidden state connected to output nodes.
226    ///
227    /// # Arguments
228    /// * `input_size` - The number of input nodes.
229    /// * `output_size` - The number of output nodes.
230    /// * `store` - The node store.
231    ///
232    /// # Returns
233    /// A new GRU graph.
234    pub fn gru(input_size: usize, output_size: usize, values: impl Into<NodeStore<T>>) -> Graph<T> {
235        let builder = NodeBuilder::new(values);
236
237        let input = builder.input(input_size);
238        let output = builder.output(output_size);
239
240        let hidden = builder.vertices_with_arity(1, Arity::Any);
241
242        let update = builder.vertices_with_arity(1, Arity::Any);
243        let reset = builder.vertices_with_arity(1, Arity::Any);
244        let candidate = builder.vertices_with_arity(1, Arity::Any);
245
246        let blend = builder.vertices_with_arity(1, Arity::Any);
247        let gate_flip = builder.vertices_with_arity(1, Arity::Any);
248
249        GraphAggregate::new()
250            .many_to_one(&input, &reset)
251            .many_to_one(&input, &update)
252            .many_to_one(&input, &candidate)
253            .one_to_one(&hidden, &reset)
254            .one_to_one(&hidden, &update)
255            .one_to_one(&hidden, &candidate)
256            .one_to_one(&update, &blend)
257            .one_to_one(&candidate, &blend)
258            .one_to_one(&reset, &hidden)
259            .one_to_one(&update, &gate_flip)
260            .one_to_one(&hidden, &gate_flip)
261            .one_to_one(&gate_flip, &hidden)
262            .one_to_one(&blend, &hidden)
263            .one_to_many(&hidden, &output)
264            .build()
265    }
266
267    /// Creates a 2D mesh graph with bidirectional connections between neighboring nodes.
268    /// The graph will have the following structure:
269    /// - Input nodes connected to the first row of mesh nodes.
270    /// - Each mesh node connected to its neighbors (up, down, left, right).
271    /// - Last row of mesh nodes connected to output nodes.
272    ///
273    /// # Arguments
274    /// * `width` - The number of nodes in the horizontal dimension.
275    /// * `height` - The number of nodes in the vertical dimension.
276    /// * `values` - The values to initialize the nodes with.
277    ///
278    /// # Returns
279    /// A new 2D mesh graph.
280    pub fn mesh(
281        input_size: usize,
282        output_size: usize,
283        width: usize,
284        height: usize,
285        values: impl Into<NodeStore<T>>,
286    ) -> Graph<T> {
287        let builder = NodeBuilder::new(values);
288
289        let inputs = builder.input(input_size);
290        let outputs = builder.output(output_size);
291        let nodes = (0..width * height)
292            .map(|_| builder.vertices(1))
293            .collect::<Vec<Vec<GraphNode<T>>>>();
294
295        let mut aggregate = GraphAggregate::new();
296
297        for y in 0..height {
298            for x in 0..width {
299                let index = y * width + x;
300                let current = &nodes[index];
301
302                if x + 1 < width {
303                    let right = &nodes[y * width + (x + 1)];
304                    aggregate = aggregate.one_to_one(current, right);
305                }
306
307                if y + 1 < height {
308                    let down = &nodes[(y + 1) * width + x];
309                    aggregate = aggregate.one_to_one(current, down);
310                }
311            }
312        }
313
314        aggregate
315            .many_to_one(&inputs, &nodes[0])
316            .one_to_many(&nodes[nodes.len() - 1], &outputs)
317            .build()
318    }
319}
320
321/// A simple builder struct for constructing nodes of a certain type. This is pretty much just a
322/// quality of life struct that removes boilerplate code when creating collections of nodes.
323pub struct NodeBuilder<T> {
324    store: NodeStore<T>,
325}
326
327impl<T: Clone + Default> NodeBuilder<T> {
328    pub fn new(store: impl Into<NodeStore<T>>) -> Self {
329        NodeBuilder {
330            store: store.into(),
331        }
332    }
333
334    pub fn input(&self, size: usize) -> Vec<GraphNode<T>> {
335        self.new_nodes(NodeType::Input, size, Arity::Zero)
336    }
337
338    pub fn output(&self, size: usize) -> Vec<GraphNode<T>> {
339        self.new_nodes(NodeType::Output, size, Arity::Any)
340    }
341
342    pub fn edge(&self, size: usize) -> Vec<GraphNode<T>> {
343        self.new_nodes(NodeType::Edge, size, Arity::Exact(1))
344    }
345
346    pub fn vertices(&self, size: usize) -> Vec<GraphNode<T>> {
347        self.new_nodes(NodeType::Vertex, size, Arity::Any)
348    }
349
350    pub fn vertices_with_arity(&self, size: usize, arity: Arity) -> Vec<GraphNode<T>> {
351        (0..size)
352            .filter_map(|idx| {
353                self.store
354                    .new_instance((idx, NodeType::Vertex, |a| a == arity))
355            })
356            .collect()
357    }
358
359    fn new_nodes(
360        &self,
361        node_type: NodeType,
362        size: usize,
363        fallback_arity: Arity,
364    ) -> Vec<GraphNode<T>> {
365        if self.store.contains_type(node_type) {
366            (0..size)
367                .filter_map(|idx| self.store.new_instance((idx, node_type)))
368                .collect()
369        } else {
370            (0..size)
371                .filter_map(|idx| {
372                    self.store
373                        .new_instance((idx, node_type, |arity| arity == fallback_arity))
374                })
375                .collect()
376        }
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383    use crate::{Node, Op, node_store};
384    use radiate_core::Valid;
385
386    #[test]
387    fn test_graph_builder() {
388        let graph = Graph::directed(3, 3, Op::sigmoid());
389
390        assert_eq!(graph.len(), 6);
391
392        for node in graph.iter() {
393            if node.node_type() == NodeType::Input {
394                assert_eq!(node.arity(), Arity::Zero);
395                assert_eq!(node.incoming().iter().count(), 0);
396                assert_eq!(node.outgoing().iter().count(), 3);
397            } else if node.node_type() == NodeType::Output {
398                assert_eq!(node.arity(), Arity::Any);
399                assert_eq!(node.incoming().iter().count(), 3);
400                assert_eq!(node.outgoing().iter().count(), 0);
401                assert_eq!(node.value(), &Op::sigmoid());
402            }
403        }
404    }
405
406    #[test]
407    fn test_graph_builder_recurrent() {
408        let graph = Graph::recurrent(3, 3, Op::sigmoid());
409
410        assert_eq!(graph.len(), 9);
411
412        for node in graph.iter() {
413            if node.node_type() == NodeType::Input {
414                assert_eq!(node.arity(), Arity::Zero);
415                assert_eq!(node.incoming().iter().count(), 0);
416                assert_eq!(node.outgoing().iter().count(), 1);
417            } else if node.node_type() == NodeType::Vertex {
418                assert_eq!(node.arity(), Arity::Any);
419                assert!(node.is_recurrent());
420                assert_eq!(node.value(), &Op::sigmoid());
421            } else if node.node_type() == NodeType::Output {
422                assert_eq!(node.arity(), Arity::Any);
423                assert_eq!(node.incoming().iter().count(), 3);
424                assert_eq!(node.outgoing().iter().count(), 0);
425                assert_eq!(node.value(), &Op::sigmoid());
426            }
427        }
428    }
429
430    #[test]
431    fn test_graph_builder_with_no_any() {
432        let graph = Graph::directed(3, 3, Op::add());
433
434        assert_eq!(graph.len(), 6);
435        assert!(graph.is_valid());
436    }
437
438    #[test]
439    fn test_graph_builder_weighted() {
440        let store = vec![
441            (NodeType::Input, vec![Op::var(0), Op::var(1), Op::var(2)]),
442            (NodeType::Output, vec![Op::sigmoid()]),
443            (NodeType::Edge, vec![Op::weight_with(1.0)]),
444        ];
445
446        let graph = Graph::weighted_directed(3, 3, store);
447
448        assert_eq!(graph.len(), 15);
449        assert!(graph.is_valid());
450
451        for node in graph.iter() {
452            if node.node_type() == NodeType::Input {
453                assert_eq!(node.arity(), Arity::Zero);
454                assert_eq!(node.incoming().iter().count(), 0);
455                assert_eq!(node.outgoing().iter().count(), 3);
456            } else if node.node_type() == NodeType::Edge {
457                assert_eq!(node.arity(), Arity::Exact(1));
458                assert_eq!(node.incoming().iter().count(), 1);
459                assert_eq!(node.outgoing().iter().count(), 1);
460                assert_eq!(node.value(), &Op::weight_with(1.0));
461            } else if node.node_type() == NodeType::Output {
462                assert_eq!(node.arity(), Arity::Any);
463                assert_eq!(node.incoming().iter().count(), 3);
464                assert_eq!(node.outgoing().iter().count(), 0);
465                assert_eq!(node.value(), &Op::sigmoid());
466            }
467        }
468    }
469
470    #[test]
471    fn test_graph_builder_weighted_recurrent() {
472        let store = node_store![
473            Input => vec![Op::var(0), Op::var(1), Op::var(2)],
474            Output => vec![Op::sigmoid()],
475            Edge => vec![Op::weight_with(1.0)]
476        ];
477
478        let graph = Graph::weighted_recurrent(3, 3, store);
479
480        assert_eq!(graph.len(), 18);
481        assert!(graph.is_valid());
482
483        for node in graph.iter() {
484            if node.node_type() == NodeType::Input {
485                assert_eq!(node.arity(), Arity::Zero);
486                assert_eq!(node.incoming().iter().count(), 0);
487                assert_eq!(node.outgoing().iter().count(), 1);
488            } else if node.node_type() == NodeType::Edge {
489                assert_eq!(node.arity(), Arity::Exact(1));
490                assert_eq!(node.incoming().iter().count(), 1);
491                assert_eq!(node.outgoing().iter().count(), 1);
492                assert_eq!(node.value(), &Op::weight_with(1.0));
493            } else if node.node_type() == NodeType::Output {
494                assert_eq!(node.arity(), Arity::Any);
495                assert_eq!(node.incoming().iter().count(), 3);
496                assert_eq!(node.outgoing().iter().count(), 0);
497                assert_eq!(node.value(), &Op::sigmoid());
498            } else if node.node_type() == NodeType::Vertex {
499                assert_eq!(node.arity(), Arity::Any);
500                assert!(node.is_recurrent());
501                assert_eq!(node.value(), &Op::sigmoid());
502            }
503        }
504    }
505
506    #[test]
507    fn test_graph_builder_lstm() {
508        let store = node_store![
509            Input => vec![Op::var(0)],
510            Output => vec![Op::sigmoid()],
511            Vertex => vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()],
512            Edge => vec![Op::weight_with(1.0)]
513        ];
514
515        let graph = Graph::lstm(1, 1, store);
516        assert_eq!(graph.len(), 8);
517        assert!(graph.is_valid());
518
519        for node in graph.iter() {
520            if node.node_type() == NodeType::Input {
521                assert_eq!(node.arity(), Arity::Zero);
522                assert_eq!(node.incoming().iter().count(), 0);
523                assert_eq!(node.outgoing().iter().count(), 4);
524            } else if node.node_type() == NodeType::Output {
525                assert_eq!(node.arity(), Arity::Any);
526                assert_eq!(node.incoming().iter().count(), 1);
527                assert_eq!(node.outgoing().iter().count(), 0);
528                assert_eq!(node.value(), &Op::sigmoid());
529            } else if node.node_type() == NodeType::Vertex {
530                assert_eq!(node.arity(), Arity::Any);
531                assert!(
532                    vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()].contains(&node.value())
533                );
534            } else if node.node_type() == NodeType::Edge {
535                assert_eq!(node.arity(), Arity::Exact(1));
536                assert_eq!(node.incoming().iter().count(), 1);
537                assert_eq!(node.outgoing().iter().count(), 1);
538                assert_eq!(node.value(), &Op::weight_with(1.0));
539            }
540        }
541    }
542
543    #[test]
544    fn test_graph_builder_gru() {
545        let store = node_store![
546            Input => vec![Op::var(0)],
547            Output => vec![Op::sigmoid()],
548            Vertex => vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()],
549            Edge => vec![Op::weight_with(1.0)]
550        ];
551
552        let graph = Graph::gru(1, 1, store);
553
554        assert_eq!(graph.len(), 8);
555        assert!(graph.is_valid());
556
557        for node in graph.iter() {
558            if node.node_type() == NodeType::Input {
559                assert_eq!(node.arity(), Arity::Zero);
560                assert_eq!(node.incoming().iter().count(), 0);
561                assert_eq!(node.outgoing().iter().count(), 3);
562            } else if node.node_type() == NodeType::Output {
563                assert_eq!(node.arity(), Arity::Any);
564                assert_eq!(node.incoming().iter().count(), 1);
565                assert_eq!(node.outgoing().iter().count(), 0);
566                assert_eq!(node.value(), &Op::sigmoid());
567            } else if node.node_type() == NodeType::Vertex {
568                assert_eq!(node.arity(), Arity::Any);
569                assert!(
570                    vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()].contains(&node.value())
571                );
572            } else if node.node_type() == NodeType::Edge {
573                assert_eq!(node.arity(), Arity::Exact(1));
574                assert_eq!(node.incoming().iter().count(), 1);
575                assert_eq!(node.outgoing().iter().count(), 1);
576                assert_eq!(node.value(), &Op::weight_with(1.0));
577            }
578        }
579    }
580}
radiate_gp/collections/graphs/builder.rs

radiate_gp/collections/graphs/
builder.rs
