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<(NodeType, Vec<Op<f32>>)> = 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<(NodeType, Vec<Op<f32>>)> = 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 vertices = builder.vertices(input_size);
99        let output = builder.output(output_size);
100
101        GraphAggregate::new()
102            .one_to_one(&input, &vertices)
103            .cycle(&vertices)
104            .all_to_all(&vertices, &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 [
191            cell_state,
192            hidden_state,
193            forget_gate,
194            input_gate,
195            output_gate,
196            candidate,
197        ] = builder.gates::<6>();
198
199        GraphAggregate::new()
200            .all_to_all(&input, &forget_gate)
201            .all_to_all(&input, &input_gate)
202            .all_to_all(&input, &output_gate)
203            .all_to_all(&input, &candidate)
204            .one_to_one(&hidden_state, &forget_gate)
205            .one_to_one(&hidden_state, &input_gate)
206            .one_to_one(&hidden_state, &output_gate)
207            .one_to_one(&hidden_state, &candidate)
208            .one_to_one(&forget_gate, &cell_state)
209            .one_to_one(&input_gate, &candidate)
210            .one_to_one(&candidate, &cell_state)
211            .one_to_one(&cell_state, &hidden_state)
212            .one_to_one(&output_gate, &hidden_state)
213            .all_to_all(&hidden_state, &output)
214            .build()
215    }
216
217    /// Creates a Gated Recurrent Unit (GRU) graph with the given input and output sizes.
218    /// The graph will have the following structure:
219    /// - Input nodes connected to reset, update, and candidate gates.
220    /// - Hidden state connected to reset, update, and candidate gates.
221    /// - Reset gate connected to hidden state.
222    /// - Update gate connected to blend and gate flip.
223    /// - Candidate connected to blend.
224    /// - Blend connected to hidden state.
225    /// - Gate flip connected to hidden state.
226    /// - Hidden state connected to output nodes.
227    ///
228    /// # Arguments
229    /// * `input_size` - The number of input nodes.
230    /// * `output_size` - The number of output nodes.
231    /// * `store` - The node store.
232    ///
233    /// # Returns
234    /// A new GRU graph.
235    pub fn gru(input_size: usize, output_size: usize, values: impl Into<NodeStore<T>>) -> Graph<T> {
236        let builder = NodeBuilder::new(values);
237
238        let input = builder.input(input_size);
239        let output = builder.output(output_size);
240
241        let [hidden, update, reset, candidate, blend, gate_flip] = builder.gates::<6>();
242
243        GraphAggregate::new()
244            .many_to_one(&input, &reset)
245            .many_to_one(&input, &update)
246            .many_to_one(&input, &candidate)
247            .one_to_one(&hidden, &reset)
248            .one_to_one(&hidden, &update)
249            .one_to_one(&hidden, &candidate)
250            .one_to_one(&update, &blend)
251            .one_to_one(&candidate, &blend)
252            .one_to_one(&reset, &hidden)
253            .one_to_one(&update, &gate_flip)
254            .one_to_one(&hidden, &gate_flip)
255            .one_to_one(&gate_flip, &hidden)
256            .one_to_one(&blend, &hidden)
257            .one_to_many(&hidden, &output)
258            .build()
259    }
260
261    /// Creates a 2D mesh graph with bidirectional connections between neighboring nodes.
262    /// The graph will have the following structure:
263    /// - Input nodes connected to the first row of mesh nodes.
264    /// - Each mesh node connected to its neighbors (up, down, left, right).
265    /// - Last row of mesh nodes connected to output nodes.
266    ///
267    /// # Arguments
268    /// * `width` - The number of nodes in the horizontal dimension.
269    /// * `height` - The number of nodes in the vertical dimension.
270    /// * `values` - The values to initialize the nodes with.
271    ///
272    /// # Returns
273    /// A new 2D mesh graph.
274    pub fn mesh(
275        input_size: usize,
276        output_size: usize,
277        width: usize,
278        height: usize,
279        values: impl Into<NodeStore<T>>,
280    ) -> Graph<T> {
281        let builder = NodeBuilder::new(values);
282
283        let inputs = builder.input(input_size);
284        let outputs = builder.output(output_size);
285        let nodes = (0..width * height)
286            .map(|_| builder.vertex())
287            .collect::<Vec<Vec<GraphNode<T>>>>();
288
289        let mut aggregate = GraphAggregate::new();
290
291        for y in 0..height {
292            for x in 0..width {
293                let index = y * width + x;
294                let current = &nodes[index];
295
296                if x + 1 < width {
297                    let right = &nodes[y * width + (x + 1)];
298                    aggregate = aggregate.one_to_one(current, right);
299                }
300
301                if y + 1 < height {
302                    let down = &nodes[(y + 1) * width + x];
303                    aggregate = aggregate.one_to_one(current, down);
304                }
305            }
306        }
307
308        aggregate
309            .many_to_one(&inputs, &nodes[0])
310            .one_to_many(&nodes[nodes.len() - 1], &outputs)
311            .build()
312    }
313}
314
315/// A simple builder struct for constructing nodes of a certain type. This is pretty much just a
316/// quality of life struct that removes boilerplate code when creating collections of nodes.
317pub struct NodeBuilder<T> {
318    store: NodeStore<T>,
319}
320
321impl<T: Clone + Default> NodeBuilder<T> {
322    pub fn new(store: impl Into<NodeStore<T>>) -> Self {
323        NodeBuilder {
324            store: store.into(),
325        }
326    }
327
328    pub fn input(&self, size: usize) -> Vec<GraphNode<T>> {
329        self.new_nodes(NodeType::Input, size, Arity::Zero)
330    }
331
332    pub fn output(&self, size: usize) -> Vec<GraphNode<T>> {
333        self.new_nodes(NodeType::Output, size, Arity::Any)
334    }
335
336    pub fn edge(&self, size: usize) -> Vec<GraphNode<T>> {
337        self.new_nodes(NodeType::Edge, size, Arity::Exact(1))
338    }
339
340    pub fn vertex(&self) -> Vec<GraphNode<T>> {
341        self.vertices(1)
342    }
343
344    pub fn vertices(&self, size: usize) -> Vec<GraphNode<T>> {
345        self.new_nodes(NodeType::Vertex, size, Arity::Any)
346    }
347
348    pub fn vertices_with_arity(&self, size: usize, arity: Arity) -> Vec<GraphNode<T>> {
349        (0..size)
350            .filter_map(|idx| {
351                self.store
352                    .new_instance((idx, NodeType::Vertex, |a| a == arity))
353            })
354            .collect()
355    }
356
357    pub fn gate(&self) -> Vec<GraphNode<T>> {
358        self.vertices_with_arity(1, Arity::Any)
359    }
360
361    pub fn gates<const N: usize>(&self) -> [Vec<GraphNode<T>>; N] {
362        std::array::from_fn(|_| self.gate())
363    }
364
365    fn new_nodes(
366        &self,
367        node_type: NodeType,
368        size: usize,
369        fallback_arity: Arity,
370    ) -> Vec<GraphNode<T>> {
371        if self.store.contains_type(node_type) {
372            (0..size)
373                .filter_map(|idx| self.store.new_instance((idx, node_type)))
374                .collect()
375        } else {
376            (0..size)
377                .filter_map(|idx| {
378                    self.store
379                        .new_instance((idx, node_type, |arity| arity == fallback_arity))
380                })
381                .collect()
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests {
388    use super::*;
389    use crate::{Node, Op, node_store};
390    use radiate_core::Valid;
391
392    #[test]
393    fn test_graph_builder() {
394        let graph = Graph::directed(3, 3, Op::<f32>::sigmoid());
395
396        assert_eq!(graph.len(), 6);
397
398        for node in graph.iter() {
399            if node.node_type() == NodeType::Input {
400                assert_eq!(node.arity(), Arity::Zero);
401                assert_eq!(node.incoming().iter().count(), 0);
402                assert_eq!(node.outgoing().iter().count(), 3);
403            } else if node.node_type() == NodeType::Output {
404                assert_eq!(node.arity(), Arity::Any);
405                assert_eq!(node.incoming().iter().count(), 3);
406                assert_eq!(node.outgoing().iter().count(), 0);
407                assert_eq!(node.value(), &Op::sigmoid());
408            }
409        }
410    }
411
412    #[test]
413    fn test_graph_builder_recurrent() {
414        let graph = Graph::recurrent(3, 3, Op::<f32>::sigmoid());
415
416        assert_eq!(graph.len(), 9);
417
418        for node in graph.iter() {
419            if node.node_type() == NodeType::Input {
420                assert_eq!(node.arity(), Arity::Zero);
421                assert_eq!(node.incoming().iter().count(), 0);
422                assert_eq!(node.outgoing().iter().count(), 1);
423            } else if node.node_type() == NodeType::Vertex {
424                assert_eq!(node.arity(), Arity::Any);
425                assert!(node.is_recurrent());
426                assert_eq!(node.value(), &Op::sigmoid());
427            } else if node.node_type() == NodeType::Output {
428                assert_eq!(node.arity(), Arity::Any);
429                assert_eq!(node.incoming().iter().count(), 3);
430                assert_eq!(node.outgoing().iter().count(), 0);
431                assert_eq!(node.value(), &Op::sigmoid());
432            }
433        }
434    }
435
436    #[test]
437    fn test_graph_builder_with_no_any() {
438        let graph = Graph::directed(3, 3, Op::<f32>::add());
439
440        assert_eq!(graph.len(), 6);
441        assert!(graph.is_valid());
442    }
443
444    #[test]
445    fn test_graph_builder_weighted() {
446        let store = vec![
447            (NodeType::Input, vec![Op::var(0), Op::var(1), Op::var(2)]),
448            (NodeType::Output, vec![Op::sigmoid()]),
449            (NodeType::Edge, vec![Op::weight_with(1.0)]),
450        ];
451
452        let graph = Graph::weighted_directed(3, 3, store);
453
454        assert_eq!(graph.len(), 15);
455        assert!(graph.is_valid());
456
457        for node in graph.iter() {
458            if node.node_type() == NodeType::Input {
459                assert_eq!(node.arity(), Arity::Zero);
460                assert_eq!(node.incoming().iter().count(), 0);
461                assert_eq!(node.outgoing().iter().count(), 3);
462            } else if node.node_type() == NodeType::Edge {
463                assert_eq!(node.arity(), Arity::Exact(1));
464                assert_eq!(node.incoming().iter().count(), 1);
465                assert_eq!(node.outgoing().iter().count(), 1);
466                assert_eq!(node.value(), &Op::weight_with(1.0));
467            } else if node.node_type() == NodeType::Output {
468                assert_eq!(node.arity(), Arity::Any);
469                assert_eq!(node.incoming().iter().count(), 3);
470                assert_eq!(node.outgoing().iter().count(), 0);
471                assert_eq!(node.value(), &Op::sigmoid());
472            }
473        }
474    }
475
476    #[test]
477    fn test_graph_builder_weighted_recurrent() {
478        let store = node_store![
479            Input => vec![Op::var(0), Op::var(1), Op::var(2)],
480            Output => vec![Op::sigmoid()],
481            Edge => vec![Op::weight_with(1.0)]
482        ];
483
484        let graph = Graph::weighted_recurrent(3, 3, store);
485
486        assert_eq!(graph.len(), 18);
487        assert!(graph.is_valid());
488
489        for node in graph.iter() {
490            if node.node_type() == NodeType::Input {
491                assert_eq!(node.arity(), Arity::Zero);
492                assert_eq!(node.incoming().iter().count(), 0);
493                assert_eq!(node.outgoing().iter().count(), 1);
494            } else if node.node_type() == NodeType::Edge {
495                assert_eq!(node.arity(), Arity::Exact(1));
496                assert_eq!(node.incoming().iter().count(), 1);
497                assert_eq!(node.outgoing().iter().count(), 1);
498                assert_eq!(node.value(), &Op::weight_with(1.0));
499            } else if node.node_type() == NodeType::Output {
500                assert_eq!(node.arity(), Arity::Any);
501                assert_eq!(node.incoming().iter().count(), 3);
502                assert_eq!(node.outgoing().iter().count(), 0);
503                assert_eq!(node.value(), &Op::sigmoid());
504            } else if node.node_type() == NodeType::Vertex {
505                assert_eq!(node.arity(), Arity::Any);
506                assert!(node.is_recurrent());
507                assert_eq!(node.value(), &Op::sigmoid());
508            }
509        }
510    }
511
512    #[test]
513    fn test_graph_builder_lstm() {
514        let store = node_store![
515            Input => vec![Op::var(0)],
516            Output => vec![Op::sigmoid()],
517            Vertex => vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()],
518            Edge => vec![Op::weight_with(1.0)]
519        ];
520
521        let graph = Graph::lstm(1, 1, store);
522        assert_eq!(graph.len(), 8);
523        assert!(graph.is_valid());
524
525        for node in graph.iter() {
526            if node.node_type() == NodeType::Input {
527                assert_eq!(node.arity(), Arity::Zero);
528                assert_eq!(node.incoming().iter().count(), 0);
529                assert_eq!(node.outgoing().iter().count(), 4);
530            } else if node.node_type() == NodeType::Output {
531                assert_eq!(node.arity(), Arity::Any);
532                assert_eq!(node.incoming().iter().count(), 1);
533                assert_eq!(node.outgoing().iter().count(), 0);
534                assert_eq!(node.value(), &Op::sigmoid());
535            } else if node.node_type() == NodeType::Vertex {
536                assert_eq!(node.arity(), Arity::Any);
537                assert!(
538                    vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()].contains(&node.value())
539                );
540            } else if node.node_type() == NodeType::Edge {
541                assert_eq!(node.arity(), Arity::Exact(1));
542                assert_eq!(node.incoming().iter().count(), 1);
543                assert_eq!(node.outgoing().iter().count(), 1);
544                assert_eq!(node.value(), &Op::weight_with(1.0));
545            }
546        }
547    }
548
549    #[test]
550    fn test_graph_builder_gru() {
551        let store = node_store![
552            Input => vec![Op::var(0)],
553            Output => vec![Op::sigmoid()],
554            Vertex => vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()],
555            Edge => vec![Op::weight_with(1.0)]
556        ];
557
558        let graph = Graph::gru(1, 1, store);
559
560        assert_eq!(graph.len(), 8);
561        assert!(graph.is_valid());
562
563        for node in graph.iter() {
564            if node.node_type() == NodeType::Input {
565                assert_eq!(node.arity(), Arity::Zero);
566                assert_eq!(node.incoming().iter().count(), 0);
567                assert_eq!(node.outgoing().iter().count(), 3);
568            } else if node.node_type() == NodeType::Output {
569                assert_eq!(node.arity(), Arity::Any);
570                assert_eq!(node.incoming().iter().count(), 1);
571                assert_eq!(node.outgoing().iter().count(), 0);
572                assert_eq!(node.value(), &Op::sigmoid());
573            } else if node.node_type() == NodeType::Vertex {
574                assert_eq!(node.arity(), Arity::Any);
575                assert!(
576                    vec![Op::sigmoid(), Op::tanh(), Op::mul(), Op::add()].contains(&node.value())
577                );
578            } else if node.node_type() == NodeType::Edge {
579                assert_eq!(node.arity(), Arity::Exact(1));
580                assert_eq!(node.incoming().iter().count(), 1);
581                assert_eq!(node.outgoing().iter().count(), 1);
582                assert_eq!(node.value(), &Op::weight_with(1.0));
583            }
584        }
585    }
586}
radiate_gp/collections/graphs/builder.rs

radiate_gp/collections/graphs/
builder.rs
