tool/core/
matrix.rs

1use crate::{List, VectorsGeneric};
2use na::{
3    storage::Storage, ClosedAdd, ClosedDiv, ClosedMul, ClosedSub, Dim, Dynamic, Matrix, Matrix1xX,
4    MatrixSlice, RealField, Scalar, U3,
5};
6use num_traits::{cast, NumCast, Signed, Zero};
7use std::cmp::PartialOrd;
8
9/// Get the slice of the whole matrix.
10pub fn slice<T, R, C, S>(
11    matrix: &Matrix<T, R, C, S>,
12) -> MatrixSlice<'_, T, Dynamic, Dynamic, S::RStride, S::CStride>
13where
14    T: Scalar,
15    R: Dim,
16    C: Dim,
17    S: Storage<T, R, C>,
18{
19    matrix.slice((0, 0), matrix.shape())
20}
21
22/// Get the number of [`Vector`][crate::Vector]s.
23pub fn number_vectors<T, S>(vectors: &VectorsGeneric<T, S>) -> usize
24where
25    T: Scalar,
26    S: Storage<T, U3, Dynamic>,
27{
28    vectors.ncols()
29}
30
31/// Compute the required size to go from `start` to `end` with `step`, including the end point (last
32/// time step can be smaller).
33pub fn size_range_with_step<T>(start: T, end: T, step: T) -> usize
34where
35    T: Scalar + NumCast + Copy + ClosedAdd + ClosedSub + ClosedDiv + ClosedMul + PartialOrd,
36{
37    if Zero::is_zero(&cast::<T, f64>(step).unwrap())
38        || Zero::is_zero(&cast::<T, f64>(end - start).unwrap())
39        || (Signed::is_positive(&cast::<T, f64>(end - start).unwrap())
40            && Signed::is_negative(&cast::<T, f64>(step).unwrap()))
41        || (Signed::is_negative(&cast::<T, f64>(end - start).unwrap())
42            && Signed::is_positive(&cast::<T, f64>(step).unwrap()))
43    {
44        return 1;
45    }
46    let mut size: usize = cast((end - start) / step).unwrap();
47    if start + cast::<usize, T>(size).unwrap() * step < end {
48        size += 1;
49    }
50    size + 1
51}
52
53/// Create a [`List`] from `start` to `end` with `step`. The last step can be smaller to include
54/// `end`.
55pub fn linspace<T>(start: T, end: T, step: T) -> List<T>
56where
57    T: Scalar + NumCast + ClosedSub + ClosedDiv + ClosedMul + ClosedAdd + PartialOrd + Copy,
58{
59    let size = size_range_with_step(start, end, step);
60    let mut vector = Matrix1xX::from_fn(size, |_, j| start + step * cast(j).unwrap());
61    if vector[size - 1] > end {
62        vector[size - 1] = end;
63    }
64    vector
65}
66
67/// Clip all elements of a [`List`] between `min` and `max`. A `None` value indicates no limit.
68pub fn clip<T>(list: &List<T>, min: Option<T>, max: Option<T>) -> List<T>
69where
70    T: RealField,
71{
72    let mut work_list = list.clone();
73    for element in work_list.iter_mut() {
74        if let Some(mini) = min {
75            if *element < mini {
76                *element = mini
77            };
78        }
79        if let Some(max) = max {
80            if *element > max {
81                *element = max
82            };
83        }
84    }
85    work_list
86}
87
88/// Compute the element-wise power of a [`List`].
89pub fn pows<T>(list: &List<T>, power: i32) -> List<T>
90where
91    T: RealField,
92{
93    List::from_iterator(list.len(), list.iter().map(|x| x.powi(power)))
94}