concision-core 0.3.1

/*
    Appellation: layout <module>
    Contrib: @FL03
*/
use super::ModelFeatures;
use crate::models::{ModelFormat, RawModelLayout, RawModelLayoutMut};

/// verify if the input and hidden dimensions are compatible by checking:
///
/// 1. they have the same dimensionality
/// 2. if the the number of dimensions is greater than one, the hidden layer should be square
/// 3. the finaly dimension of the input is equal to one hidden dimension
fn _verify_input_and_hidden_shape<D>(input: D, hidden: D) -> bool
where
    D: ndarray::Dimension,
{
    let lhs = input.as_array_view();
    let rhs = hidden.as_array_view();
    // check that the input and hidden dimensions are compatible
    input.ndim() != hidden.ndim()
        && lhs[input.ndim() - 1] != rhs[0]
        && rhs.iter().all(|&v| v == rhs[0])
}

impl ModelFeatures {
    pub fn from_shape_and_size(shape: &[usize], size: usize) -> Self {
        let input = shape[0];
        let output = *shape.last().unwrap();
        let hidden = if shape.len() > 2 {
            shape[1]
        } else {
            (size - (input * output)) / (input + output)
        };
        let layers = if shape.len() > 2 { shape.len() - 2 } else { 1 };
        Self::new(input, hidden, output, layers)
    }
    /// creates a new instance of [`ModelFeatures`] for a neural network with `n` layers. If
    /// the number of layers is `<=1` then the [`ModelFormat`] is automatically
    /// configured as a _shallow_ neural network.
    pub const fn new(input: usize, hidden: usize, output: usize, layers: usize) -> Self {
        let inner = ModelFormat::new(hidden, layers);
        Self {
            input,
            output,
            inner,
        }
    }
    /// creates a new instance of [`ModelFeatures`] for a deep neural network, using the given
    /// input, hidden, and output features with the given number of hidden layers
    pub const fn deep(input: usize, hidden: usize, output: usize, layers: usize) -> Self {
        Self {
            input,
            output,
            inner: ModelFormat::Deep { hidden, layers },
        }
    }
    /// returns a new instance of [`ModelFeatures`] for a shallow neural network, using the
    /// given features for the input, output, and single hidden layer
    pub const fn shallow(input: usize, hidden: usize, output: usize) -> Self {
        Self {
            input,
            output,
            inner: ModelFormat::Shallow { hidden },
        }
    }
    pub fn from_layout<L>(layout: L) -> Self
    where
        L: RawModelLayout,
    {
        Self {
            input: layout.input(),
            inner: ModelFormat::new(layout.hidden(), layout.depth()),
            output: layout.output(),
        }
    }
    /// returns a copy of the input features for the model
    pub const fn input(&self) -> usize {
        self.input
    }
    /// returns a mutable reference to the input features for the model
    pub const fn input_mut(&mut self) -> &mut usize {
        &mut self.input
    }
    /// returns a copy of the inner format for the model
    pub const fn inner(&self) -> ModelFormat {
        self.inner
    }
    /// returns a mutable reference to the inner format for the model
    pub const fn inner_mut(&mut self) -> &mut ModelFormat {
        &mut self.inner
    }
    /// returns a copy of the hidden features for the model
    pub const fn hidden(&self) -> usize {
        self.inner().hidden()
    }
    /// returns a mutable reference to the hidden features for the model
    pub const fn hidden_mut(&mut self) -> &mut usize {
        self.inner_mut().hidden_mut()
    }
    /// returns a copy of the number of hidden layers for the model
    pub const fn layers(&self) -> usize {
        self.inner().layers()
    }
    /// returns a mutable reference to the number of hidden layers for the model
    pub const fn layers_mut(&mut self) -> &mut usize {
        self.inner_mut().layers_mut()
    }
    /// returns a copy of the output features for the model
    pub const fn output(&self) -> usize {
        self.output
    }
    /// returns a mutable reference to the output features for the model
    pub const fn output_mut(&mut self) -> &mut usize {
        &mut self.output
    }
    #[inline]
    /// sets the input features for the model
    pub fn set_input(&mut self, input: usize) -> &mut Self {
        self.input = input;
        self
    }
    #[inline]
    /// sets the hidden features for the model
    pub fn set_hidden(&mut self, hidden: usize) -> &mut Self {
        self.inner_mut().set_hidden(hidden);
        self
    }
    #[inline]
    /// sets the number of hidden layers for the model
    pub fn set_layers(&mut self, layers: usize) -> &mut Self {
        self.inner_mut().set_layers(layers);
        self
    }
    #[inline]
    /// sets the output features for the model
    pub fn set_output(&mut self, output: usize) -> &mut Self {
        self.output = output;
        self
    }
    /// consumes the current instance and returns a new instance with the given input
    pub fn with_input(self, input: usize) -> Self {
        Self { input, ..self }
    }
    /// consumes the current instance and returns a new instance with the given hidden
    /// features
    pub fn with_hidden(self, hidden: usize) -> Self {
        Self {
            inner: self.inner.with_hidden(hidden),
            ..self
        }
    }
    /// consumes the current instance and returns a new instance with the given number of
    /// hidden layers
    pub fn with_layers(self, layers: usize) -> Self {
        Self {
            inner: self.inner.with_layers(layers),
            ..self
        }
    }
    /// consumes the current instance and returns a new instance with the given output
    /// features
    pub fn with_output(self, output: usize) -> Self {
        Self { output, ..self }
    }
    /// the dimension of the input layer; (input, hidden)
    pub fn dim_input(&self) -> (usize, usize) {
        (self.input(), self.hidden())
    }
    /// the dimension of the hidden layers; (hidden, hidden)
    pub fn dim_hidden(&self) -> (usize, usize) {
        (self.hidden(), self.hidden())
    }
    /// the dimension of the output layer; (hidden, output)
    pub fn dim_output(&self) -> (usize, usize) {
        (self.hidden(), self.output())
    }
    /// the total number of parameters in the model
    pub fn size(&self) -> usize {
        self.size_input() + self.size_hidden() + self.size_output()
    }
    /// the total number of input parameters in the model
    pub fn size_input(&self) -> usize {
        self.input() * self.hidden()
    }
    /// the total number of hidden parameters in the model
    pub fn size_hidden(&self) -> usize {
        self.hidden() * self.hidden() * self.layers()
    }
    /// the total number of output parameters in the model
    pub fn size_output(&self) -> usize {
        self.hidden() * self.output()
    }
}

impl RawModelLayout for ModelFeatures {
    fn input(&self) -> usize {
        self.input()
    }

    fn hidden(&self) -> usize {
        self.hidden()
    }

    fn depth(&self) -> usize {
        self.layers()
    }

    fn output(&self) -> usize {
        self.output()
    }
}

impl RawModelLayoutMut for ModelFeatures {
    fn input_mut(&mut self) -> &mut usize {
        self.input_mut()
    }

    fn hidden_mut(&mut self) -> &mut usize {
        self.hidden_mut()
    }

    fn layers_mut(&mut self) -> &mut usize {
        self.layers_mut()
    }

    fn output_mut(&mut self) -> &mut usize {
        self.output_mut()
    }
}

impl Default for ModelFeatures {
    fn default() -> Self {
        Self {
            input: 16,
            inner: ModelFormat::Deep {
                hidden: 16,
                layers: 1,
            },
            output: 16,
        }
    }
}

impl core::fmt::Display for ModelFeatures {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str(&format!(
            "{{ input: {}, hidden: {}, layers: {}, output: {} }}",
            self.input(),
            self.hidden(),
            self.layers(),
            self.output()
        ))
    }
}