tensors-rs 0.1.2

//! Owned two-dimensional row-major arrays.

use core::ops::{Index, IndexMut};

use crate::error::{Error, Result};
use crate::numeric::Float;
use crate::rand::SmallRng;
use crate::view2::{ArrayView2, ArrayViewMut2};

/// Owned 2D row-major array.
#[derive(Clone, Debug, PartialEq)]
pub struct Array2<T> {
    data: Vec<T>,
    rows: usize,
    cols: usize,
}

impl<T> Array2<T> {
    /// Build an array from row-major data.
    pub fn from_vec(shape: [usize; 2], data: Vec<T>) -> Result<Self> {
        let expected = shape[0]
            .checked_mul(shape[1])
            .ok_or(Error::DimensionTooLarge)?;
        if data.len() != expected {
            return Err(Error::shape(vec![expected], vec![data.len()]));
        }
        Ok(Self {
            data,
            rows: shape[0],
            cols: shape[1],
        })
    }

    /// Build an array from a function over `(row, col)`.
    pub fn from_fn(shape: [usize; 2], mut f: impl FnMut(usize, usize) -> T) -> Self {
        let len = shape[0] * shape[1];
        let mut data = Vec::with_capacity(len);
        for i in 0..shape[0] {
            for j in 0..shape[1] {
                data.push(f(i, j));
            }
        }
        Self {
            data,
            rows: shape[0],
            cols: shape[1],
        }
    }

    /// Fallibly build an array from a function over `(row, col)`.
    pub fn try_from_fn(shape: [usize; 2], mut f: impl FnMut(usize, usize) -> T) -> Result<Self> {
        let len = checked_len(shape)?;
        let mut data = Vec::new();
        data.try_reserve_exact(len)
            .map_err(|_| Error::AllocationFailed)?;
        for i in 0..shape[0] {
            for j in 0..shape[1] {
                data.push(f(i, j));
            }
        }
        Ok(Self {
            data,
            rows: shape[0],
            cols: shape[1],
        })
    }

    /// Shape as `[rows, cols]`.
    #[inline]
    pub fn shape(&self) -> [usize; 2] {
        [self.rows, self.cols]
    }

    /// Number of rows.
    #[inline]
    pub fn rows(&self) -> usize {
        self.rows
    }

    /// Number of columns.
    #[inline]
    pub fn cols(&self) -> usize {
        self.cols
    }

    /// Row-major strides in elements.
    #[inline]
    pub fn strides(&self) -> [isize; 2] {
        [self.cols as isize, 1]
    }

    /// Distance in elements between consecutive rows.
    #[inline]
    pub fn row_stride(&self) -> isize {
        self.cols as isize
    }

    /// Distance in elements between consecutive columns.
    #[inline]
    pub fn col_stride(&self) -> isize {
        1
    }

    /// Leading dimension for row-major storage.
    #[inline]
    pub fn leading_dimension(&self) -> isize {
        self.cols as isize
    }

    /// Number of elements.
    #[inline]
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Whether the array has zero elements.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Owned arrays are compact row-major contiguous.
    #[inline]
    pub fn is_contiguous(&self) -> bool {
        true
    }

    /// Borrow the row-major backing slice.
    #[inline]
    pub fn as_slice(&self) -> &[T] {
        &self.data
    }

