ndspec/core/
ndarray_ext.rs

1use ndarray::{Array, Array1, Array2, Dimension, Ix1, Ix2};
2use num::complex::Complex;
3use num::traits::Float;
4
5/// Trait for converting angles between radians and degrees.
6///
7/// This trait is implemented for both 1D and 2D arrays.
8/// - `to_radians` converts angles from degrees to radians.
9/// - `to_degrees` converts angles from radians to degrees.
10pub trait Angles {
11    fn to_radians(&self) -> Self;
12    fn to_degrees(&self) -> Self;
13}
14
15/// Implement the `Angles` trait for 1D arrays of floating-point numbers.
16impl<T: Float + ndarray::ScalarOperand> Angles for Array1<T> {
17    fn to_radians(&self) -> Self {
18        self.mapv(|angle| angle.to_radians())
19    }
20
21    fn to_degrees(&self) -> Self {
22        self.mapv(|angle| angle.to_degrees())
23    }
24}
25
26/// Implement the `Angles` trait for 2D arrays of floating-point numbers.
27impl<T: Float + ndarray::ScalarOperand> Angles for Array2<T> {
28    fn to_radians(&self) -> Self {
29        self.mapv(|angle| angle.to_radians())
30    }
31
32    fn to_degrees(&self) -> Self {
33        self.mapv(|angle| angle.to_degrees())
34    }
35}
36
37/// Trait for converting between real and complex numbers.
38///
39/// This trait is implemented for both 1D and 2D arrays.
40pub trait ComplexExt<T, D: Dimension> {
41    fn abs(&self) -> Array<T, D>;
42    fn imag(&self) -> Array<T, D>;
43    fn real(&self) -> Array<T, D>;
44    fn phase(&self) -> Array<T, D>;
45    fn from_real_imag(real: Array<T, D>, imag: Array<T, D>) -> Array<Complex<T>, D>;
46}
47
48/// Implement the `ComplexExt` trait for 1D arrays.
49impl<T: Float> ComplexExt<T, Ix1> for Array1<Complex<T>> {
50    fn abs(&self) -> Array<T, Ix1> {
51        self.map(|c| c.norm())
52    }
53
54    fn imag(&self) -> Array<T, Ix1> {
55        self.map(|c| c.im)
56    }
57
58    fn real(&self) -> Array<T, Ix1> {
59        self.map(|c| c.re)
60    }
61
62    fn phase(&self) -> Array<T, Ix1> {
63        self.map(|c| c.arg())
64    }
65
66    fn from_real_imag(real: Array<T, Ix1>, imag: Array<T, Ix1>) -> Array<Complex<T>, Ix1> {
67        let vec: Vec<_> = real
68            .iter()
69            .zip(imag.iter())
70            .map(|(&r, &i)| Complex::new(r, i))
71            .collect();
72        Array::from_shape_vec(real.raw_dim(), vec).unwrap()
73    }
74}
75
76/// Implement the `ComplexExt` trait for 2D arrays.
77impl<T: Float> ComplexExt<T, Ix2> for Array2<Complex<T>> {
78    fn abs(&self) -> Array2<T> {
79        self.map(|c| c.norm())
80    }
81
82    fn imag(&self) -> Array2<T> {
83        self.map(|c| c.im)
84    }
85
86    fn real(&self) -> Array2<T> {
87        self.map(|c| c.re)
88    }
89
90    fn phase(&self) -> Array2<T> {
91        self.map(|c| c.arg())
92    }
93
94    fn from_real_imag(real: Array<T, Ix2>, imag: Array<T, Ix2>) -> Array2<Complex<T>> {
95        let vec: Vec<_> = real
96            .iter()
97            .zip(imag.iter())
98            .map(|(&r, &i)| Complex::new(r, i))
99            .collect();
100        Array::from_shape_vec(real.raw_dim(), vec).unwrap()
101    }
102}
103
104/// `FromVecVec` is a trait for converting a vector of vectors (2D vector) into a 2D array.
105///
106/// This trait is generic over a type `T` that implements `Clone`.
107/// The `from_vec_vec` method takes a vector of vectors (`Vec<Vec<T>>`) and
108/// returns a 2D array (`Array2<T>`).
109///
110/// # Panics
111///
112/// The `from_vec_vec` method will panic if the provided vector of vectors
113/// cannot be shaped into a 2D array.
114///
115pub trait FromVecVec<T: Clone> {
116    fn from_vec_vec(vec_vec: Vec<Vec<T>>) -> Array2<T> {
117        let rows = vec_vec.len();
118        let cols = vec_vec[0].len();
119
120        Array2::from_shape_vec((rows, cols), vec_vec.into_iter().flatten().collect())
121            .expect("Failed to create Array2, incorrect shape")
122    }
123}
124
125/// Implementation of `FromVecVec` for `Array2<f32>`.
126impl FromVecVec<f32> for Array2<f32> {}
127
128/// Implementation of `FromVecVec` for `Array2<f64>`.
129impl FromVecVec<f64> for Array2<f64> {}
130
131/// Implementation of `FromVecVec` for `Array2<Complex<f32>>`.
132impl FromVecVec<Complex<f32>> for Array2<Complex<f32>> {}
133
134/// Implementation of `FromVecVec` for `Array2<Complex<f64>>`.
135impl FromVecVec<Complex<f64>> for Array2<Complex<f64>> {}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use ndarray::{arr1, arr2};
141    use num_complex::Complex64;
142
143    #[test]
144    fn test_from_real_imag_1x3() {
145        let real = arr1(&[1.0, 2.0, 3.0]);
146        let imag = arr1(&[4.0, 5.0, 6.0]);
147        let expected = arr1(&[
148            Complex64::new(1.0, 4.0),
149            Complex64::new(2.0, 5.0),
150            Complex64::new(3.0, 6.0),
151        ]);
152        assert_eq!(Array1::from_real_imag(real, imag), expected);
153    }
154
155    #[test]
156    fn test_from_real_imag_2x2() {
157        let real = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
158        let imag = arr2(&[[5.0, 6.0], [7.0, 8.0]]);
159        let expected = arr2(&[
160            [Complex64::new(1.0, 5.0), Complex64::new(2.0, 6.0)],
161            [Complex64::new(3.0, 7.0), Complex64::new(4.0, 8.0)],
162        ]);
163        assert_eq!(Array2::from_real_imag(real, imag), expected);
164    }
165
166    #[test]
167    fn test_from_real_imag_2x3() {
168        let real = arr2(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
169        let imag = arr2(&[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]);
170        let expected = arr2(&[
171            [
172                Complex64::new(1.0, 7.0),
173                Complex64::new(2.0, 8.0),
174                Complex64::new(3.0, 9.0),
175            ],
176            [
177                Complex64::new(4.0, 10.0),
178                Complex64::new(5.0, 11.0),
179                Complex64::new(6.0, 12.0),
180            ],
181        ]);
182        assert_eq!(Array2::from_real_imag(real, imag), expected);
183    }
184}
ndspec/core/ndarray_ext.rs

ndspec/core/
ndarray_ext.rs
