Implement Tensor struct

Cargo.toml

@@ -4,8 +4,9 @@ version = "0.1.0"
 edition = "2021"
 description = "A tensor library for scientific computing in Rust"
 license = "MIT"
-license-file = "LICENSE"
 homepage = "https://github.com/Rust-Scientific-Computing/feotensor"
 repository = "https://github.com/Rust-Scientific-Computing/feotensor"
 
 [dependencies]
+itertools = "0.13.0"
+num = "0.4.3"

src/axes.rs

@@ -0,0 +1 @@
+pub type Axes = Vec<usize>;

src/coordinate.rs

@@ -0,0 +1,131 @@
+use std::fmt;
+use std::ops::{Index, IndexMut};
+
+use crate::error::ShapeError;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Coordinate {
+    indices: Vec<usize>,
+}
+
+impl Coordinate {
+    pub fn new(indices: Vec<usize>) -> Result<Self, ShapeError> {
+        if indices.is_empty() {
+            return Err(ShapeError::new("Coordinate cannot be empty"));
+        }
+        Ok(Self { indices })
+    }
+
+    pub fn order(&self) -> usize {
+        self.indices.len()
+    }
+
+    pub fn iter(&self) -> std::slice::Iter<'_, usize> {
+        self.indices.iter()
+    }
+
+    pub fn insert(&self, index: usize, axis: usize) -> Self {
+        let mut new_indices = self.indices.clone();
+        new_indices.insert(index, axis);
+        Self {
+            indices: new_indices,
+        }
+    }
+}
+
+impl Index<usize> for Coordinate {
+    type Output = usize;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.indices[index]
+    }
+}
+
+impl IndexMut<usize> for Coordinate {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        &mut self.indices[index]
+    }
+}
+
+impl fmt::Display for Coordinate {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use itertools::Itertools;
+        let idxs = self.indices.iter().map(|&x| format!("{}", x)).join(", ");
+        write!(f, "({})", idxs)
+    }
+}
+
+#[macro_export]
+macro_rules! coord {
+    ($($index:expr),*) => {
+        {
+            use $crate::coordinate::Coordinate;
+            Coordinate::new(vec![$($index),*])
+        }
+    };
+
+    ($index:expr; $count:expr) => {
+        {
+            use $crate::coordinate::Coordinate;
+            Coordinate::new(vec![$index; $count])
+        }
+    };
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_order() {
+        let coord = coord![1, 2, 3].unwrap();
+        assert_eq!(coord.order(), 3);
+    }
+
+    #[test]
+    fn test_iter() {
+        let coord = coord![1, 2, 3].unwrap();
+        let mut iter = coord.iter();
+        assert_eq!(iter.next(), Some(&1));
+        assert_eq!(iter.next(), Some(&2));
+        assert_eq!(iter.next(), Some(&3));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_insert() {
+        let coord = coord![1, 2, 3].unwrap();
+        let new_coord = coord.insert(1, 4);
+        assert_eq!(new_coord, coord![1, 4, 2, 3].unwrap());
+    }
+
+    #[test]
+    fn test_index() {
+        let coord = coord![1, 2, 3].unwrap();
+        assert_eq!(coord[0], 1);
+        assert_eq!(coord[1], 2);
+        assert_eq!(coord[2], 3);
+    }
+
+    #[test]
+    fn test_index_mut() {
+        let mut coord = coord![1, 2, 3].unwrap();
+        coord[1] = 4;
+        assert_eq!(coord[1], 4);
+    }
+
+    #[test]
+    fn test_display() {
+        let coord = coord![1, 2, 3].unwrap();
+        assert_eq!(format!("{}", coord), "(1, 2, 3)");
+    }
+
+    #[test]
+    fn test_coord_macro() {
+        let coord = coord![1, 2, 3].unwrap();
+        assert_eq!(coord, Coordinate::new(vec![1, 2, 3]).unwrap());
+
+        let coord_repeated = coord![1; 3].unwrap();
+        assert_eq!(coord_repeated, Coordinate::new(vec![1, 1, 1]).unwrap());
+    }
+}

src/error.rs

@@ -0,0 +1,20 @@
+use std::fmt;
+
+#[derive(Debug, Clone)]
+pub struct ShapeError {
+    reason: String,
+}
+
+impl ShapeError {
+    pub fn new(reason: &str) -> Self {
+        ShapeError {
+            reason: reason.to_string(),
+        }
+    }
+}
+
+impl fmt::Display for ShapeError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "ShapeError: {}", self.reason)
+    }
+}

src/iter.rs

