eryon_surface/model/

model.rs

/*
    Appellation: mlp <module>
    Contrib: @FL03
*/
use super::{FftAttention, SurfaceModelConfig};
use cnc::prelude::{ModelFeatures, ModelParams, Params};
use num_traits::FromPrimitive;

/// A multi-layer perceptron implementation
#[derive(Clone, Debug)]
pub struct SurfaceModel<A = f64> {
    pub(crate) config: SurfaceModelConfig<A>,
    pub(crate) features: ModelFeatures,
    pub(crate) attention: Option<FftAttention<A>>,
    pub(crate) params: ModelParams<A>,
    pub(crate) velocities: ModelParams<A>,
    pub(crate) prev_target_norm: Option<A>,
}

impl<A> SurfaceModel<A> {
    /// create a new instance of the model; all parameters are initialized to their defaults
    pub fn default(config: SurfaceModelConfig<A>, features: ModelFeatures) -> Self
    where
        A: Clone + Default,
    {
        Self {
            attention: None,
            config,
            features,
            params: ModelParams::default(features),
            velocities: ModelParams::default(features),
            prev_target_norm: None,
        }
    }
    /// create a new instance of the model; all parameters are initialized to 1
    pub fn ones(config: SurfaceModelConfig<A>, features: ModelFeatures) -> Self
    where
        A: Clone + num_traits::One,
    {
        Self {
            attention: None,
            config,
            features,
            params: ModelParams::ones(features),
            velocities: ModelParams::ones(features),
            prev_target_norm: None,
        }
    }
    /// create a new instance of the model; all parameters are initialized to 0
    pub fn zeros(config: SurfaceModelConfig<A>, features: ModelFeatures) -> Self
    where
        A: Clone + num_traits::Zero,
    {
        Self {
            attention: None,
            config,
            features,
            params: ModelParams::zeros(features),
            velocities: ModelParams::zeros(features),
            prev_target_norm: None,
        }
    }
    /// returns an immutable reference to the attention mechanism of the model
    pub const fn attention(&self) -> Option<&FftAttention<A>> {
        self.attention.as_ref()
    }
    /// returns a mutable reference to the attention mechanism of the model
    pub const fn attention_mut(&mut self) -> Option<&mut FftAttention<A>> {
        self.attention.as_mut()
    }
    /// returns an immutable reference to the model's configuration
    pub const fn config(&self) -> &SurfaceModelConfig<A> {
        &self.config
    }
    /// returns a mutable reference to the model's configuration
    pub const fn config_mut(&mut self) -> &mut SurfaceModelConfig<A> {
        &mut self.config
    }
    /// returns an immutable reference to the model's features
    pub const fn features(&self) -> ModelFeatures {
        self.features
    }
    /// returns a mutable reference to the model's features
    pub const fn features_mut(&mut self) -> &mut ModelFeatures {
        &mut self.features
    }
    /// returns an immutable reference to the model's parameters
    pub const fn params(&self) -> &ModelParams<A> {
        &self.params
    }
    /// returns a mutable reference to the model's parameters
    pub const fn params_mut(&mut self) -> &mut ModelParams<A> {
        &mut self.params
    }
    /// returns an immutable reference to the model's velocities
    pub const fn velocities(&self) -> &ModelParams<A> {
        &self.velocities
    }
    /// returns a mutable reference to the model's velocities
    pub const fn velocities_mut(&mut self) -> &mut ModelParams<A> {
        &mut self.velocities
    }
    /// returns an immutable reference to the input layer of the model
    pub const fn input(&self) -> &Params<A> {
        self.params().input()
    }
    /// returns a mutable reference to the input layer of the model
    pub const fn input_mut(&mut self) -> &mut Params<A> {
        self.params_mut().input_mut()
    }
    /// returns an immutable reference to the hidden layers of the model
    pub const fn hidden(&self) -> &Vec<Params<A>> {
        self.params().hidden()
    }
    /// returns an immutable reference to the hidden layers of the model as a slice
    pub const fn hidden_mut(&mut self) -> &mut Vec<Params<A>> {
        self.params_mut().hidden_mut()
    }
    /// returns an immutable reference to the hidden layers of the model as a slice
    #[inline]
    pub fn hidden_as_slice(&self) -> &[Params<A>] {
        self.params().hidden_as_slice()
    }
    /// returns an immutable reference to the output layer of the model
    pub const fn output(&self) -> &Params<A> {
        self.params().output()
    }
    /// returns a mutable reference to the output layer of the model
    pub const fn output_mut(&mut self) -> &mut Params<A> {
        self.params_mut().output_mut()
    }
    /// returns an immutable reference to the decay of the model
    pub const fn decay(&self) -> &A {
        self.config().decay()
    }
    /// returns an immutable reference to the learning rate of the model
    pub const fn learning_rate(&self) -> &A {
        self.config().learning_rate()
    }
    /// returns an immutable reference to the momentum of the model
    pub const fn momentum(&self) -> &A {
        self.config().momentum()
    }
    /// set the model's attention mechanism
    #[inline]
    pub fn set_attention(&mut self, attention: Option<FftAttention<A>>) -> &mut Self {
        self.attention = attention;
        self
    }
    /// set the model's configuration
    #[inline]
    pub fn set_config(&mut self, config: SurfaceModelConfig<A>) -> &mut Self {
        *self.config_mut() = config;
        self
    }
    /// set the model's features
    #[inline]
    pub fn set_features(&mut self, features: ModelFeatures) -> &mut Self {
        *self.features_mut() = features;
        self
    }
    /// set the model's parameters
    #[inline]
    pub fn set_params(&mut self, params: ModelParams<A>) -> &mut Self {
        *self.params_mut() = params;
        self
    }
    /// set the input layer of the model
    #[inline]
    pub fn set_input(&mut self, input: Params<A>) -> &mut Self {
        self.params_mut().set_input(input);
        self
    }
    /// set the hidden layers of the model
    #[inline]
    pub fn set_hidden<I>(&mut self, iter: I) -> &mut Self
    where
        I: IntoIterator<Item = Params<A>>,
    {
        self.params_mut().set_hidden(Vec::from_iter(iter));
        self
    }
    /// set the output layer of the model
    #[inline]
    pub fn set_output(&mut self, output: Params<A>) -> &mut Self {
        self.params_mut().set_output(output);
        self
    }
    /// set the decay for the model
    #[inline]
    pub fn set_decay(&mut self, decay: A) -> &mut Self {
        self.config_mut().set_decay(decay);
        self
    }
    /// set the learning rate for the model
    #[inline]
    pub fn set_learning_rate(&mut self, learning_rate: A) -> &mut Self {
        self.config_mut().set_learning_rate(learning_rate);
        self
    }
    /// set the model's momentum
    #[inline]
    pub fn set_momentum(&mut self, momentum: A) -> &mut Self {
        self.config_mut().set_momentum(momentum);
        self
    }
    /// create a new instance of the model with the specified attention mechanism
    #[inline]
    pub fn with_attention(self) -> Self
    where
        A: FromPrimitive,
    {
        Self {
            attention: Some(FftAttention::default()),
            ..self
        }
    }
    /// consumes the current instance and returns another with the specified configuration
    pub fn with_config(self, config: SurfaceModelConfig<A>) -> Self {
        Self { config, ..self }
    }
    /// consumes the current instance and returns another with the given features
    pub fn with_features(self, features: ModelFeatures) -> Self {
        Self { features, ..self }
    }
    /// consumes the current instance and returns another with the specified input layer
    pub fn with_input(self, input: Params<A>) -> Self {
        Self {
            params: self.params.with_input(input),
            ..self
        }
    }
    /// consumes the current instance and returns another with the specified hidden layers
    pub fn with_hidden<I>(self, iter: I) -> Self
    where
        I: IntoIterator<Item = Params<A>>,
    {
        Self {
            params: self.params.with_hidden(iter),
            ..self
        }
    }
    /// consumes the current instance and returns another with the specified output layer
    pub fn with_output(self, output: Params<A>) -> Self {
        Self {
            params: self.params.with_output(output),
            ..self
        }
    }
    /// returns true if the model has an attention mechanism
    pub fn has_attention(&self) -> bool {
        self.attention().is_some()
    }
    /// returns true if the model has no attention mechanism
    pub fn has_no_attention(&self) -> bool {
        self.attention().is_none()
    }
}

impl<A> Default for SurfaceModel<A>
where
    A: Clone + Default + FromPrimitive,
{
    fn default() -> Self {
        let config = SurfaceModelConfig::default();
        let features = ModelFeatures::default();
        Self::default(config, features)
    }
}

impl<A> PartialEq for SurfaceModel<A>
where
    A: PartialEq,
    ModelParams<A>: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.config == other.config
            && self.features == other.features
            && self.params == other.params
            && self.velocities == other.velocities
            && self.attention == other.attention
    }
}