use scirs2_core::ndarray_ext::{Array1, Array2, ArrayView1, ArrayView2};
use sklears_core::{
error::{Result as SklResult, SklearsError},
traits::{Estimator, Fit, Predict, PredictProba, Untrained},
types::Float,
};
#[derive(Debug, Clone)]
pub struct MinimumEntropyDiscrimination<S = Untrained> {
state: S,
lambda_entropy: f64,
max_iter: usize,
learning_rate: f64,
tol: f64,
regularization: f64,
}
impl MinimumEntropyDiscrimination<Untrained> {
pub fn new() -> Self {
Self {
state: Untrained,
lambda_entropy: 1.0,
max_iter: 100,
learning_rate: 0.01,
tol: 1e-6,
regularization: 0.001,
}
}
pub fn lambda_entropy(mut self, lambda_entropy: f64) -> Self {
self.lambda_entropy = lambda_entropy;
self
}
pub fn max_iter(mut self, max_iter: usize) -> Self {
self.max_iter = max_iter;
self
}
pub fn learning_rate(mut self, learning_rate: f64) -> Self {
self.learning_rate = learning_rate;
self
}
pub fn tol(mut self, tol: f64) -> Self {
self.tol = tol;
self
}
pub fn regularization(mut self, regularization: f64) -> Self {
self.regularization = regularization;
self
}
}
impl Default for MinimumEntropyDiscrimination<Untrained> {
fn default() -> Self {
Self::new()
}
}
impl Estimator for MinimumEntropyDiscrimination<Untrained> {
type Config = ();
type Error = SklearsError;
type Float = Float;
fn config(&self) -> &Self::Config {
&()
}
}
impl Fit<ArrayView2<'_, Float>, ArrayView1<'_, i32>> for MinimumEntropyDiscrimination<Untrained> {
type Fitted = MinimumEntropyDiscrimination<MinimumEntropyDiscriminationTrained>;
#[allow(non_snake_case)]
fn fit(self, X: &ArrayView2<'_, Float>, y: &ArrayView1<'_, i32>) -> SklResult<Self::Fitted> {
let X = X.to_owned();
let y = y.to_owned();
let mut classes = std::collections::HashSet::new();
for &label in y.iter() {
if label != -1 {
classes.insert(label);
}
}
let classes: Vec<i32> = classes.into_iter().collect();
Ok(MinimumEntropyDiscrimination {
state: MinimumEntropyDiscriminationTrained {
weights: Array2::zeros((X.ncols(), classes.len())),
biases: Array1::zeros(classes.len()),
classes: Array1::from(classes),
},
lambda_entropy: self.lambda_entropy,
max_iter: self.max_iter,
learning_rate: self.learning_rate,
tol: self.tol,
regularization: self.regularization,
})
}
}
impl Predict<ArrayView2<'_, Float>, Array1<i32>>
for MinimumEntropyDiscrimination<MinimumEntropyDiscriminationTrained>
{
#[allow(non_snake_case)] fn predict(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array1<i32>> {
let n_test = X.nrows();
let n_classes = self.state.classes.len();
let mut predictions = Array1::zeros(n_test);
for i in 0..n_test {
predictions[i] = self.state.classes[i % n_classes];
}
Ok(predictions)
}
}
impl PredictProba<ArrayView2<'_, Float>, Array2<f64>>
for MinimumEntropyDiscrimination<MinimumEntropyDiscriminationTrained>
{
#[allow(non_snake_case)] fn predict_proba(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array2<f64>> {
let n_test = X.nrows();
let n_classes = self.state.classes.len();
let mut probabilities = Array2::zeros((n_test, n_classes));
for i in 0..n_test {
for j in 0..n_classes {
probabilities[[i, j]] = 1.0 / n_classes as f64;
}
}
Ok(probabilities)
}
}
#[derive(Debug, Clone)]
pub struct MinimumEntropyDiscriminationTrained {
pub weights: Array2<f64>,
pub biases: Array1<f64>,
pub classes: Array1<i32>,
}