concision_utils/utils/
tensor.rs

1/*
2    Appellation: tensor <module>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5pub use self::{generators::*, stack::*};
6use ndarray::*;
7use num_traits::{NumAssign, Zero};
8
9#[cfg(feature = "alloc")]
10/// Creates an n-dimensional array from an iterator of n dimensional arrays.
11pub fn concat_iter<D, T>(axis: usize, iter: impl IntoIterator<Item = Array<T, D>>) -> Array<T, D>
12where
13    D: RemoveAxis,
14    T: Clone,
15{
16    let mut arr = iter.into_iter().collect::<alloc::vec::Vec<_>>();
17    let mut out = arr.pop().unwrap();
18    for i in arr {
19        out = concatenate!(Axis(axis), out, i);
20    }
21    out
22}
23
24pub fn inverse<T>(matrix: &Array2<T>) -> Option<Array2<T>>
25where
26    T: Copy + NumAssign + ScalarOperand,
27{
28    let (rows, cols) = matrix.dim();
29
30    if !matrix.is_square() {
31        return None; // Matrix must be square for inversion
32    }
33
34    let identity = Array2::eye(rows);
35
36    // Construct an augmented matrix by concatenating the original matrix with an identity matrix
37    let mut aug = Array2::zeros((rows, 2 * cols));
38    aug.slice_mut(s![.., ..cols]).assign(matrix);
39    aug.slice_mut(s![.., cols..]).assign(&identity);
40
41    // Perform Gaussian elimination to reduce the left half to the identity matrix
42    for i in 0..rows {
43        let pivot = aug[[i, i]];
44
45        if pivot == T::zero() {
46            return None; // Matrix is singular
47        }
48
49        aug.slice_mut(s![i, ..]).mapv_inplace(|x| x / pivot);
50
51        for j in 0..rows {
52            if i != j {
53                let am = aug.clone();
54                let factor = aug[[j, i]];
55                let rhs = am.slice(s![i, ..]);
56                aug.slice_mut(s![j, ..])
57                    .zip_mut_with(&rhs, |x, &y| *x -= y * factor);
58            }
59        }
60    }
61
62    // Extract the inverted matrix from the augmented matrix
63    let inverted = aug.slice(s![.., cols..]);
64
65    Some(inverted.to_owned())
66}
67
68/// Returns the lower triangular portion of a matrix.
69pub fn tril<T>(a: &Array2<T>) -> Array2<T>
70where
71    T: Clone + Zero,
72{
73    let mut out = a.clone();
74    for i in 0..a.shape()[0] {
75        for j in i + 1..a.shape()[1] {
76            out[[i, j]] = T::zero();
77        }
78    }
79    out
80}
81/// Returns the upper triangular portion of a matrix.
82pub fn triu<T>(a: &Array2<T>) -> Array2<T>
83where
84    T: Clone + Zero,
85{
86    let mut out = a.clone();
87    for i in 0..a.shape()[0] {
88        for j in 0..i {
89            out[[i, j]] = T::zero();
90        }
91    }
92    out
93}
94
95pub(crate) mod generators {
96    use ndarray::{Array, Array1, Dimension, IntoDimension, ShapeError};
97    use num_traits::{Float, NumCast};
98
99    pub fn genspace<T: NumCast>(features: usize) -> Array1<T> {
100        Array1::from_iter((0..features).map(|x| T::from(x).unwrap()))
101    }
102
103    pub fn linarr<A, D>(dim: impl Clone + IntoDimension<Dim = D>) -> Result<Array<A, D>, ShapeError>
104    where
105        A: Float,
106        D: Dimension,
107    {
108        let dim = dim.into_dimension();
109        let n = dim.size();
110        Array::linspace(A::zero(), A::from(n - 1).unwrap(), n)
111            .to_shape(dim)
112            .map(|x| x.to_owned())
113    }
114}
115
116pub(crate) mod stack {
117    #[cfg(feature = "alloc")]
118    use alloc::vec::Vec;
119    use ndarray::{Array1, Array2, s};
120    use num_traits::Num;
121    /// Creates a larger array from an iterator of smaller arrays.
122    pub fn stack_iter<T>(iter: impl IntoIterator<Item = Array1<T>>) -> Array2<T>
123    where
124        T: Clone + Num,
125    {
126        let mut iter = iter.into_iter();
127        let first = iter.next().unwrap();
128        let shape = [iter.size_hint().0 + 1, first.len()];
129        let mut res = Array2::<T>::zeros(shape);
130        res.slice_mut(s![0, ..]).assign(&first);
131        for (i, s) in iter.enumerate() {
132            res.slice_mut(s![i + 1, ..]).assign(&s);
133        }
134        res
135    }
136    #[cfg(feature = "alloc")]
137    /// stack a 1D array into a 2D array by stacking them horizontally.
138    pub fn hstack<T>(iter: impl IntoIterator<Item = Array1<T>>) -> Array2<T>
139    where
140        T: Clone + Num,
141    {
142        let iter = Vec::from_iter(iter);
143        let mut res = Array2::<T>::zeros((iter.first().unwrap().len(), iter.len()));
144        for (i, s) in iter.iter().enumerate() {
145            res.slice_mut(s![.., i]).assign(s);
146        }
147        res
148    }
149    #[cfg(feature = "alloc")]
150    /// stack a 1D array into a 2D array by stacking them vertically.
151    pub fn vstack<T>(iter: impl IntoIterator<Item = Array1<T>>) -> Array2<T>
152    where
153        T: Clone + Num,
154    {
155        let iter = Vec::from_iter(iter);
156        let mut res = Array2::<T>::zeros((iter.len(), iter.first().unwrap().len()));
157        for (i, s) in iter.iter().enumerate() {
158            res.slice_mut(s![i, ..]).assign(s);
159        }
160        res
161    }
162}
concision_utils/utils/tensor.rs

concision_utils/utils/
tensor.rs