@@ -0,0 +1,88 @@
+use crate::coord;
+use crate::coordinate::Coordinate;
+use crate::shape::Shape;
+use std::cmp::max;
+
+pub struct IndexIterator {
+    shape: Shape,
+    current: Coordinate,
+    done: bool,
+}
+
+impl IndexIterator {
+    pub fn new(shape: &Shape) -> Self {
+        // (shape.order() == 0) => `next` returns None before `current` is used
+        let current = coord![0; max(shape.order(), 1)].unwrap();
+        IndexIterator {
+            shape: shape.clone(),
+            current,
+            done: false,
+        }
+    }
+}
+
+impl Iterator for IndexIterator {
+    type Item = Coordinate;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done || self.shape.order() == 0 {
+            return None;
+        }
+
+        let result = self.current.clone();
+
+        for i in (0..self.shape.order()).rev() {
+            if self.current[i] + 1 < self.shape[i] {
+                self.current[i] += 1;
+                break;
+            } else {
+                self.current[i] = 0;
+                if i == 0 {
+                    self.done = true;
+                }
+            }
+        }
+
+        Some(result)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::shape;
+
+    #[test]
+    fn test_index_iterator() {
+        let shape = shape![2, 3].unwrap();
+        let mut iter = IndexIterator::new(&shape);
+
+        assert_eq!(iter.next(), Some(coord![0, 0].unwrap()));
+        assert_eq!(iter.next(), Some(coord![0, 1].unwrap()));
+        assert_eq!(iter.next(), Some(coord![0, 2].unwrap()));
+        assert_eq!(iter.next(), Some(coord![1, 0].unwrap()));
+        assert_eq!(iter.next(), Some(coord![1, 1].unwrap()));
+        assert_eq!(iter.next(), Some(coord![1, 2].unwrap()));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_index_iterator_single_dimension() {
+        let shape = shape![4].unwrap();
+        let mut iter = IndexIterator::new(&shape);
+
+        assert_eq!(iter.next(), Some(coord![0].unwrap()));
+        assert_eq!(iter.next(), Some(coord![1].unwrap()));
+        assert_eq!(iter.next(), Some(coord![2].unwrap()));
+        assert_eq!(iter.next(), Some(coord![3].unwrap()));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_index_iterator_empty_tensor() {
+        let shape = shape![].unwrap();
+        let mut iter = IndexIterator::new(&shape);
+
+        assert_eq!(iter.next(), None);
+    }
+}

src/lib.rs

@@ -1,14 +1,9 @@
-pub fn add(left: usize, right: usize) -> usize {
-    left + right
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn it_works() {
-        let result = add(2, 2);
-        assert_eq!(result, 4);
-    }
-}
+pub mod axes;
+pub mod coordinate;
+pub mod error;
+pub mod iter;
+pub mod matrix;
+pub mod shape;
+pub mod storage;
+pub mod tensor;
+pub mod vector;

src/shape.rs

@@ -0,0 +1,114 @@
+use std::fmt;
+use std::ops::Index;
+
+use crate::error::ShapeError;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Shape {
+    dims: Vec<usize>,
+}
+
+impl Shape {
+    pub fn new(dims: Vec<usize>) -> Result<Shape, ShapeError> {
+        if dims.iter().any(|&x| x == 0) {
+            return Err(ShapeError::new("Dimension cannot be zero"));
+        }
+        Ok(Shape { dims })
+    }
+
+    pub fn size(&self) -> usize {
+        self.dims.iter().product()
+    }
+
+    pub fn order(&self) -> usize {
+        self.dims.len()
+    }
+
+    pub fn stack(&self, rhs: &Shape) -> Shape {
+        let mut new_dims = self.dims.clone();
+        new_dims.extend(rhs.dims.iter());
+        Shape { dims: new_dims }
+    }
+}
+
+impl Index<usize> for Shape {
+    type Output = usize;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.dims[index]
+    }
+}
+
+impl Index<std::ops::RangeFrom<usize>> for Shape {
+    type Output = [usize];
+
+    fn index(&self, index: std::ops::RangeFrom<usize>) -> &Self::Output {
+        &self.dims[index]
+    }
+}
+
+impl fmt::Display for Shape {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use itertools::Itertools;
+        let dims = self.dims.iter().map(|&x| format!("{}", x)).join(", ");
+        write!(f, "({})", dims)
+    }
+}
+
+#[macro_export]
+macro_rules! shape {
+    ($($dim:expr),*) => {
+        {
+            use $crate::shape::Shape;
+            Shape::new(vec![$($dim),*])
+        }
+    };
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_shape_new() {
+        let dims = vec![2, 3, 4];
+        let shape = Shape::new(dims.clone()).unwrap();
+        assert_eq!(shape.dims, dims);
+    }
+
+    #[test]
+    fn test_shape_new_with_zero_dimension() {
+        let dims = vec![2, 0, 4];
+        let result = Shape::new(dims);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_shape_size() {
+        let shape = Shape::new(vec![2, 3, 4]).unwrap();
+        assert_eq!(shape.size(), 24);
+    }
+
+    #[test]
+    fn test_shape_display() {
+        let shape = Shape::new(vec![2, 3, 4]).unwrap();
+        assert_eq!(format!("{}", shape), "(2, 3, 4)");
+    }
+
+    #[test]
+    fn test_shape_macro() {
+        let shape = shape![2, 3, 4].unwrap();
+        assert_eq!(shape.dims, vec![2, 3, 4]);
+
+        let shape = shape![1].unwrap();
+        assert_eq!(shape.dims, vec![1]);
+    }
+
+    #[test]
+    fn test_shape_extend() {
+        let shape1 = Shape::new(vec![2, 3]).unwrap();
+        let shape2 = Shape::new(vec![4, 5]).unwrap();
+        let extended_shape = shape1.stack(&shape2);
+        assert_eq!(extended_shape.dims, vec![2, 3, 4, 5]);
+    }
+}

src/storage.rs

@@ -0,0 +1,178 @@
+use std::ops::{Index, IndexMut};
+
+use crate::coordinate::Coordinate;
+use crate::error::ShapeError;
+use crate::shape::Shape;
+
+#[derive(Debug, PartialEq)]
+pub struct DynamicStorage<T> {
+    data: Vec<T>,
+}
+
+impl<T> DynamicStorage<T> {
+    pub fn new(data: Vec<T>) -> Self {
+        Self { data }
+    }
+
+    /// For the row-wise maths see: https://bit.ly/3KQjPa3
+    pub fn flatten(&self, coord: &Coordinate, shape: &Shape) -> Result<usize, ShapeError> {
+        if coord.order() != shape.order() {
+            let msg = format!("incorrect order ({} vs {}).", coord.order(), shape.order());
+            return Err(ShapeError::new(msg.as_str()));
+        }
+
+        for (i, &dim) in coord.iter().enumerate() {
+            if dim >= shape[i] {
+                return Err(ShapeError::new(
+                    format!("out of bounds for dimension {}", i).as_str(),
+                ));
+            }
+        }
+
+        let mut index = 0;
+        for k in 0..shape.order() {
+            let stride = shape[k + 1..].iter().product::<usize>();
+            index += coord[k] * stride;
+        }
+        Ok(index)
+    }
+
+    pub fn size(&self) -> usize {
+        self.data.len()
+    }
+
+    pub fn iter(&self) -> std::slice::Iter<'_, T> {
+        self.data.iter()
+    }
+
+    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
+        self.data.iter_mut()
+    }
+}
+
+impl<T> Index<usize> for DynamicStorage<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.data[index]
+    }
+}
+
+impl<T> Index<std::ops::Range<usize>> for DynamicStorage<T> {
+    type Output = [T];
+
+    fn index(&self, range: std::ops::Range<usize>) -> &Self::Output {
+        &self.data[range]
+    }
+}
+
+impl<T> IndexMut<usize> for DynamicStorage<T> {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        &mut self.data[index]
+    }
+}
+
+impl<T> IntoIterator for DynamicStorage<T> {
+    type Item = T;
+    type IntoIter = std::vec::IntoIter<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.data.into_iter()
+    }
+}
+
+impl<'a, T> IntoIterator for &'a DynamicStorage<T> {
+    type Item = &'a T;
+    type IntoIter = std::slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.data.iter()
+    }
+}
+
+impl<'a, T> IntoIterator for &'a mut DynamicStorage<T> {
+    type Item = &'a mut T;
+    type IntoIter = std::slice::IterMut<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.data.iter_mut()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_iter() {
+        let storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        let mut iter = storage.iter();
+        assert_eq!(iter.next(), Some(&1));
+        assert_eq!(iter.next(), Some(&2));
+        assert_eq!(iter.next(), Some(&3));
+        assert_eq!(iter.next(), Some(&4));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_iter_mut() {
+        let mut storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        {
+            let mut iter = storage.iter_mut();
+            if let Some(x) = iter.next() {
+                *x = 10;
+            }
+        }
+        assert_eq!(storage.data, vec![10, 2, 3, 4]);
+    }
+
+    #[test]
+    fn test_index() {
+        let storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        assert_eq!(storage[0], 1);
+        assert_eq!(storage[1], 2);
+        assert_eq!(storage[2], 3);
+        assert_eq!(storage[3], 4);
+    }
+
+    #[test]
+    fn test_index_mut() {
+        let mut storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        storage[0] = 10;
+        assert_eq!(storage[0], 10);
+    }
+
+    #[test]
+    fn test_into_iter() {
+        let storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        let mut iter = storage.into_iter();
+        assert_eq!(iter.next(), Some(1));
+        assert_eq!(iter.next(), Some(2));
+        assert_eq!(iter.next(), Some(3));
+        assert_eq!(iter.next(), Some(4));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_into_iter_ref() {
+        let storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        let mut iter = (&storage).into_iter();
+        assert_eq!(iter.next(), Some(&1));
+        assert_eq!(iter.next(), Some(&2));
+        assert_eq!(iter.next(), Some(&3));
+        assert_eq!(iter.next(), Some(&4));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_into_iter_mut() {
+        let mut storage = DynamicStorage::new(vec![1, 2, 3, 4]);
+        {
+            let mut iter = (&mut storage).into_iter();
+            if let Some(x) = iter.next() {
+                *x = 10;
+            }
+        }
+        assert_eq!(storage.data, vec![10, 2, 3, 4]);
+    }
+}