use super::{DataSet, FeosError, Loss};
use feos_core::Residual;
use ndarray::{Array1, ArrayView1, Axis, arr1, concatenate};
use std::fmt;
use std::fmt::Display;
use std::fmt::Write;
pub struct Estimator<E: Residual> {
data: Vec<Arc<dyn DataSet<E>>>,
weights: Vec<f64>,
losses: Vec<Loss>,
}
impl<E: Residual> Estimator<E> {
pub fn new(data: Vec<Arc<dyn DataSet<E>>>, weights: Vec<f64>, losses: Vec<Loss>) -> Self {
Self {
data,
weights,
losses,
}
}
pub fn add_data(&mut self, data: &Arc<dyn DataSet<E>>, weight: f64, loss: Loss) {
self.data.push(data.clone());
self.weights.push(weight);
self.losses.push(loss);
}
pub fn cost(&self, eos: &Arc<E>) -> Result<Array1<f64>, FeosError> {
let w = arr1(&self.weights) / self.weights.iter().sum::<f64>();
let predictions = self
.data
.iter()
.enumerate()
.map(|(i, d)| Ok(d.cost(eos, self.losses[i])? * w[i]))
.collect::<Result<Vec<_>, FeosError>>()?;
let aview: Vec<ArrayView1<f64>> = predictions.iter().map(|pi| pi.view()).collect();
Ok(concatenate(Axis(0), &aview)?)
}
pub fn predict(&self, eos: &Arc<E>) -> Result<Vec<Array1<f64>>, FeosError> {
self.data.iter().map(|d| d.predict(eos)).collect()
}
pub fn relative_difference(&self, eos: &Arc<E>) -> Result<Vec<Array1<f64>>, FeosError> {
self.data
.iter()
.map(|d| d.relative_difference(eos))
.collect()
}
pub fn mean_absolute_relative_difference(
&self,
eos: &Arc<E>,
) -> Result<Array1<f64>, FeosError> {
self.data
.iter()
.map(|d| d.mean_absolute_relative_difference(eos))
.collect()
}
pub fn datasets(&self) -> Vec<Arc<dyn DataSet<E>>> {
self.data.to_vec()
}
pub fn _repr_markdownn_(&self) -> String {
let mut f = String::new();
write!(f, "| target | input | datapoints |\n|:-|:-|:-|").unwrap();
for d in self.data.iter() {
write!(
f,
"\n|{}|{}|{}|",
d.target_str(),
d.input_str().join(", "),
d.datapoints()
)
.unwrap();
}
f
}
}
impl<E: Residual> Display for Estimator<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for d in self.data.iter() {
writeln!(f, "{d}")?;
}
Ok(())
}
}