use super::{DataSet, Loss};
use crate::{Graph, GraphChromosome, GraphEvaluator, Op, Tree, TreeChromosome, eval::EvalIntoMut};
use radiate_core::{BatchFitnessFunction, Genotype, fitness::FitnessFunction};
use std::cell::RefCell;
thread_local! {
static LOSS_BUFFER: RefCell<Vec<f32>> = RefCell::new(Vec::new());
}
#[derive(Clone)]
pub struct Regression {
data_set: DataSet<f32>,
loss: Loss,
}
impl Regression {
pub fn new(sample_set: impl Into<DataSet<f32>>, loss: Loss) -> Self {
Regression {
data_set: sample_set.into(),
loss,
}
}
#[inline]
fn calc_into_buff_mut<EV>(&self, eval: &mut EV) -> f32
where
EV: EvalIntoMut<[f32], [f32]>,
{
let out_len = self.data_set.shape().2;
LOSS_BUFFER.with(|cell| {
let mut buf = cell.borrow_mut();
if buf.len() < out_len {
buf.resize(out_len, 0.0);
}
self.loss
.calculate(&self.data_set, &mut buf[..out_len], |x, y| {
eval.eval_into_mut(x, y)
})
})
}
}
impl<'a> FitnessFunction<&'a Genotype<GraphChromosome<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, input: &'a Genotype<GraphChromosome<Op<f32>>>) -> f32 {
let mut evaluator = GraphEvaluator::new(&input[0]);
self.calc_into_buff_mut(&mut evaluator)
}
}
impl FitnessFunction<Graph<Op<f32>>, f32> for Regression {
#[inline]
fn evaluate(&self, input: Graph<Op<f32>>) -> f32 {
let mut evaluator = GraphEvaluator::new(&input);
self.calc_into_buff_mut(&mut evaluator)
}
}
impl BatchFitnessFunction<Graph<Op<f32>>, f32> for Regression {
#[inline]
fn evaluate(&self, inputs: Vec<Graph<Op<f32>>>) -> Vec<f32> {
let mut results = Vec::with_capacity(inputs.len());
for input in inputs {
let mut evaluator = GraphEvaluator::new(&input);
results.push(self.calc_into_buff_mut(&mut evaluator));
}
results
}
}
impl<'a> BatchFitnessFunction<&'a Genotype<GraphChromosome<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, inputs: Vec<&'a Genotype<GraphChromosome<Op<f32>>>>) -> Vec<f32> {
let mut results = Vec::with_capacity(inputs.len());
for input in inputs {
let mut evaluator = GraphEvaluator::new(&input[0]);
results.push(self.calc_into_buff_mut(&mut evaluator));
}
results
}
}
impl FitnessFunction<Tree<Op<f32>>, f32> for Regression {
#[inline]
fn evaluate(&self, mut input: Tree<Op<f32>>) -> f32 {
self.calc_into_buff_mut(&mut input)
}
}
impl FitnessFunction<Vec<Tree<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, mut input: Vec<Tree<Op<f32>>>) -> f32 {
self.calc_into_buff_mut(&mut input)
}
}
impl BatchFitnessFunction<Tree<Op<f32>>, f32> for Regression {
#[inline]
fn evaluate(&self, mut inputs: Vec<Tree<Op<f32>>>) -> Vec<f32> {
let mut results = Vec::with_capacity(inputs.len());
for input in inputs.iter_mut() {
results.push(self.calc_into_buff_mut(input));
}
results
}
}
impl BatchFitnessFunction<Vec<Tree<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, mut inputs: Vec<Vec<Tree<Op<f32>>>>) -> Vec<f32> {
let mut results = Vec::with_capacity(inputs.len());
for input in inputs.iter_mut() {
results.push(self.calc_into_buff_mut(input));
}
results
}
}
impl<'a> FitnessFunction<&'a Genotype<TreeChromosome<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, input: &'a Genotype<TreeChromosome<Op<f32>>>) -> f32 {
let roots = input.iter().map(|c| c.root()).collect::<Vec<_>>();
self.calc_into_buff_mut(&mut roots.as_slice())
}
}
impl<'a> BatchFitnessFunction<&'a Genotype<TreeChromosome<Op<f32>>>, f32> for Regression {
#[inline]
fn evaluate(&self, inputs: Vec<&'a Genotype<TreeChromosome<Op<f32>>>>) -> Vec<f32> {
let mut results = Vec::with_capacity(inputs.len());
for input in inputs {
let roots = input.iter().map(|c| c.root()).collect::<Vec<_>>();
results.push(self.calc_into_buff_mut(&mut roots.as_slice()));
}
results
}
}