radiate_gp/collections/trees/
builder.rs

1use crate::collections::{Tree, TreeNode};
2use crate::node::Node;
3use crate::{Arity, Factory, NodeStore, NodeType};
4use radiate_core::Valid;
5
6const NUM_CHILDREN_ANY: usize = 2;
7
8impl<T: Clone + Default> Tree<T> {
9    /// Create a tree with the given depth, where each node is a random node from the node store.
10    /// This obeys the rules of the [NodeStore]'s [NodeType]'s arity, and will create a tree
11    /// that is as balanced as possible.
12    ///
13    /// Note that the root node will try to be a [NodeType::Root] if it is available in the
14    /// [NodeStore], otherwise it will be a [NodeType::Vertex]. This allows caller's to specify what
15    /// the root node is if desired, otherwise it will be a random vertex node from the [NodeStore].
16    ///
17    /// # The [NodeStore] must contain at least one [NodeType::Root] or one [NodeType::Vertex]
18    ///
19    /// # Arguments
20    /// * `depth` - The depth of the tree.
21    /// * `nodes` - The node store to use for the tree.
22    ///
23    /// # Returns
24    /// A tree with the given depth, where each node is a random node from the node store.
25    pub fn with_depth(depth: usize, nodes: impl Into<NodeStore<T>>) -> Self {
26        let store = nodes.into();
27
28        let root = if store.contains_type(NodeType::Root) {
29            store.new_instance(NodeType::Root)
30        } else {
31            store.new_instance(NodeType::Vertex)
32        };
33
34        if let Some(mut root_node) = root {
35            if root_node.arity() == Arity::Any {
36                for _ in 0..NUM_CHILDREN_ANY {
37                    if let Some(child) = Self::grow(depth - 1, &store) {
38                        root_node.add_child(child);
39                    }
40                }
41            } else {
42                for _ in 0..*root_node.arity() {
43                    if let Some(child) = Self::grow(depth - 1, &store) {
44                        root_node.add_child(child);
45                    }
46                }
47            }
48
49            Tree::new(root_node)
50        } else {
51            Tree::default()
52        }
53    }
54
55    /// Recursively grow a tree from the given depth, where each node is a random node from the
56    /// node store. If the depth is 0, then a leaf node is returned. Otherwise, a vertex node is
57    /// returned with children that are grown from the given depth.
58    /// This obeys the rules of the [NodeStore]'s [NodeType]'s arity, and will create a tree
59    /// that is as balanced as possible.
60    ///
61    /// # Arguments
62    /// * `current_depth` - The current depth of the tree.
63    /// * `store` - The node store to use for the tree.
64    ///
65    /// # Returns
66    /// A tree node with the given depth, where each node is a random node from the node store.
67    pub(crate) fn grow(current_depth: usize, store: &NodeStore<T>) -> Option<TreeNode<T>> {
68        if current_depth == 0 {
69            return store.new_instance(NodeType::Leaf);
70        }
71
72        let mut parent = store.new_instance(NodeType::Vertex)?;
73        let num_children = match parent.arity() {
74            Arity::Zero => 0,
75            Arity::Exact(n) => n,
76            Arity::Any => NUM_CHILDREN_ANY,
77        };
78
79        for _ in 0..num_children {
80            let child = Self::grow(current_depth - 1, store)?;
81            parent.add_child(child);
82        }
83
84        Some(parent)
85    }
86
87    #[allow(dead_code)]
88    pub(crate) fn repair_node(node: &mut TreeNode<T>, store: &NodeStore<T>) {
89        if node.children().is_none() && node.is_valid() {
90            return;
91        }
92
93        let num_children = match node.arity() {
94            Arity::Zero => 0,
95            Arity::Exact(n) => n,
96            Arity::Any => NUM_CHILDREN_ANY,
97        };
98
99        let current_num_children = node.children().map_or(0, |c| c.len());
100
101        if current_num_children < num_children {
102            for _ in 0..(num_children - current_num_children) {
103                if let Some(leaf) = store.new_instance(NodeType::Leaf) {
104                    node.add_child(leaf);
105                }
106            }
107        } else if current_num_children > num_children {
108            for _ in 0..(current_num_children - num_children) {
109                node.detach(current_num_children - 1);
110            }
111        }
112
113        if let Some(children) = node.children_mut() {
114            for child in children.iter_mut() {
115                Self::repair_node(child, store);
116            }
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::{Op, TreeIterator};
125
126    #[test]
127    fn test_tree_builder_depth_two() {
128        let store = vec![
129            (NodeType::Vertex, vec![Op::add(), Op::sub(), Op::mul()]),
130            (NodeType::Leaf, vec![Op::constant(1.0), Op::constant(2.0)]),
131        ];
132        let tree = Tree::with_depth(2, store);
133
134        assert!(tree.root().is_some());
135        assert_eq!(tree.root().unwrap().children().unwrap().len(), 2);
136        assert_eq!(tree.height(), 2);
137        assert_eq!(tree.size(), 7);
138
139        for node in tree.iter_breadth_first() {
140            if node.arity() == Arity::Any {
141                assert_eq!(node.children().map(|c| c.len()), Some(2));
142            } else if let Arity::Exact(n) = node.arity() {
143                assert_eq!(node.children().map(|c| c.len()), Some(n));
144            } else {
145                assert_eq!(node.children(), None);
146            }
147        }
148    }
149
150    #[test]
151    fn test_tree_builder_depth_three() {
152        // just a quality of life test to make sure the builder is working.
153        // The above test should be good enough, but just for peace of mind.
154        let store = vec![
155            (NodeType::Vertex, vec![Op::add(), Op::sub(), Op::mul()]),
156            (NodeType::Leaf, vec![Op::constant(1.0), Op::constant(2.0)]),
157        ];
158        let tree = Tree::with_depth(3, store);
159
160        assert!(tree.root().is_some());
161        assert_eq!(tree.root().unwrap().children().unwrap().len(), 2);
162        assert_eq!(tree.height(), 3);
163        assert_eq!(tree.size(), 15);
164
165        for node in tree.iter_breadth_first() {
166            if node.arity() == Arity::Any {
167                assert_eq!(node.children().map(|c| c.len()), Some(2));
168            } else if let Arity::Exact(n) = node.arity() {
169                assert_eq!(node.children().map(|c| c.len()), Some(n));
170            } else {
171                assert_eq!(node.children(), None);
172            }
173        }
174    }
175
176    #[test]
177    fn test_vertex_with_any_arity_builds_correct_depth() {
178        let tree = Tree::with_depth(
179            2,
180            vec![
181                (
182                    NodeType::Vertex,
183                    vec![Op::sigmoid(), Op::relu(), Op::tanh()],
184                ),
185                (NodeType::Leaf, vec![Op::constant(1.0), Op::constant(2.0)]),
186            ],
187        );
188
189        assert!(tree.root().is_some());
190        assert_eq!(tree.root().unwrap().children().unwrap().len(), 2);
191        assert_eq!(tree.height(), 2);
192        assert_eq!(tree.size(), 7);
193
194        for node in tree.iter_breadth_first() {
195            if node.arity() == Arity::Any {
196                assert_eq!(node.children().map(|c| c.len()), Some(2));
197            } else if let Arity::Exact(n) = node.arity() {
198                assert_eq!(node.children().map(|c| c.len()), Some(n));
199            } else {
200                assert_eq!(node.children(), None);
201            }
202        }
203    }
204}
radiate_gp/collections/trees/builder.rs

radiate_gp/collections/trees/
builder.rs
