radiate_gp/collections/trees/
codec.rs

1use crate::NodeStore;
2use crate::collections::{Tree, TreeChromosome, TreeNode};
3use radiate_core::{Codec, Genotype};
4use std::sync::Arc;
5
6type Constraint<N> = Arc<dyn Fn(&N) -> bool + Send + Sync>;
7
8#[derive(Clone)]
9pub struct TreeCodec<T: Clone, D = Vec<Tree<T>>> {
10    depth: usize,
11    num_trees: usize,
12    store: Option<NodeStore<T>>,
13    constraint: Option<Constraint<TreeNode<T>>>,
14    template: Option<Tree<T>>,
15    _marker: std::marker::PhantomData<D>,
16}
17
18impl<T: Clone + Default> TreeCodec<T> {
19    pub fn single(depth: usize, store: impl Into<NodeStore<T>>) -> TreeCodec<T, Tree<T>> {
20        TreeCodec {
21            depth,
22            num_trees: 1,
23            store: Some(store.into()),
24            constraint: None,
25            template: None,
26            _marker: std::marker::PhantomData,
27        }
28    }
29
30    pub fn multi_root(
31        depth: usize,
32        num_trees: usize,
33        store: impl Into<NodeStore<T>>,
34    ) -> TreeCodec<T, Vec<Tree<T>>> {
35        TreeCodec {
36            depth,
37            num_trees,
38            store: Some(store.into()),
39            constraint: None,
40            template: None,
41            _marker: std::marker::PhantomData,
42        }
43    }
44}
45
46impl<T: Clone, D> TreeCodec<T, D> {
47    pub fn constraint<F>(mut self, constraint: F) -> Self
48    where
49        F: Fn(&TreeNode<T>) -> bool + Send + Sync + 'static,
50    {
51        self.constraint = Some(Arc::new(constraint));
52        self
53    }
54
55    pub fn with_tree(mut self, template: impl Into<Tree<T>>) -> Self {
56        self.template = Some(template.into());
57        self
58    }
59}
60
61impl<T> Codec<TreeChromosome<T>, Vec<Tree<T>>> for TreeCodec<T, Vec<Tree<T>>>
62where
63    T: Clone + PartialEq + Default,
64{
65    fn encode(&self) -> Genotype<TreeChromosome<T>> {
66        if let Some(store) = &self.store {
67            let new_chromosomes = (0..self.num_trees)
68                .map(|_| match self.template.as_ref() {
69                    Some(template) => template.clone(),
70                    None => Tree::with_depth(self.depth, store),
71                })
72                .filter_map(|tree| tree.take_root().map(|root| vec![root]))
73                .map(|node| TreeChromosome::new(node, Some(store.clone()), self.constraint.clone()))
74                .collect::<Vec<TreeChromosome<T>>>();
75
76            return Genotype::new(new_chromosomes);
77        }
78
79        Genotype::default()
80    }
81
82    fn decode(&self, genotype: &Genotype<TreeChromosome<T>>) -> Vec<Tree<T>> {
83        genotype
84            .iter()
85            .map(|chromosome| Tree::new(chromosome.root().clone()))
86            .collect()
87    }
88}
89
90impl<T> Codec<TreeChromosome<T>, Tree<T>> for TreeCodec<T, Tree<T>>
91where
92    T: Clone + PartialEq + Default,
93{
94    fn encode(&self) -> Genotype<TreeChromosome<T>> {
95        if let Some(store) = &self.store {
96            let tree = match self.template.as_ref() {
97                Some(template) => template.clone(),
98                None => Tree::with_depth(self.depth, store),
99            };
100
101            let new_chromosome = tree
102                .take_root()
103                .map(|root| vec![root])
104                .map(|tree| TreeChromosome::new(tree, Some(store.clone()), self.constraint.clone()))
105                .unwrap_or_else(|| TreeChromosome::new(vec![], None, self.constraint.clone()));
106
107            return Genotype::new(vec![new_chromosome]);
108        }
109
110        Genotype::new(vec![])
111    }
112
113    fn decode(&self, genotype: &Genotype<TreeChromosome<T>>) -> Tree<T> {
114        genotype
115            .iter()
116            .next()
117            .map(|chromosome| Tree::new(chromosome.root().clone()))
118            .unwrap_or_default()
119    }
120}
121
122unsafe impl<T: Send + Clone> Send for TreeCodec<T> {}
123unsafe impl<T: Sync + Clone> Sync for TreeCodec<T> {}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use crate::{NodeType, ops::Op};
129    use radiate_core::codecs::Codec;
130
131    #[test]
132    fn test_tree_codec() {
133        let store = vec![
134            (NodeType::Root, vec![Op::add(), Op::sub()]),
135            (NodeType::Vertex, vec![Op::add(), Op::sub(), Op::mul()]),
136            (NodeType::Leaf, vec![Op::constant(1.0), Op::constant(2.0)]),
137        ];
138        let codec = TreeCodec::single(3, store);
139
140        let genotype = codec.encode();
141        let tree = codec.decode(&genotype);
142
143        assert_eq!(tree.root().map(|root| root.height()), Some(3));
144        assert!(tree.root().is_some());
145    }
146
147    #[test]
148    fn test_tree_codec_multi() {
149        let store = vec![
150            (NodeType::Root, vec![Op::add(), Op::sub()]),
151            (NodeType::Vertex, vec![Op::add(), Op::sub(), Op::mul()]),
152            (NodeType::Leaf, vec![Op::constant(1.0), Op::constant(2.0)]),
153        ];
154        let codec = TreeCodec::multi_root(3, 2, store);
155
156        let genotype = codec.encode();
157        let trees = codec.decode(&genotype);
158
159        assert_eq!(trees.len(), 2);
160        assert_eq!(trees[0].root().map(|root| root.height()), Some(3));
161        assert_eq!(trees[1].root().map(|root| root.height()), Some(3));
162        assert!(trees[0].root().is_some());
163        assert!(trees[1].root().is_some());
164    }
165}
radiate_gp/collections/trees/codec.rs

radiate_gp/collections/trees/
codec.rs
