use std::marker::PhantomData;
use serde::{Deserialize, Serialize};
use crate::algorithm::neighbour::{KNNAlgorithm, KNNAlgorithmName};
use crate::api::{Predictor, SupervisedEstimator};
use crate::error::Failed;
use crate::linalg::{row_iter, Matrix};
use crate::math::distance::euclidian::Euclidian;
use crate::math::distance::{Distance, Distances};
use crate::math::num::RealNumber;
use crate::neighbors::KNNWeightFunction;
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct KNNClassifierParameters<T: RealNumber, D: Distance<Vec<T>, T>> {
pub distance: D,
pub algorithm: KNNAlgorithmName,
pub weight: KNNWeightFunction,
pub k: usize,
t: PhantomData<T>,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct KNNClassifier<T: RealNumber, D: Distance<Vec<T>, T>> {
classes: Vec<T>,
y: Vec<usize>,
knn_algorithm: KNNAlgorithm<T, D>,
weight: KNNWeightFunction,
k: usize,
}
impl<T: RealNumber, D: Distance<Vec<T>, T>> KNNClassifierParameters<T, D> {
pub fn with_k(mut self, k: usize) -> Self {
self.k = k;
self
}
pub fn with_distance<DD: Distance<Vec<T>, T>>(
self,
distance: DD,
) -> KNNClassifierParameters<T, DD> {
KNNClassifierParameters {
distance,
algorithm: self.algorithm,
weight: self.weight,
k: self.k,
t: PhantomData,
}
}
pub fn with_algorithm(mut self, algorithm: KNNAlgorithmName) -> Self {
self.algorithm = algorithm;
self
}
pub fn with_weight(mut self, weight: KNNWeightFunction) -> Self {
self.weight = weight;
self
}
}
impl<T: RealNumber> Default for KNNClassifierParameters<T, Euclidian> {
fn default() -> Self {
KNNClassifierParameters {
distance: Distances::euclidian(),
algorithm: KNNAlgorithmName::CoverTree,
weight: KNNWeightFunction::Uniform,
k: 3,
t: PhantomData,
}
}
}
impl<T: RealNumber, D: Distance<Vec<T>, T>> PartialEq for KNNClassifier<T, D> {
fn eq(&self, other: &Self) -> bool {
if self.classes.len() != other.classes.len()
|| self.k != other.k
|| self.y.len() != other.y.len()
{
false
} else {
for i in 0..self.classes.len() {
if (self.classes[i] - other.classes[i]).abs() > T::epsilon() {
return false;
}
}
for i in 0..self.y.len() {
if self.y[i] != other.y[i] {
return false;
}
}
true
}
}
}
impl<T: RealNumber, M: Matrix<T>, D: Distance<Vec<T>, T>>
SupervisedEstimator<M, M::RowVector, KNNClassifierParameters<T, D>> for KNNClassifier<T, D>
{
fn fit(
x: &M,
y: &M::RowVector,
parameters: KNNClassifierParameters<T, D>,
) -> Result<Self, Failed> {
KNNClassifier::fit(x, y, parameters)
}
}
impl<T: RealNumber, M: Matrix<T>, D: Distance<Vec<T>, T>> Predictor<M, M::RowVector>
for KNNClassifier<T, D>
{
fn predict(&self, x: &M) -> Result<M::RowVector, Failed> {
self.predict(x)
}
}
impl<T: RealNumber, D: Distance<Vec<T>, T>> KNNClassifier<T, D> {
pub fn fit<M: Matrix<T>>(
x: &M,
y: &M::RowVector,
parameters: KNNClassifierParameters<T, D>,
) -> Result<KNNClassifier<T, D>, Failed> {
let y_m = M::from_row_vector(y.clone());
let (_, y_n) = y_m.shape();
let (x_n, _) = x.shape();
let data = row_iter(x).collect();
let mut yi: Vec<usize> = vec![0; y_n];
let classes = y_m.unique();
for (i, yi_i) in yi.iter_mut().enumerate().take(y_n) {
let yc = y_m.get(0, i);
*yi_i = classes.iter().position(|c| yc == *c).unwrap();
}
if x_n != y_n {
return Err(Failed::fit(&format!(
"Size of x should equal size of y; |x|=[{}], |y|=[{}]",
x_n, y_n
)));
}
if parameters.k <= 1 {
return Err(Failed::fit(&format!(
"k should be > 1, k=[{}]",
parameters.k
)));
}
Ok(KNNClassifier {
classes,
y: yi,
k: parameters.k,
knn_algorithm: parameters.algorithm.fit(data, parameters.distance)?,
weight: parameters.weight,
})
}
pub fn predict<M: Matrix<T>>(&self, x: &M) -> Result<M::RowVector, Failed> {
let mut result = M::zeros(1, x.shape().0);
for (i, x) in row_iter(x).enumerate() {
result.set(0, i, self.classes[self.predict_for_row(x)?]);
}
Ok(result.to_row_vector())
}
fn predict_for_row(&self, x: Vec<T>) -> Result<usize, Failed> {
let search_result = self.knn_algorithm.find(&x, self.k)?;
let weights = self
.weight
.calc_weights(search_result.iter().map(|v| v.1).collect());
let w_sum = weights.iter().copied().sum();
let mut c = vec![T::zero(); self.classes.len()];
let mut max_c = T::zero();
let mut max_i = 0;
for (r, w) in search_result.iter().zip(weights.iter()) {
c[self.y[r.0]] += *w / w_sum;
if c[self.y[r.0]] > max_c {
max_c = c[self.y[r.0]];
max_i = self.y[r.0];
}
}
Ok(max_i)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::linalg::naive::dense_matrix::DenseMatrix;
#[test]
fn knn_fit_predict() {
let x =
DenseMatrix::from_2d_array(&[&[1., 2.], &[3., 4.], &[5., 6.], &[7., 8.], &[9., 10.]]);
let y = vec![2., 2., 2., 3., 3.];
let knn = KNNClassifier::fit(&x, &y, Default::default()).unwrap();
let y_hat = knn.predict(&x).unwrap();
assert_eq!(5, Vec::len(&y_hat));
assert_eq!(y.to_vec(), y_hat);
}
#[test]
fn knn_fit_predict_weighted() {
let x = DenseMatrix::from_2d_array(&[&[1.], &[2.], &[3.], &[4.], &[5.]]);
let y = vec![2., 2., 2., 3., 3.];
let knn = KNNClassifier::fit(
&x,
&y,
KNNClassifierParameters::default()
.with_k(5)
.with_algorithm(KNNAlgorithmName::LinearSearch)
.with_weight(KNNWeightFunction::Distance),
)
.unwrap();
let y_hat = knn.predict(&DenseMatrix::from_2d_array(&[&[4.1]])).unwrap();
assert_eq!(vec![3.0], y_hat);
}
#[test]
fn serde() {
let x =
DenseMatrix::from_2d_array(&[&[1., 2.], &[3., 4.], &[5., 6.], &[7., 8.], &[9., 10.]]);
let y = vec![2., 2., 2., 3., 3.];
let knn = KNNClassifier::fit(&x, &y, Default::default()).unwrap();
let deserialized_knn = bincode::deserialize(&bincode::serialize(&knn).unwrap()).unwrap();
assert_eq!(knn, deserialized_knn);
}
}