Skip to main content

radiate_gp/regression/
fitness.rs

1use super::{DataSet, Loss};
2use crate::{
3    Graph, GraphChromosome, GraphEvaluator, Op, Tree, TreeChromosome, eval::EvalIntoMut,
4    ops::OpFloat,
5};
6use radiate_core::{BatchFitnessFunction, Genotype, Score, fitness::FitnessFunction};
7
8#[derive(Clone)]
9pub struct Regression<F: OpFloat> {
10    data_set: DataSet<F>,
11    loss: Loss,
12}
13
14impl<F: OpFloat> Regression<F> {
15    pub fn new(sample_set: impl Into<DataSet<F>>, loss: Loss) -> Self {
16        Regression {
17            data_set: sample_set.into(),
18            loss,
19        }
20    }
21
22    #[inline]
23    fn calc_into_buff_mut<EV>(&self, eval: &mut EV) -> F
24    where
25        EV: EvalIntoMut<[F], [F]>,
26    {
27        let out_len = self.data_set.shape().2;
28        F::with_loss_buffer(|buf| {
29            if buf.len() < out_len {
30                buf.resize(out_len, F::ZERO);
31            }
32
33            self.calc_with_buf(eval, &mut buf[..out_len])
34        })
35    }
36
37    #[inline]
38    fn calc_with_buf<EV>(&self, eval: &mut EV, buf: &mut [F]) -> F
39    where
40        EV: EvalIntoMut<[F], [F]>,
41    {
42        self.loss
43            .calculate(&self.data_set, buf, |x, y| eval.eval_into_mut(x, y))
44    }
45
46    fn with_buff<E, O>(&self, eval_fn: E) -> O
47    where
48        E: FnOnce(&mut [F]) -> O,
49    {
50        let out_len = self.data_set.shape().2;
51        F::with_loss_buffer(|buf| {
52            if buf.len() < out_len {
53                buf.resize(out_len, F::ZERO);
54            }
55
56            eval_fn(&mut buf[..out_len])
57        })
58    }
59}
60
61impl<'a, F> FitnessFunction<&'a Genotype<GraphChromosome<Op<F>>>, F> for Regression<F>
62where
63    F: OpFloat + Into<Score>,
64{
65    #[inline]
66    fn evaluate(&self, input: &'a Genotype<GraphChromosome<Op<F>>>) -> F {
67        let mut evaluator = GraphEvaluator::new(&input[0]);
68        self.calc_into_buff_mut(&mut evaluator)
69    }
70}
71
72impl<F> FitnessFunction<Graph<Op<F>>, F> for Regression<F>
73where
74    F: OpFloat + Into<Score>,
75{
76    #[inline]
77    fn evaluate(&self, input: Graph<Op<F>>) -> F {
78        let mut evaluator = GraphEvaluator::new(&input);
79        self.calc_into_buff_mut(&mut evaluator)
80    }
81}
82
83impl<F> BatchFitnessFunction<Graph<Op<F>>, F> for Regression<F>
84where
85    F: OpFloat + Into<Score>,
86{
87    #[inline]
88    fn evaluate(&self, inputs: Vec<Graph<Op<F>>>) -> Vec<F> {
89        self.with_buff(move |buf| {
90            let mut results = Vec::with_capacity(inputs.len());
91            for input in inputs.iter() {
92                let mut evaluator = GraphEvaluator::new(&input);
93                results.push(self.calc_with_buf(&mut evaluator, buf));
94            }
95
96            results
97        })
98    }
99}
100
101impl<'a, F> BatchFitnessFunction<&'a Genotype<GraphChromosome<Op<F>>>, F> for Regression<F>
102where
103    F: OpFloat + Into<Score>,
104{
105    #[inline]
106    fn evaluate(&self, inputs: Vec<&'a Genotype<GraphChromosome<Op<F>>>>) -> Vec<F> {
107        self.with_buff(move |buf| {
108            let mut results = Vec::with_capacity(inputs.len());
109            for input in inputs.iter() {
110                let mut evaluator = GraphEvaluator::new(&input[0]);
111                results.push(self.calc_with_buf(&mut evaluator, buf));
112            }
113
114            results
115        })
116    }
117}
118
119/// --- Trees ---
120impl<F> FitnessFunction<Tree<Op<F>>, F> for Regression<F>
121where
122    F: OpFloat + Into<Score>,
123{
124    #[inline]
125    fn evaluate(&self, mut input: Tree<Op<F>>) -> F {
126        self.calc_into_buff_mut(&mut input)
127    }
128}
129
130impl<F> FitnessFunction<Vec<Tree<Op<F>>>, F> for Regression<F>
131where
132    F: OpFloat + Into<Score>,
133{
134    #[inline]
135    fn evaluate(&self, mut input: Vec<Tree<Op<F>>>) -> F {
136        self.calc_into_buff_mut(&mut input)
137    }
138}
139
140impl<F> BatchFitnessFunction<Tree<Op<F>>, F> for Regression<F>
141where
142    F: OpFloat + Into<Score>,
143{
144    #[inline]
145    fn evaluate(&self, mut inputs: Vec<Tree<Op<F>>>) -> Vec<F> {
146        self.with_buff(move |buf| {
147            let mut results = Vec::with_capacity(inputs.len());
148            for input in inputs.iter_mut() {
149                results.push(self.calc_with_buf(input, buf));
150            }
151
152            results
153        })
154    }
155}
156
157impl<F> BatchFitnessFunction<Vec<Tree<Op<F>>>, F> for Regression<F>
158where
159    F: OpFloat + Into<Score>,
160{
161    #[inline]
162    fn evaluate(&self, mut inputs: Vec<Vec<Tree<Op<F>>>>) -> Vec<F> {
163        self.with_buff(move |buf| {
164            let mut results = Vec::with_capacity(inputs.len());
165            for input in inputs.iter_mut() {
166                results.push(self.calc_with_buf(input, buf));
167            }
168
169            results
170        })
171    }
172}
173
174impl<'a, F> FitnessFunction<&'a Genotype<TreeChromosome<Op<F>>>, F> for Regression<F>
175where
176    F: OpFloat + Into<Score>,
177{
178    #[inline]
179    fn evaluate(&self, input: &'a Genotype<TreeChromosome<Op<F>>>) -> F {
180        let roots = input.iter().map(|c| c.root()).collect::<Vec<_>>();
181        self.calc_into_buff_mut(&mut roots.as_slice())
182    }
183}
184
185impl<'a, F> BatchFitnessFunction<&'a Genotype<TreeChromosome<Op<F>>>, F> for Regression<F>
186where
187    F: OpFloat + Into<Score>,
188{
189    #[inline]
190    fn evaluate(&self, inputs: Vec<&'a Genotype<TreeChromosome<Op<F>>>>) -> Vec<F> {
191        self.with_buff(move |buf| {
192            let mut results = Vec::with_capacity(inputs.len());
193            for input in inputs.iter() {
194                let roots = input.iter().map(|c| c.root()).collect::<Vec<_>>();
195                results.push(self.calc_with_buf(&mut roots.as_slice(), buf));
196            }
197
198            results
199        })
200    }
201}