sklears-semi-supervised 0.2.0

//! Relation Networks implementation

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,
};

/// Relation Networks for Few-Shot Learning
///
/// Relation Networks learn to predict relation scores between query and support samples.
/// The network consists of an embedding module and a relation module that learns
/// to compare objects and output relation scores.
#[derive(Debug, Clone)]
pub struct RelationNetworks<S = Untrained> {
    state: S,
    embedding_dim: usize,
    relation_dim: usize,
    hidden_layers: Vec<usize>,
    learning_rate: f64,
    n_episodes: usize,
    n_way: usize,
    n_shot: usize,
    n_query: usize,
}

impl RelationNetworks<Untrained> {
    /// Create a new RelationNetworks instance
    pub fn new() -> Self {
        Self {
            state: Untrained,
            embedding_dim: 64,
            relation_dim: 8,
            hidden_layers: vec![64, 64],
            learning_rate: 0.001,
            n_episodes: 100,
            n_way: 5,
            n_shot: 1,
            n_query: 15,
        }
    }

    /// Set the embedding dimensionality
    pub fn embedding_dim(mut self, embedding_dim: usize) -> Self {
        self.embedding_dim = embedding_dim;
        self
    }

    /// Set the relation module dimensionality
    pub fn relation_dim(mut self, relation_dim: usize) -> Self {
        self.relation_dim = relation_dim;
        self
    }

    /// Set the hidden layer dimensions
    pub fn hidden_layers(mut self, hidden_layers: Vec<usize>) -> Self {
        self.hidden_layers = hidden_layers;
        self
    }

    /// Set the learning rate
    pub fn learning_rate(mut self, learning_rate: f64) -> Self {
        self.learning_rate = learning_rate;
        self
    }

    /// Set the number of training episodes
    pub fn n_episodes(mut self, n_episodes: usize) -> Self {
        self.n_episodes = n_episodes;
        self
    }

    /// Set the number of classes per episode (N-way)
    pub fn n_way(mut self, n_way: usize) -> Self {
        self.n_way = n_way;
        self
    }

    /// Set the number of support examples per class (N-shot)
    pub fn n_shot(mut self, n_shot: usize) -> Self {
        self.n_shot = n_shot;
        self
    }

    /// Set the number of query examples per class
    pub fn n_query(mut self, n_query: usize) -> Self {
        self.n_query = n_query;
        self
    }
}

impl Default for RelationNetworks<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for RelationNetworks<Untrained> {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &()
    }
}

impl Fit<ArrayView2<'_, Float>, ArrayView1<'_, i32>> for RelationNetworks<Untrained> {
    type Fitted = RelationNetworks<RelationNetworksTrained>;

    #[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();

        // Get unique classes
        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(RelationNetworks {
            state: RelationNetworksTrained {
                embedding_weights: Array2::zeros((X.ncols(), self.embedding_dim)),
                relation_weights: Array2::zeros((self.embedding_dim * 2, self.relation_dim)),
                classes: Array1::from(classes),
            },
            embedding_dim: self.embedding_dim,
            relation_dim: self.relation_dim,
            hidden_layers: self.hidden_layers,
            learning_rate: self.learning_rate,
            n_episodes: self.n_episodes,
            n_way: self.n_way,
            n_shot: self.n_shot,
            n_query: self.n_query,
        })
    }
}

impl Predict<ArrayView2<'_, Float>, Array1<i32>> for RelationNetworks<RelationNetworksTrained> {
    #[allow(non_snake_case)] // standard ML notation
    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 RelationNetworks<RelationNetworksTrained>
{
    #[allow(non_snake_case)] // standard ML notation
    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)
    }
}

/// Trained state for RelationNetworks
#[derive(Debug, Clone)]
pub struct RelationNetworksTrained {
    /// embedding_weights
    pub embedding_weights: Array2<f64>,
    /// relation_weights
    pub relation_weights: Array2<f64>,
    /// classes
    pub classes: Array1<i32>,
}