use ndarray::prelude::*;
use crate::{
datasets::{CatSample, CatTable, Dataset},
models::Labelled,
types::{Labels, Set, States},
};
pub type CatWtdSample = (CatSample, f64);
#[derive(Clone, Debug)]
pub struct CatWtdTable {
dataset: CatTable,
weights: Array1<f64>,
}
impl Labelled for CatWtdTable {
#[inline]
fn labels(&self) -> &Labels {
self.dataset.labels()
}
}
impl CatWtdTable {
pub fn new(dataset: CatTable, weights: Array1<f64>) -> Self {
assert_eq!(
dataset.values().nrows(),
weights.len(),
"The number of weights must be equal to the number of samples."
);
assert!(
weights.iter().all(|&w| w.is_finite()),
"All weights must be finite."
);
Self { dataset, weights }
}
#[inline]
pub const fn states(&self) -> &States {
self.dataset.states()
}
#[inline]
pub const fn shape(&self) -> &Array1<usize> {
self.dataset.shape()
}
#[inline]
pub const fn weights(&self) -> &Array1<f64> {
&self.weights
}
}
impl Dataset for CatWtdTable {
type Values = CatTable;
#[inline]
fn values(&self) -> &Self::Values {
&self.dataset
}
#[inline]
fn sample_size(&self) -> f64 {
self.weights.sum()
}
fn select(&self, x: &Set<usize>) -> Self {
let dataset = self.dataset.select(x);
let weights = self.weights.clone();
Self::new(dataset, weights)
}
}