concision_core/utils/
tensor.rs

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