1use ndarray::*;
2use num_traits::{Float, FromPrimitive};
3use rand::distributions::{uniform::SampleUniform, Distribution, Uniform};
4use rand::rngs::StdRng;
5use rand::SeedableRng;
6
7#[derive(Clone, Copy, PartialEq, Eq)]
8pub enum ShiftDirect {
9 Left,
10 Right,
11}
12
13#[derive(Clone)]
14pub struct Padded2D<T: Float> {
15 pub data: Array2<T>,
16 pub raw_size: (usize, usize),
17 pub padding_size: usize,
18}
19
20impl<T: Float> Padded2D<T> {
21 pub fn from_arr(a: &Array2<T>, padding_size: usize) -> Self {
22 let shape = a.shape();
23 let raw_size = (shape[0], shape[1]);
24 let data = a.padding_zero(padding_size);
25 Self {
26 data,
27 raw_size,
28 padding_size,
29 }
30 }
31
32 pub fn raw_view(&self) -> ArrayView2<T> {
33 self.shift_view(0, ShiftDirect::Left).unwrap()
34 }
35
36 pub fn shift_view(&self, l: usize, shift_dir: ShiftDirect) -> Result<ArrayView2<T>, &str> {
37 let padding_size = self.padding_size;
38 if padding_size < l {
39 return Err("Shitf size l is larger than padding size!");
40 }
41 let raw_size = self.raw_size;
42 let start = match shift_dir {
43 ShiftDirect::Left => padding_size + l,
44 ShiftDirect::Right => padding_size - l,
45 };
46 Ok(self.data.slice(s![.., start..start + raw_size.1]))
47 }
48}
49
50pub trait Padding2D {
51 fn padding_zero(&self, padding_size: usize) -> Self;
52 fn into_padding_zero(&mut self, padding_size: usize);
53}
54
55impl<T> Padding2D for Array2<T>
56where
57 T: Float,
58{
59 fn padding_zero(&self, padding_size: usize) -> Self {
61 let mut a = self.clone();
62 a.into_padding_zero(padding_size);
63 a
64 }
65
66 fn into_padding_zero(&mut self, padding_size: usize) {
67 let shape = self.shape();
68 let mut a: Array2<T> = Array2::zeros((shape[0], shape[1] + padding_size * 2));
69 let mut a_smut = a.slice_mut(s![.., padding_size..padding_size + shape[1]]);
70 Zip::from(&mut a_smut).and(&self.to_owned()).apply(|a, &b| {
71 *a = b;
72 });
73 *self = a;
74 }
75}
76
77pub trait Tensor3D<T>
78where
79 T: Float + std::ops::AddAssign + 'static,
80{
81 fn convolution(&self, rhs: &Padded2D<T>) -> Array2<T>;
82 fn transpose_convolution(&self, rhs: &Padded2D<T>) -> Array2<T>;
83 fn transpose_convolution_tail(&self, rhs: &Padded2D<T>) -> Array2<T>;
84 fn transpose_convolution_headtail(&self, rhs: &Padded2D<T>) -> Array2<T>;
85}
86
87impl<T> Tensor3D<T> for Array3<T>
88where
89 T: Float + std::ops::AddAssign + std::ops::MulAssign + 'static,
90{
91 fn convolution(&self, rhs: &Padded2D<T>) -> Array2<T> {
92 let shape3d = self.shape();
93 let n = shape3d[0];
94 let l = shape3d[2];
95 let t = rhs.raw_size.1;
96 let mut arr = Array2::zeros((n, t));
97 for i in 0..l {
98 let brr = self
99 .slice(s![.., .., i])
100 .dot(&rhs.shift_view(i, ShiftDirect::Right).unwrap());
101 Zip::from(&mut arr).and(&brr).apply(|a, &b| {
102 *a += b;
103 });
104 }
105 arr
106 }
107
108 fn transpose_convolution(&self, rhs: &Padded2D<T>) -> Array2<T> {
109 let shape3d = self.shape();
110 let k = shape3d[1];
111 let l = shape3d[2];
112 let t = rhs.raw_size.1;
113 let mut arr = Array2::zeros((k, t));
114 for i in 0..l {
115 let brr = self
116 .slice(s![.., .., i])
117 .t()
118 .dot(&rhs.shift_view(i, ShiftDirect::Left).unwrap());
119 Zip::from(&mut arr).and(&brr).apply(|a, &b| {
120 *a += b;
121 });
122 }
123 arr
124 }
125
126 fn transpose_convolution_tail(&self, rhs: &Padded2D<T>) -> Array2<T> {
127 let l = self.shape()[2];
128 let t = rhs.raw_size.1;
129 let mut arr = self.transpose_convolution(rhs);
130 for i in 1..l {
131 Zip::from(&mut arr.slice_mut(s![.., t - i])).apply(|a| {
132 *a *= T::from(l).unwrap() / T::from(i).unwrap();
133 })
134 }
135 arr
136 }
137
138 fn transpose_convolution_headtail(&self, rhs: &Padded2D<T>) -> Array2<T> {
139 let l = self.shape()[2];
140 let l2 = l * l;
141 let t = rhs.raw_size.1;
142 let mut arr = self.transpose_convolution(rhs);
143 for i in 1..l {
144 let j = l - i;
145 let active = l2 - j * (j + 1) / 2;
146 Zip::from(&mut arr.slice_mut(s![.., t - i])).apply(|a| {
147 *a *= T::from(l).unwrap() / T::from(i).unwrap();
148 });
149 Zip::from(&mut arr.slice_mut(s![.., i - 1])).apply(|a| {
150 *a *= T::from(l2).unwrap() / T::from(active).unwrap();
151 });
152 }
153 arr
154 }
155}
156
157pub trait Smooth {
158 fn smooth(&self, l: usize) -> Self;
159 fn smooth_x(rhs: &Self, l: usize) -> Self;
160 fn smooth_headtail(&self, l: usize) -> Self;
161 fn smooth_x_headtail(rhs: &Self, l: usize) -> Self;
162}
163
164impl<T> Smooth for Array2<T>
165where
166 T: Float + Clone + std::ops::MulAssign,
167{
168 fn smooth(&self, l: usize) -> Self {
169 let shape = self.shape();
170 let rows = shape[0];
171 let cols = shape[1];
172 let mut arr = Array2::zeros((rows, cols));
173 let mut brr: Array1<T> = Array1::zeros(rows);
174 for i in 0..l - 1 {
175 brr = brr + self.slice(s![.., i]);
176 }
177 for i in 0..cols {
178 let j = i + l - 1;
179 if j < cols {
180 brr = brr + self.slice(s![.., j]);
181 }
182 if i >= l {
183 brr = brr - self.slice(s![.., i - l]);
184 }
185 Zip::from(&mut arr.slice_mut(s![.., i]))
186 .and(&brr)
187 .apply(|a, &b| {
188 *a = b;
189 });
190 }
191 arr
192 }
193
194 fn smooth_x(rhs: &Self, l: usize) -> Self {
195 let shape = rhs.shape();
196 let rows = shape[0];
197 let cols = shape[1];
198 let mut arr = Array2::zeros((rows, cols));
199 let mut brr: Array1<T> = Array1::zeros(cols);
200 for i in 0..l - 1 {
201 brr = brr + rhs.slice(s![i, ..]);
202 }
203 for i in 0..rows {
204 let j = i + l - 1;
205 if j < rows {
206 brr = brr + rhs.slice(s![j, ..]);
207 }
208 if i >= l {
209 brr = brr - rhs.slice(s![i - l, ..]);
210 }
211 Zip::from(&mut arr.slice_mut(s![i, ..]))
212 .and(&brr)
213 .apply(|a, &b| {
214 *a = b;
215 });
216 }
217 arr
218 }
219
220 fn smooth_headtail(&self, l: usize) -> Self {
221 let mut arr = self.smooth(l);
222 let cols = self.shape()[1];
223 let mut active = l;
224 let all = T::from(2 * l - 1).unwrap();
225 for i in 0..l-1 {
226 Zip::from(&mut arr.slice_mut(s![.., i]))
227 .apply(|a| {
228 *a *= all / T::from(active).unwrap();
229 });
230 Zip::from(&mut arr.slice_mut(s![.., cols - i - 1]))
231 .apply(|a| {
232 *a *= all / T::from(active).unwrap();
233 });
234 active += 1;
235 }
236 arr
237 }
238
239 fn smooth_x_headtail(rhs: &Self, l: usize) -> Self {
240 let mut arr = Self::smooth_x(rhs, l);
241 let rows = rhs.shape()[0];
242 let mut active = l;
243 let all = T::from(2 * l - 1).unwrap();
244 for i in 0..l-1 {
245 Zip::from(&mut arr.slice_mut(s![i, ..]))
246 .apply(|a| {
247 *a *= all / T::from(active).unwrap();
248 });
249 Zip::from(&mut arr.slice_mut(s![rows - i - 1, ..]))
250 .apply(|a| {
251 *a *= all / T::from(active).unwrap();
252 });
253 active += 1;
254 }
255 arr
256 }
257}
258
259pub trait RandomArray {
260 type ShapeTuple;
261 fn random(shape: Self::ShapeTuple, random_seed: u8) -> Self;
262 fn shuffle(&self, random_seed: u8) -> Self;
263}
264
265impl<T: Float + SampleUniform> RandomArray for Array2<T> {
266 type ShapeTuple = (usize, usize);
267
268 fn random(shape: Self::ShapeTuple, random_seed: u8) -> Self {
269 let mut arr = Array2::zeros(shape);
270 let mut rng = StdRng::from_seed([random_seed; 32]);
271 let ud = Uniform::new::<T, T>(T::zero(), T::one());
272 Zip::from(&mut arr).apply(|a| {
273 *a = ud.sample(&mut rng);
274 });
275 arr
276 }
277
278 fn shuffle(&self, random_seed: u8) -> Self {
279 let shape = self.shape();
280 let y_size = shape[0];
281 let x_size = shape[1];
282 let mut rng = StdRng::from_seed([random_seed; 32]);
283 let ud = Uniform::new::<usize, usize>(0, x_size);
284 let mut arr = Array2::zeros((y_size, x_size));
285 for y in 0..y_size {
286 let i = ud.sample(&mut rng);
287 Zip::from(&mut arr.slice_mut(s![y, i..]))
288 .and(&self.slice(s![y, ..x_size - i]))
289 .apply(|a, &b| {
290 *a = b;
291 });
292 Zip::from(&mut arr.slice_mut(s![y, ..i]))
293 .and(&self.slice(s![y, x_size - i..]))
294 .apply(|a, &b| {
295 *a = b;
296 });
297 }
298 arr
299 }
300}
301
302impl<T: Float + SampleUniform> RandomArray for Array3<T> {
303 type ShapeTuple = (usize, usize, usize);
304
305 fn random(shape: Self::ShapeTuple, random_seed: u8) -> Self {
306 let mut arr = Array3::zeros(shape);
307 let mut rng = StdRng::from_seed([random_seed; 32]);
308 let ud = Uniform::new::<T, T>(T::zero(), T::one());
309 Zip::from(&mut arr).apply(|a| {
310 *a = ud.sample(&mut rng);
311 });
312 arr
313 }
314
315 fn shuffle(&self, random_seed: u8) -> Self {
316 let shape = self.shape();
317 let n = shape[0];
318 let k = shape[1];
319 let l = shape[2];
320 let mut rng = StdRng::from_seed([random_seed; 32]);
321 let ud = Uniform::new::<usize, usize>(0, l);
322 let mut arr = Array3::zeros((n, k, l));
323 for ki in 0..k {
324 for ni in 0..n {
325 let i = ud.sample(&mut rng);
326 Zip::from(&mut arr.slice_mut(s![ni, ki, i..]))
327 .and(&self.slice(s![ni, ki, ..l - i]))
328 .apply(|a, &b| {
329 *a = b;
330 });
331 Zip::from(&mut arr.slice_mut(s![ni, ki, ..i]))
332 .and(&self.slice(s![ni, ki, l - i..]))
333 .apply(|a, &b| {
334 *a = b;
335 });
336 }
337 }
338 arr
339 }
340}
341
342pub fn frobenius_norm_2<T: Float>(a: &Array2<T>) -> T {
343 a.fold(T::zero(), |m, &x| m + x * x)
344}
345
346pub fn norm_1st_ij<T: Float>(a: &Array2<T>) -> T {
347 a.sum() - a.diag().sum()
348}
349
350pub trait Skewness<T> {
351 fn skewness_core(&self) -> Array1<T>;
352}
353
354impl<T: Float + FromPrimitive> Skewness<T> for Array2<T> {
355 fn skewness_core(&self) -> Array1<T> {
356 let n = self.shape()[0];
357 let mu = self.mean_axis(Axis(1)).unwrap().into_shape((n, 1)).unwrap();
358 let sigma = self.std_axis(Axis(1), T::one()).into_shape((n, 1)).unwrap();
359 ((self - &mu) / sigma).fold_axis(Axis(1), T::zero(), |&m, &x| m + x * x * x)
360 }
361}
362
363pub trait OrdArray<T: Clone, D: RemoveAxis> {
364 fn max_axis(&self, axis: Axis) -> Array<T, D::Smaller>;
365 fn min_axis(&self, axis: Axis) -> Array<T, D::Smaller>;
366}
367
368impl<T: Float, D: RemoveAxis> OrdArray<T, D> for Array<T, D> {
369 fn max_axis(&self, axis: Axis) -> Array<T, D::Smaller> {
370 self.fold_axis(axis, T::neg_infinity(), |&m, &x| if m < x { x } else { m })
371 }
372
373 fn min_axis(&self, axis: Axis) -> Array<T, D::Smaller> {
374 self.fold_axis(axis, T::infinity(), |&m, &x| if m > x { x } else { m })
375 }
376}
377
378pub fn reconstruction_cost<T: Float>(x_bar: &Array2<T>, x: &Array2<T>) -> T {
379 frobenius_norm_2(&(x_bar - x))
380}
381
382pub fn xortho_cost<T: Float + std::ops::MulAssign + 'static>(wtx: &Array2<T>, h: &Array2<T>, l: usize) -> T {
383 norm_1st_ij(&wtx.smooth(l).dot(&h.t()))
384}