    /// Borrow the row-major backing slice mutably.
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        &mut self.data
    }

    /// Consume into the backing vector.
    pub fn into_vec(self) -> Vec<T> {
        self.data
    }

    /// Immutable strided view.
    pub fn view(&self) -> ArrayView2<'_, T> {
        ArrayView2::from_raw_parts(&self.data, self.shape(), self.strides(), 0)
    }

    /// Mutable strided view.
    pub fn view_mut(&mut self) -> ArrayViewMut2<'_, T> {
        ArrayViewMut2::from_raw_parts(
            &mut self.data,
            [self.rows, self.cols],
            [self.cols as isize, 1],
            0,
        )
    }

    /// Transpose view without copying.
    pub fn transpose_view(&self) -> ArrayView2<'_, T> {
        self.view().transpose()
    }

    /// Get an element reference.
    #[inline]
    pub fn get(&self, row: usize, col: usize) -> Option<&T> {
        (row < self.rows && col < self.cols).then(|| &self.data[row * self.cols + col])
    }

    /// Get a mutable element reference.
    #[inline]
    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut T> {
        (row < self.rows && col < self.cols).then(|| &mut self.data[row * self.cols + col])
    }

    /// Row as a one-row matrix view.
    pub fn row(&self, row: usize) -> Result<ArrayView2<'_, T>> {
        self.view().row(row)
    }

    /// Borrow a row as a contiguous slice.
    pub fn row_slice(&self, row: usize) -> Result<&[T]> {
        if row >= self.rows {
            return Err(Error::IndexOutOfBounds);
        }
        let start = row * self.cols;
        Ok(&self.data[start..start + self.cols])
    }

    /// Mutable row as a one-row matrix view.
    pub fn row_mut(&mut self, row: usize) -> Result<ArrayViewMut2<'_, T>> {
        if row >= self.rows {
            return Err(Error::IndexOutOfBounds);
        }
        Ok(ArrayViewMut2::from_raw_parts(
            &mut self.data,
            [1, self.cols],
            [self.cols as isize, 1],
            (row * self.cols) as isize,
        ))
    }

    /// Borrow a row as a mutable contiguous slice.
    pub fn row_slice_mut(&mut self, row: usize) -> Result<&mut [T]> {
        if row >= self.rows {
            return Err(Error::IndexOutOfBounds);
        }
        let start = row * self.cols;
        Ok(&mut self.data[start..start + self.cols])
    }

    /// Column as an `rows x 1` matrix view.
    pub fn col(&self, col: usize) -> Result<ArrayView2<'_, T>> {
        self.view().col(col)
    }

    /// Mutable column as an `rows x 1` matrix view.
    pub fn col_mut(&mut self, col: usize) -> Result<ArrayViewMut2<'_, T>> {
        if col >= self.cols {
            return Err(Error::IndexOutOfBounds);
        }
        Ok(ArrayViewMut2::from_raw_parts(
            &mut self.data,
            [self.rows, 1],
            [self.cols as isize, 1],
            col as isize,
        ))
    }

    /// Half-open row slice.
    pub fn rows_range(&self, start: usize, end: usize) -> Result<ArrayView2<'_, T>> {
        self.view().rows_range(start, end)
    }

    /// Mutable half-open row slice.
    pub fn rows_range_mut(&mut self, start: usize, end: usize) -> Result<ArrayViewMut2<'_, T>> {
        if start > end || end > self.rows {
            return Err(Error::IndexOutOfBounds);
        }
        Ok(ArrayViewMut2::from_raw_parts(
            &mut self.data,
            [end - start, self.cols],
            [self.cols as isize, 1],
            (start * self.cols) as isize,
        ))
    }

    /// Half-open column slice.
    pub fn cols_range(&self, start: usize, end: usize) -> Result<ArrayView2<'_, T>> {
        self.view().cols_range(start, end)
    }

    /// Mutable half-open column slice.
    pub fn cols_range_mut(&mut self, start: usize, end: usize) -> Result<ArrayViewMut2<'_, T>> {
        if start > end || end > self.cols {
            return Err(Error::IndexOutOfBounds);
        }
        Ok(ArrayViewMut2::from_raw_parts(
            &mut self.data,
            [self.rows, end - start],
            [self.cols as isize, 1],
            start as isize,
        ))
    }

    /// Reshape without changing storage order.
    pub fn reshape(mut self, shape: [usize; 2]) -> Result<Self> {
        let expected = shape[0]
            .checked_mul(shape[1])
            .ok_or(Error::DimensionTooLarge)?;
        if expected != self.data.len() {
            return Err(Error::shape(vec![self.data.len()], vec![expected]));
        }
        self.rows = shape[0];
        self.cols = shape[1];
        Ok(self)
    }
}

impl<T: Clone> Array2<T> {
    /// Fill a new array with `value`.
    pub fn filled(shape: [usize; 2], value: T) -> Self {
        Self {
            data: vec![value; shape[0] * shape[1]],
            rows: shape[0],
            cols: shape[1],
        }
    }

    /// Fallibly fill a new array with `value`.
    pub fn try_filled(shape: [usize; 2], value: T) -> Result<Self> {
        let len = checked_len(shape)?;
        let mut data = Vec::new();
        data.try_reserve_exact(len)
            .map_err(|_| Error::AllocationFailed)?;
        data.resize(len, value);
        Ok(Self {
            data,
            rows: shape[0],
            cols: shape[1],
        })
    }

    /// Clone into compact row-major storage.
    pub fn clone_contiguous(view: ArrayView2<'_, T>) -> Self {
        Self::from_fn(view.shape(), |i, j| view[(i, j)].clone())
    }

    /// Copy this array into compact row-major storage.
    pub fn to_row_major(&self) -> Self {
        self.clone()
    }

