onnx-export-rs 0.1.1

use ndarray::{Array1, Array2};

use crate::{Error, Result};

/// Neighbor weighting used during k-NN inference.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum KnnWeight {
    /// Every selected neighbor receives equal weight.
    #[default]
    Uniform,
    /// Neighbors are weighted by inverse Euclidean distance.
    Distance,
}

/// Fitted k-nearest-neighbor regression state.
#[derive(Clone, Debug, PartialEq)]
pub struct KnnRegressor {
    /// Training samples shaped `[samples, features]`.
    pub samples: Array2<f64>,
    /// Target associated with every training sample.
    pub targets: Array1<f64>,
    /// Number of neighbors.
    pub k: usize,
    /// Neighbor weighting strategy.
    pub weight: KnnWeight,
}

impl KnnRegressor {
    /// Creates validated k-NN regression state.
    pub fn new(
        samples: Array2<f64>,
        targets: Array1<f64>,
        k: usize,
        weight: KnnWeight,
    ) -> Result<Self> {
        validate(&samples, targets.len(), k)?;
        if targets.iter().any(|value| !value.is_finite()) {
            return Err(Error::InvalidModel("k-NN targets must be finite".into()));
        }
        Ok(Self {
            samples,
            targets,
            k,
            weight,
        })
    }
}

/// Fitted k-nearest-neighbor classification state.
#[derive(Clone, Debug, PartialEq)]
pub struct KnnClassifier {
    /// Training samples shaped `[samples, features]`.
    pub samples: Array2<f64>,
    /// Zero-based class index associated with every sample.
    pub target_indices: Vec<i64>,
    /// External integer label for every class index.
    pub class_labels: Vec<i64>,
    /// Number of neighbors.
    pub k: usize,
    /// Neighbor weighting strategy.
    pub weight: KnnWeight,
}

impl KnnClassifier {
    /// Creates validated k-NN classification state.
    pub fn new(
        samples: Array2<f64>,
        target_indices: Vec<i64>,
        class_labels: Vec<i64>,
        k: usize,
        weight: KnnWeight,
    ) -> Result<Self> {
        validate(&samples, target_indices.len(), k)?;
        if class_labels.is_empty()
            || target_indices
                .iter()
                .any(|&index| index < 0 || index as usize >= class_labels.len())
        {
            return Err(Error::InvalidModel("invalid k-NN class indices".into()));
        }
        Ok(Self {
            samples,
            target_indices,
            class_labels,
            k,
            weight,
        })
    }
}

fn validate(samples: &Array2<f64>, targets: usize, k: usize) -> Result<()> {
    if samples.nrows() == 0
        || samples.ncols() == 0
        || targets != samples.nrows()
        || k == 0
        || k > samples.nrows()
        || samples.iter().any(|value| !value.is_finite())
    {
        return Err(Error::InvalidModel("invalid k-NN fitted state".into()));
    }
    Ok(())
}