use ndarray::{Array, Array1, Array2, Dimension, Ix1, Ix2};
use num::complex::Complex;
use num::traits::Float;
pub trait Angles {
fn to_radians(&self) -> Self;
fn to_degrees(&self) -> Self;
}
impl<T: Float + ndarray::ScalarOperand> Angles for Array1<T> {
fn to_radians(&self) -> Self {
self.mapv(|angle| angle.to_radians())
}
fn to_degrees(&self) -> Self {
self.mapv(|angle| angle.to_degrees())
}
}
impl<T: Float + ndarray::ScalarOperand> Angles for Array2<T> {
fn to_radians(&self) -> Self {
self.mapv(|angle| angle.to_radians())
}
fn to_degrees(&self) -> Self {
self.mapv(|angle| angle.to_degrees())
}
}
pub trait ComplexExt<T, D: Dimension> {
fn abs(&self) -> Array<T, D>;
fn imag(&self) -> Array<T, D>;
fn real(&self) -> Array<T, D>;
fn phase(&self) -> Array<T, D>;
fn from_real_imag(real: Array<T, D>, imag: Array<T, D>) -> Array<Complex<T>, D>;
}
impl<T: Float> ComplexExt<T, Ix1> for Array1<Complex<T>> {
fn abs(&self) -> Array<T, Ix1> {
self.map(|c| c.norm())
}
fn imag(&self) -> Array<T, Ix1> {
self.map(|c| c.im)
}
fn real(&self) -> Array<T, Ix1> {
self.map(|c| c.re)
}
fn phase(&self) -> Array<T, Ix1> {
self.map(|c| c.arg())
}
fn from_real_imag(real: Array<T, Ix1>, imag: Array<T, Ix1>) -> Array<Complex<T>, Ix1> {
let vec: Vec<_> = real
.iter()
.zip(imag.iter())
.map(|(&r, &i)| Complex::new(r, i))
.collect();
Array::from_shape_vec(real.raw_dim(), vec).unwrap()
}
}
impl<T: Float> ComplexExt<T, Ix2> for Array2<Complex<T>> {
fn abs(&self) -> Array2<T> {
self.map(|c| c.norm())
}
fn imag(&self) -> Array2<T> {
self.map(|c| c.im)
}
fn real(&self) -> Array2<T> {
self.map(|c| c.re)
}
fn phase(&self) -> Array2<T> {
self.map(|c| c.arg())
}
fn from_real_imag(real: Array<T, Ix2>, imag: Array<T, Ix2>) -> Array2<Complex<T>> {
let vec: Vec<_> = real
.iter()
.zip(imag.iter())
.map(|(&r, &i)| Complex::new(r, i))
.collect();
Array::from_shape_vec(real.raw_dim(), vec).unwrap()
}
}
pub trait FromVecVec<T: Clone> {
fn from_vec_vec(vec_vec: Vec<Vec<T>>) -> Array2<T> {
let rows = vec_vec.len();
let cols = vec_vec[0].len();
Array2::from_shape_vec((rows, cols), vec_vec.into_iter().flatten().collect())
.expect("Failed to create Array2, incorrect shape")
}
}
impl FromVecVec<f32> for Array2<f32> {}
impl FromVecVec<f64> for Array2<f64> {}
impl FromVecVec<Complex<f32>> for Array2<Complex<f32>> {}
impl FromVecVec<Complex<f64>> for Array2<Complex<f64>> {}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::{arr1, arr2};
use num_complex::Complex64;
#[test]
fn test_from_real_imag_1x3() {
let real = arr1(&[1.0, 2.0, 3.0]);
let imag = arr1(&[4.0, 5.0, 6.0]);
let expected = arr1(&[
Complex64::new(1.0, 4.0),
Complex64::new(2.0, 5.0),
Complex64::new(3.0, 6.0),
]);
assert_eq!(Array1::from_real_imag(real, imag), expected);
}
#[test]
fn test_from_real_imag_2x2() {
let real = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
let imag = arr2(&[[5.0, 6.0], [7.0, 8.0]]);
let expected = arr2(&[
[Complex64::new(1.0, 5.0), Complex64::new(2.0, 6.0)],
[Complex64::new(3.0, 7.0), Complex64::new(4.0, 8.0)],
]);
assert_eq!(Array2::from_real_imag(real, imag), expected);
}
#[test]
fn test_from_real_imag_2x3() {
let real = arr2(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
let imag = arr2(&[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]);
let expected = arr2(&[
[
Complex64::new(1.0, 7.0),
Complex64::new(2.0, 8.0),
Complex64::new(3.0, 9.0),
],
[
Complex64::new(4.0, 10.0),
Complex64::new(5.0, 11.0),
Complex64::new(6.0, 12.0),
],
]);
assert_eq!(Array2::from_real_imag(real, imag), expected);
}
}