Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,410 @@
|
||||
//! K-Nearest Neighbors implementations
|
||||
//!
|
||||
//! KNN is a lazy learning algorithm that makes predictions based on the
|
||||
//! k nearest neighbors in the feature space.
|
||||
|
||||
use crate::error::Result;
|
||||
use rtx_tensor::Tensor;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Distance metrics for KNN
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum Distance {
|
||||
Euclidean,
|
||||
Manhattan,
|
||||
Minkowski { p: f32 },
|
||||
}
|
||||
|
||||
/// K-Nearest Neighbors Classifier
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct KNeighborsClassifier {
|
||||
/// Number of neighbors
|
||||
n_neighbors: usize,
|
||||
/// Distance metric
|
||||
metric: Distance,
|
||||
/// Training data
|
||||
x_train: Option<Tensor>,
|
||||
/// Training labels
|
||||
y_train: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl KNeighborsClassifier {
|
||||
/// Create new KNN classifier
|
||||
pub fn new(n_neighbors: usize) -> Self {
|
||||
Self {
|
||||
n_neighbors,
|
||||
metric: Distance::Euclidean,
|
||||
x_train: None,
|
||||
y_train: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set distance metric
|
||||
pub fn metric(mut self, metric: Distance) -> Self {
|
||||
self.metric = metric;
|
||||
self
|
||||
}
|
||||
|
||||
/// Fit the classifier (store training data)
|
||||
pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if y.shape().ndim() != 1 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"y must be 1-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_samples = x.shape().dims()[0];
|
||||
if y.shape().dims()[0] != n_samples {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X and y must have same number of samples".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if n_samples < self.n_neighbors {
|
||||
return Err(crate::error::MLError::invalid_input(format!(
|
||||
"Number of samples ({}) must be >= n_neighbors ({})",
|
||||
n_samples, self.n_neighbors
|
||||
)));
|
||||
}
|
||||
|
||||
// Store training data (KNN is lazy learning)
|
||||
self.x_train = Some(x.clone());
|
||||
self.y_train = Some(y.clone());
|
||||
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Predict class labels for test data
|
||||
pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let x_train = self.x_train.as_ref().ok_or_else(|| {
|
||||
crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
|
||||
})?;
|
||||
|
||||
let y_train = self.y_train.as_ref().unwrap();
|
||||
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_test_samples = x.shape().dims()[0];
|
||||
let n_features = x.shape().dims()[1];
|
||||
|
||||
if x_train.shape().dims()[1] != n_features {
|
||||
return Err(crate::error::MLError::invalid_input(format!(
|
||||
"X has {} features but model was fitted with {} features",
|
||||
n_features,
|
||||
x_train.shape().dims()[1]
|
||||
)));
|
||||
}
|
||||
|
||||
// Get data as CPU tensors
|
||||
let x_test_data = x.to_cpu()?;
|
||||
let x_train_data = x_train.to_cpu()?;
|
||||
let y_train_data = y_train.to_cpu()?;
|
||||
|
||||
let n_train_samples = x_train.shape().dims()[0];
|
||||
let mut predictions = vec![0.0f32; n_test_samples];
|
||||
|
||||
// For each test sample
|
||||
for i in 0..n_test_samples {
|
||||
// Compute distances to all training samples
|
||||
let mut distances: Vec<(f32, usize)> = Vec::with_capacity(n_train_samples);
|
||||
|
||||
for j in 0..n_train_samples {
|
||||
let distance = self.compute_distance(
|
||||
&x_test_data[i * n_features..(i + 1) * n_features],
|
||||
&x_train_data[j * n_features..(j + 1) * n_features],
|
||||
);
|
||||
distances.push((distance, j));
|
||||
}
|
||||
|
||||
// Sort by distance and take k nearest
|
||||
distances.sort_by(|a, b| a.0.total_cmp(&b.0));
|
||||
let k_nearest: Vec<usize> = distances
|
||||
.iter()
|
||||
.take(self.n_neighbors)
|
||||
.map(|(_, idx)| *idx)
|
||||
.collect();
|
||||
|
||||
// Majority vote among k nearest neighbors
|
||||
let mut class_votes: HashMap<i32, usize> = HashMap::new();
|
||||
for &neighbor_idx in &k_nearest {
|
||||
let class = y_train_data[neighbor_idx] as i32;
|
||||
*class_votes.entry(class).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
// Find class with most votes
|
||||
let predicted_class = class_votes
|
||||
.iter()
|
||||
.max_by_key(|(_, count)| **count)
|
||||
.map_or(0, |(&class, _)| class) as f32;
|
||||
|
||||
predictions[i] = predicted_class;
|
||||
}
|
||||
|
||||
Ok(Tensor::from_data(
|
||||
predictions,
|
||||
vec![n_test_samples],
|
||||
x.device(),
|
||||
)?)
|
||||
}
|
||||
|
||||
/// Compute distance between two points
|
||||
fn compute_distance(&self, point1: &[f32], point2: &[f32]) -> f32 {
|
||||
match self.metric {
|
||||
Distance::Euclidean => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).powi(2))
|
||||
.sum::<f32>()
|
||||
.sqrt(),
|
||||
Distance::Manhattan => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).abs())
|
||||
.sum::<f32>(),
|
||||
Distance::Minkowski { p } => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).abs().powf(p))
|
||||
.sum::<f32>()
|
||||
.powf(1.0 / p),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get parameters
|
||||
pub fn params(&self) -> (usize, Distance) {
|
||||
(self.n_neighbors, self.metric)
|
||||
}
|
||||
|
||||
/// Check if model is fitted
|
||||
pub fn is_fitted(&self) -> bool {
|
||||
self.x_train.is_some() && self.y_train.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for KNeighborsClassifier {
|
||||
fn default() -> Self {
|
||||
Self::new(5)
|
||||
}
|
||||
}
|
||||
|
||||
/// K-Nearest Neighbors Regressor
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct KNeighborsRegressor {
|
||||
/// Number of neighbors
|
||||
n_neighbors: usize,
|
||||
/// Distance metric
|
||||
metric: Distance,
|
||||
/// Weights for neighbors
|
||||
weights: Weights,
|
||||
/// Training data
|
||||
x_train: Option<Tensor>,
|
||||
/// Training targets
|
||||
y_train: Option<Tensor>,
|
||||
}
|
||||
|
||||
/// Weighting schemes for KNN regression
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum Weights {
|
||||
Uniform, // All neighbors have equal weight
|
||||
Distance, // Weight by inverse distance
|
||||
}
|
||||
|
||||
impl KNeighborsRegressor {
|
||||
/// Create new KNN regressor
|
||||
pub fn new(n_neighbors: usize) -> Self {
|
||||
Self {
|
||||
n_neighbors,
|
||||
metric: Distance::Euclidean,
|
||||
weights: Weights::Uniform,
|
||||
x_train: None,
|
||||
y_train: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set distance metric
|
||||
pub fn metric(mut self, metric: Distance) -> Self {
|
||||
self.metric = metric;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set weighting scheme
|
||||
pub fn weights(mut self, weights: Weights) -> Self {
|
||||
self.weights = weights;
|
||||
self
|
||||
}
|
||||
|
||||
/// Fit the regressor (store training data)
|
||||
pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if y.shape().ndim() != 1 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"y must be 1-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_samples = x.shape().dims()[0];
|
||||
if y.shape().dims()[0] != n_samples {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X and y must have same number of samples".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if n_samples < self.n_neighbors {
|
||||
return Err(crate::error::MLError::invalid_input(format!(
|
||||
"Number of samples ({}) must be >= n_neighbors ({})",
|
||||
n_samples, self.n_neighbors
|
||||
)));
|
||||
}
|
||||
|
||||
// Store training data (KNN is lazy learning)
|
||||
self.x_train = Some(x.clone());
|
||||
self.y_train = Some(y.clone());
|
||||
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Predict continuous values for test data
|
||||
pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let x_train = self.x_train.as_ref().ok_or_else(|| {
|
||||
crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
|
||||
})?;
|
||||
|
||||
let y_train = self.y_train.as_ref().unwrap();
|
||||
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_test_samples = x.shape().dims()[0];
|
||||
let n_features = x.shape().dims()[1];
|
||||
|
||||
if x_train.shape().dims()[1] != n_features {
|
||||
return Err(crate::error::MLError::invalid_input(format!(
|
||||
"X has {} features but model was fitted with {} features",
|
||||
n_features,
|
||||
x_train.shape().dims()[1]
|
||||
)));
|
||||
}
|
||||
|
||||
// Get data as CPU tensors
|
||||
let x_test_data = x.to_cpu()?;
|
||||
let x_train_data = x_train.to_cpu()?;
|
||||
let y_train_data = y_train.to_cpu()?;
|
||||
|
||||
let n_train_samples = x_train.shape().dims()[0];
|
||||
let mut predictions = vec![0.0f32; n_test_samples];
|
||||
|
||||
// For each test sample
|
||||
for i in 0..n_test_samples {
|
||||
// Compute distances to all training samples
|
||||
let mut distances: Vec<(f32, usize)> = Vec::with_capacity(n_train_samples);
|
||||
|
||||
for j in 0..n_train_samples {
|
||||
let distance = self.compute_distance(
|
||||
&x_test_data[i * n_features..(i + 1) * n_features],
|
||||
&x_train_data[j * n_features..(j + 1) * n_features],
|
||||
);
|
||||
distances.push((distance, j));
|
||||
}
|
||||
|
||||
// Sort by distance and take k nearest
|
||||
distances.sort_by(|a, b| a.0.total_cmp(&b.0));
|
||||
let k_nearest: Vec<(f32, usize)> =
|
||||
distances.into_iter().take(self.n_neighbors).collect();
|
||||
|
||||
// Compute weighted average
|
||||
let prediction = match self.weights {
|
||||
Weights::Uniform => {
|
||||
// Simple average
|
||||
let sum: f32 = k_nearest.iter().map(|(_, idx)| y_train_data[*idx]).sum();
|
||||
sum / self.n_neighbors as f32
|
||||
}
|
||||
Weights::Distance => {
|
||||
// Weighted by inverse distance
|
||||
let mut weighted_sum = 0.0f32;
|
||||
let mut weight_sum = 0.0f32;
|
||||
|
||||
for (distance, idx) in k_nearest {
|
||||
let weight = if distance > 0.0 { 1.0 / distance } else { 1e6 }; // Handle exact matches
|
||||
weighted_sum += weight * y_train_data[idx];
|
||||
weight_sum += weight;
|
||||
}
|
||||
|
||||
if weight_sum > 0.0 {
|
||||
weighted_sum / weight_sum
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
predictions[i] = prediction;
|
||||
}
|
||||
|
||||
Ok(Tensor::from_data(
|
||||
predictions,
|
||||
vec![n_test_samples],
|
||||
x.device(),
|
||||
)?)
|
||||
}
|
||||
|
||||
/// Compute distance between two points
|
||||
fn compute_distance(&self, point1: &[f32], point2: &[f32]) -> f32 {
|
||||
match self.metric {
|
||||
Distance::Euclidean => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).powi(2))
|
||||
.sum::<f32>()
|
||||
.sqrt(),
|
||||
Distance::Manhattan => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).abs())
|
||||
.sum::<f32>(),
|
||||
Distance::Minkowski { p } => point1
|
||||
.iter()
|
||||
.zip(point2.iter())
|
||||
.map(|(&x1, &x2)| (x1 - x2).abs().powf(p))
|
||||
.sum::<f32>()
|
||||
.powf(1.0 / p),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get parameters
|
||||
pub fn params(&self) -> (usize, Distance, Weights) {
|
||||
(self.n_neighbors, self.metric, self.weights)
|
||||
}
|
||||
|
||||
/// Check if model is fitted
|
||||
pub fn is_fitted(&self) -> bool {
|
||||
self.x_train.is_some() && self.y_train.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for KNeighborsRegressor {
|
||||
fn default() -> Self {
|
||||
Self::new(5)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user