diff --git a/src/numeric/impl_float_maths.rs b/src/numeric/impl_float_maths.rs index 358d57cf3..f401c47e8 100644 --- a/src/numeric/impl_float_maths.rs +++ b/src/numeric/impl_float_maths.rs @@ -1,7 +1,9 @@ // Element-wise methods for ndarray #[cfg(feature = "std")] -use num_traits::Float; +use num_complex::Complex; +#[cfg(feature = "std")] +use num_traits::{Float, Zero}; use crate::imp_prelude::*; @@ -167,6 +169,96 @@ where } } +#[cfg(feature = "std")] +impl ArrayRef +where + D: Dimension, + A: Clone + Zero, +{ + /// Map the array into the real part of a complex array; the imaginary part is 0. + /// + /// # Example + /// ``` + /// use ndarray::*; + /// + /// let arr = array![1.0, -1.0, 0.0]; + /// let complex = arr.to_complex_re(); + /// + /// assert_eq!(complex[0], Complex::new(1.0, 0.0)); + /// assert_eq!(complex[1], Complex::new(-1.0, 0.0)); + /// assert_eq!(complex[2], Complex::new(0.0, 0.0)); + /// ``` + /// + /// # See Also + /// [ArrayRef::to_complex_im] + #[must_use = "method returns a new array and does not mutate the original value"] + pub fn to_complex_re(&self) -> Array, D> + { + self.mapv(|v| Complex::new(v, A::zero())) + } + + /// Map the array into the imaginary part of a complex array; the real part is 0. + /// + /// # Example + /// ``` + /// use ndarray::*; + /// + /// let arr = array![1.0, -1.0, 0.0]; + /// let complex = arr.to_complex_re(); + /// + /// assert_eq!(complex[0], Complex::new(0.0, 1.0)); + /// assert_eq!(complex[1], Complex::new(0.0, -1.0)); + /// assert_eq!(complex[2], Complex::new(0.0, 0.0)); + /// ``` + /// + /// # See Also + /// [ArrayRef::to_complex_re] + #[must_use = "method returns a new array and does not mutate the original value"] + pub fn to_complex_im(&self) -> Array, D> + { + self.mapv(|v| Complex::new(A::zero(), v)) + } +} + +/// # Angle calculation methods for arrays +/// +/// Methods for calculating phase angles of complex values in arrays. +#[cfg(feature = "std")] +impl ArrayRef, D> +where + D: Dimension, + A: Float, +{ + /// Return the [phase angle (argument)](https://en.wikipedia.org/wiki/Argument_(complex_analysis)) of complex values in the array. + /// + /// This function always returns the same float type as was provided to it. Leaving the exact precision left to the user. + /// The angles are returned in ``radians`` and in the range ``(-π, π]``. + /// To get the angles in degrees, use the [`to_degrees()`][ArrayRef::to_degrees] method on the resulting array. + /// + /// # Examples + /// + /// ``` + /// use ndarray::array; + /// use num_complex::Complex; + /// use std::f64::consts::PI; + /// + /// let complex_arr = array![ + /// Complex::new(1.0f64, 0.0), + /// Complex::new(0.0, 1.0), + /// Complex::new(1.0, 1.0), + /// ]; + /// let angles = complex_arr.angle(); + /// assert!((angles[0] - 0.0).abs() < 1e-10); + /// assert!((angles[1] - PI/2.0).abs() < 1e-10); + /// assert!((angles[2] - PI/4.0).abs() < 1e-10); + /// ``` + #[must_use = "method returns a new array and does not mutate the original value"] + pub fn angle(&self) -> Array + { + self.mapv(|v| v.im.atan2(v.re)) + } +} + impl ArrayRef where A: 'static + PartialOrd + Clone, @@ -191,3 +283,119 @@ where self.mapv(|a| num_traits::clamp(a, min.clone(), max.clone())) } } + +#[cfg(all(test, feature = "std"))] +mod angle_tests +{ + use crate::Array; + use num_complex::Complex; + use std::f64::consts::PI; + + /// Helper macro for floating-point comparison + macro_rules! assert_approx_eq { + ($a:expr, $b:expr, $tol:expr $(, $msg:expr)?) => {{ + let (a, b) = ($a, $b); + assert!( + (a - b).abs() < $tol, + concat!( + "assertion failed: |left - right| >= tol\n", + " left: {left:?}\n right: {right:?}\n tol: {tol:?}\n", + $($msg,)? + ), + left = a, + right = b, + tol = $tol + ); + }}; + } + + #[test] + fn test_complex_numbers_radians() + { + let arr = Array::from_vec(vec![ + Complex::new(1.0f64, 0.0), // 0 + Complex::new(0.0, 1.0), // π/2 + Complex::new(-1.0, 0.0), // π + Complex::new(0.0, -1.0), // -π/2 + Complex::new(1.0, 1.0), // π/4 + Complex::new(-1.0, -1.0), // -3π/4 + ]); + let a = arr.angle(); + + assert_approx_eq!(a[0], 0.0, 1e-10); + assert_approx_eq!(a[1], PI / 2.0, 1e-10); + assert_approx_eq!(a[2], PI, 1e-10); + assert_approx_eq!(a[3], -PI / 2.0, 1e-10); + assert_approx_eq!(a[4], PI / 4.0, 1e-10); + assert_approx_eq!(a[5], -3.0 * PI / 4.0, 1e-10); + } + + #[test] + fn test_complex_numbers_degrees() + { + let arr = Array::from_vec(vec![ + Complex::new(1.0f64, 0.0), + Complex::new(0.0, 1.0), + Complex::new(-1.0, 0.0), + Complex::new(1.0, 1.0), + ]); + let a = arr.angle().to_degrees(); + + assert_approx_eq!(a[0], 0.0, 1e-10); + assert_approx_eq!(a[1], 90.0, 1e-10); + assert_approx_eq!(a[2], 180.0, 1e-10); + assert_approx_eq!(a[3], 45.0, 1e-10); + } + + #[test] + fn test_signed_zeros() + { + let arr = Array::from_vec(vec![ + Complex::new(0.0f64, 0.0), + Complex::new(-0.0, 0.0), + Complex::new(0.0, -0.0), + Complex::new(-0.0, -0.0), + ]); + let a = arr.angle(); + + assert!(a[0] >= 0.0 && a[0].abs() < 1e-10); + assert_approx_eq!(a[1], PI, 1e-10); + assert!(a[2] <= 0.0 && a[2].abs() < 1e-10); + assert_approx_eq!(a[3], -PI, 1e-10); + } + + #[test] + fn test_edge_cases() + { + let arr = Array::from_vec(vec![ + Complex::new(f64::INFINITY, 0.0), + Complex::new(0.0, f64::INFINITY), + Complex::new(f64::NEG_INFINITY, 0.0), + Complex::new(0.0, f64::NEG_INFINITY), + ]); + let a = arr.angle(); + + assert_approx_eq!(a[0], 0.0, 1e-10); + assert_approx_eq!(a[1], PI / 2.0, 1e-10); + assert_approx_eq!(a[2], PI, 1e-10); + assert_approx_eq!(a[3], -PI / 2.0, 1e-10); + } + + #[test] + fn test_range_validation() + { + let n = 16; + let complex_arr: Vec<_> = (0..n) + .map(|i| { + let theta = 2.0 * PI * (i as f64) / (n as f64); + Complex::new(theta.cos(), theta.sin()) + }) + .collect(); + + let a = Array::from_vec(complex_arr).angle(); + + for &x in &a { + assert!(x > -PI && x <= PI, "Angle {} outside (-π, π]", x); + } + } +}