concision_core/utils/
tensor.rs1#[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; }
36
37 let identity = Array2::eye(rows);
38
39 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 for i in 0..rows {
46 let pivot = aug[[i, i]];
47
48 if pivot == T::zero() {
49 return None; }
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 let inverted = aug.slice(s![.., cols..]);
67
68 Some(inverted.to_owned())
69}
70
71pub 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
87pub 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}
100pub 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 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 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 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}