    /// Copy this array into a column-major vector.
    pub fn to_col_major_vec(&self) -> Vec<T> {
        self.view().to_col_major_vec()
    }

    /// Copy values from another view with the same shape.
    pub fn copy_from_view(&mut self, other: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] = other[(i, j)].clone();
            }
        }
        Ok(())
    }
}

impl<T: Float> Array2<T> {
    /// Array filled with zeros.
    pub fn zeros(shape: [usize; 2]) -> Self {
        Self::filled(shape, T::zero())
    }

    /// Fallibly allocate an array filled with zeros.
    pub fn try_zeros(shape: [usize; 2]) -> Result<Self> {
        Self::try_filled(shape, T::zero())
    }

    /// Array filled with ones.
    pub fn ones(shape: [usize; 2]) -> Self {
        Self::filled(shape, T::one())
    }

    /// Fallibly allocate an array filled with ones.
    pub fn try_ones(shape: [usize; 2]) -> Result<Self> {
        Self::try_filled(shape, T::one())
    }

    /// Allocate another array with the same shape, filled with zeros.
    pub fn zeros_like(&self) -> Self {
        Self::zeros(self.shape())
    }

    /// Scale in place.
    pub fn scale(&mut self, alpha: T) {
        for value in &mut self.data {
            *value *= alpha;
        }
    }

    /// Return `self * alpha` without modifying `self`.
    pub fn scaled(&self, alpha: T) -> Self {
        Self::from_fn(self.shape(), |i, j| self[(i, j)] * alpha)
    }

    /// Write `self * alpha` into `out` without modifying `self`.
    pub fn scaled_into(&self, alpha: T, mut out: ArrayViewMut2<'_, T>) -> Result<()> {
        if self.shape() != out.shape() {
            return Err(Error::shape(self.shape(), out.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                out[(i, j)] = self[(i, j)] * alpha;
            }
        }
        Ok(())
    }

    /// Return `self + other` without modifying either input.
    pub fn add(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        self.zip_map(other, |left, right| left + right)
    }

    /// Write `self + other` into `out` without modifying either input.
    pub fn add_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        self.zip_map_into(other, out, |left, right| left + right)
    }

    /// Return `self - other` without modifying either input.
    pub fn sub(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        self.zip_map(other, |left, right| left - right)
    }

    /// Write `self - other` into `out` without modifying either input.
    pub fn sub_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        self.zip_map_into(other, out, |left, right| left - right)
    }

    /// Return the elementwise product `self * other` without modifying either input.
    pub fn mul(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        self.zip_map(other, |left, right| left * right)
    }

    /// Write the elementwise product into `out` without modifying either input.
    pub fn mul_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        self.zip_map_into(other, out, |left, right| left * right)
    }

    /// Return the Hadamard product without modifying either input.
    pub fn hadamard(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        self.mul(other)
    }

    /// Write the Hadamard product into `out` without modifying either input.
    pub fn hadamard_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        self.mul_into(other, out)
    }

    /// Return the elementwise quotient `self / other` without modifying either input.
    pub fn div(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        self.zip_map(other, |left, right| left / right)
    }

    /// Write the elementwise quotient into `out` without modifying either input.
    pub fn div_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        self.zip_map_into(other, out, |left, right| left / right)
    }

    /// Return `self + alpha * x` without modifying either input.
    pub fn axpy_result(&self, alpha: T, x: ArrayView2<'_, T>) -> Result<Self> {
        self.zip_map(x, |left, right| left + alpha * right)
    }

    /// Write `self + alpha * x` into `out` without modifying either input.
    pub fn axpy_into(
        &self,
        alpha: T,
        x: ArrayView2<'_, T>,
        out: ArrayViewMut2<'_, T>,
    ) -> Result<()> {
        self.zip_map_into(x, out, |left, right| left + alpha * right)
    }

    /// Return a matrix product `self @ other` without modifying either input.
    pub fn matmul(&self, other: ArrayView2<'_, T>) -> Result<Self> {
        crate::linalg::matmul(self.view(), other)
    }

    /// Write matrix product `self @ other` into `out` without modifying either input.
    pub fn matmul_into(&self, other: ArrayView2<'_, T>, out: ArrayViewMut2<'_, T>) -> Result<()> {
        crate::linalg::gemm(T::one(), self.view(), false, other, false, T::zero(), out)
    }

    /// In-place addition.
    pub fn add_assign_view(&mut self, other: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] += other[(i, j)];
            }
        }
        Ok(())
    }

    /// In-place subtraction.
    pub fn sub_assign_view(&mut self, other: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] -= other[(i, j)];
            }
        }
        Ok(())
    }

    /// Hadamard product in place.
    pub fn mul_assign_view(&mut self, other: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] *= other[(i, j)];
            }
        }
        Ok(())
    }

    /// In-place elementwise division.
    pub fn div_assign_view(&mut self, other: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] /= other[(i, j)];
            }
        }
        Ok(())
    }

    /// Compute `self += alpha * x`.
    pub fn axpy(&mut self, alpha: T, x: ArrayView2<'_, T>) -> Result<()> {
        if self.shape() != x.shape() {
            return Err(Error::shape(self.shape(), x.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] += alpha * x[(i, j)];
            }
        }
        Ok(())
    }

    /// Map values in place.
    pub fn map_inplace(&mut self, mut f: impl FnMut(T) -> T) {
        for value in &mut self.data {
            *value = f(*value);
        }
    }

    /// Zip-map another view into this array in place.
    pub fn zip_map_inplace(
        &mut self,
        other: ArrayView2<'_, T>,
        mut f: impl FnMut(T, T) -> T,
    ) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                self[(i, j)] = f(self[(i, j)], other[(i, j)]);
            }
        }
        Ok(())
    }

    /// Return an elementwise zip-map without modifying either input.
    pub fn zip_map(&self, other: ArrayView2<'_, T>, mut f: impl FnMut(T, T) -> T) -> Result<Self> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        Ok(Self::from_fn(self.shape(), |i, j| {
            f(self[(i, j)], other[(i, j)])
        }))
    }

    /// Write an elementwise zip-map into `out` without modifying either input.
    pub fn zip_map_into(
        &self,
        other: ArrayView2<'_, T>,
        mut out: ArrayViewMut2<'_, T>,
        mut f: impl FnMut(T, T) -> T,
    ) -> Result<()> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        if self.shape() != out.shape() {
            return Err(Error::shape(self.shape(), out.shape()));
        }
        for i in 0..self.rows {
            for j in 0..self.cols {
                out[(i, j)] = f(self[(i, j)], other[(i, j)]);
            }
        }
        Ok(())
    }

    /// Fill with deterministic uniform random values in `[0, 1)`.
    pub fn fill_uniform(&mut self, seed: u64) {
        let mut rng = SmallRng::new(seed);
        for value in &mut self.data {
            *value = rng.uniform();
        }
    }

    /// Fill with deterministic standard-normal random values.
    pub fn fill_randn(&mut self, seed: u64) {
        let mut rng = SmallRng::new(seed);
        for value in &mut self.data {
            *value = rng.normal();
        }
    }

    /// Sum all elements.
    pub fn sum(&self) -> T {
        self.data.iter().copied().sum()
    }

    /// Mean of all elements, or zero for an empty array.
    pub fn mean(&self) -> T {
        if self.is_empty() {
            T::zero()
        } else {
            self.sum() / T::from_f64(self.len() as f64)
        }
    }

    /// Frobenius norm.
    pub fn norm_frobenius(&self) -> T {
        self.data
            .iter()
            .copied()
            .map(|value| value * value)
            .sum::<T>()
            .sqrt()
    }

    /// Maximum absolute element, or zero for an empty array.
    pub fn max_abs(&self) -> T {
        self.data
            .iter()
            .copied()
            .map(T::abs)
            .fold(
                T::zero(),
                |best, value| if value > best { value } else { best },
            )
    }

    /// Dot product of two arrays treated as flattened row-major buffers.
    pub fn dot(&self, other: ArrayView2<'_, T>) -> Result<T> {
        if self.shape() != other.shape() {
            return Err(Error::shape(self.shape(), other.shape()));
        }
        let mut sum = T::zero();
        for i in 0..self.rows {
            for j in 0..self.cols {
                sum += self[(i, j)] * other[(i, j)];
            }
        }
        Ok(sum)
    }
}

fn checked_len(shape: [usize; 2]) -> Result<usize> {
    shape[0]
        .checked_mul(shape[1])
        .ok_or(Error::DimensionTooLarge)
}

impl<T> Index<(usize, usize)> for Array2<T> {
    type Output = T;

    fn index(&self, index: (usize, usize)) -> &Self::Output {
        self.get(index.0, index.1)
            .expect("array index out of bounds")
    }
}

impl<T> IndexMut<(usize, usize)> for Array2<T> {
    fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
        self.get_mut(index.0, index.1)
            .expect("array index out of bounds")
    }
}