Skip to content

Commit 49cbd5d

Browse files
jmg049akern40
andauthored
Add phase angle calculation functions for complex arrays (#1543)
Implement phase angle calculation for complex arrays, calculated as atan2(imaginary, real) in radians. Angles are wrapped to (-pi, pi]. --------- Co-authored-by: Adam Kern <[email protected]>
1 parent f6a3205 commit 49cbd5d

File tree

1 file changed

+211
-1
lines changed

1 file changed

+211
-1
lines changed

src/numeric/impl_float_maths.rs

Lines changed: 211 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
// Element-wise methods for ndarray
22

33
#[cfg(feature = "std")]
4-
use num_traits::Float;
4+
use num_complex::Complex;
5+
#[cfg(feature = "std")]
6+
use num_traits::{Float, Zero};
57

68
use crate::imp_prelude::*;
79

@@ -166,6 +168,98 @@ where
166168
}
167169
}
168170

171+
#[cfg(feature = "std")]
172+
impl<A, D> ArrayRef<A, D>
173+
where
174+
D: Dimension,
175+
A: Clone + Zero,
176+
{
177+
/// Map the array into the real part of a complex array; the imaginary part is 0.
178+
///
179+
/// # Example
180+
/// ```
181+
/// use ndarray::*;
182+
/// use num_complex::Complex;
183+
///
184+
/// let arr = array![1.0, -1.0, 0.0];
185+
/// let complex = arr.to_complex_re();
186+
///
187+
/// assert_eq!(complex[0], Complex::new(1.0, 0.0));
188+
/// assert_eq!(complex[1], Complex::new(-1.0, 0.0));
189+
/// assert_eq!(complex[2], Complex::new(0.0, 0.0));
190+
/// ```
191+
///
192+
/// # See Also
193+
/// [ArrayRef::to_complex_im]
194+
#[must_use = "method returns a new array and does not mutate the original value"]
195+
pub fn to_complex_re(&self) -> Array<Complex<A>, D>
196+
{
197+
self.mapv(|v| Complex::new(v, A::zero()))
198+
}
199+
200+
/// Map the array into the imaginary part of a complex array; the real part is 0.
201+
///
202+
/// # Example
203+
/// ```
204+
/// use ndarray::*;
205+
/// use num_complex::Complex;
206+
///
207+
/// let arr = array![1.0, -1.0, 0.0];
208+
/// let complex = arr.to_complex_im();
209+
///
210+
/// assert_eq!(complex[0], Complex::new(0.0, 1.0));
211+
/// assert_eq!(complex[1], Complex::new(0.0, -1.0));
212+
/// assert_eq!(complex[2], Complex::new(0.0, 0.0));
213+
/// ```
214+
///
215+
/// # See Also
216+
/// [ArrayRef::to_complex_re]
217+
#[must_use = "method returns a new array and does not mutate the original value"]
218+
pub fn to_complex_im(&self) -> Array<Complex<A>, D>
219+
{
220+
self.mapv(|v| Complex::new(A::zero(), v))
221+
}
222+
}
223+
224+
/// # Angle calculation methods for arrays
225+
///
226+
/// Methods for calculating phase angles of complex values in arrays.
227+
#[cfg(feature = "std")]
228+
impl<A, D> ArrayRef<Complex<A>, D>
229+
where
230+
D: Dimension,
231+
A: Float,
232+
{
233+
/// Return the [phase angle (argument)](https://en.wikipedia.org/wiki/Argument_(complex_analysis)) of complex values in the array.
234+
///
235+
/// This function always returns the same float type as was provided to it. Leaving the exact precision left to the user.
236+
/// The angles are returned in ``radians`` and in the range ``(-π, π]``.
237+
/// To get the angles in degrees, use the [`to_degrees()`][ArrayRef::to_degrees] method on the resulting array.
238+
///
239+
/// # Examples
240+
///
241+
/// ```
242+
/// use ndarray::array;
243+
/// use num_complex::Complex;
244+
/// use std::f64::consts::PI;
245+
///
246+
/// let complex_arr = array![
247+
/// Complex::new(1.0f64, 0.0),
248+
/// Complex::new(0.0, 1.0),
249+
/// Complex::new(1.0, 1.0),
250+
/// ];
251+
/// let angles = complex_arr.angle();
252+
/// assert!((angles[0] - 0.0).abs() < 1e-10);
253+
/// assert!((angles[1] - PI/2.0).abs() < 1e-10);
254+
/// assert!((angles[2] - PI/4.0).abs() < 1e-10);
255+
/// ```
256+
#[must_use = "method returns a new array and does not mutate the original value"]
257+
pub fn angle(&self) -> Array<A, D>
258+
{
259+
self.mapv(|v| v.im.atan2(v.re))
260+
}
261+
}
262+
169263
impl<A, D> ArrayRef<A, D>
170264
where
171265
A: 'static + PartialOrd + Clone,
@@ -190,3 +284,119 @@ where
190284
self.mapv(|a| num_traits::clamp(a, min.clone(), max.clone()))
191285
}
192286
}
287+
288+
#[cfg(all(test, feature = "std"))]
289+
mod angle_tests
290+
{
291+
use crate::Array;
292+
use num_complex::Complex;
293+
use std::f64::consts::PI;
294+
295+
/// Helper macro for floating-point comparison
296+
macro_rules! assert_approx_eq {
297+
($a:expr, $b:expr, $tol:expr $(, $msg:expr)?) => {{
298+
let (a, b) = ($a, $b);
299+
assert!(
300+
(a - b).abs() < $tol,
301+
concat!(
302+
"assertion failed: |left - right| >= tol\n",
303+
" left: {left:?}\n right: {right:?}\n tol: {tol:?}\n",
304+
$($msg,)?
305+
),
306+
left = a,
307+
right = b,
308+
tol = $tol
309+
);
310+
}};
311+
}
312+
313+
#[test]
314+
fn test_complex_numbers_radians()
315+
{
316+
let arr = Array::from_vec(vec![
317+
Complex::new(1.0f64, 0.0), // 0
318+
Complex::new(0.0, 1.0), // π/2
319+
Complex::new(-1.0, 0.0), // π
320+
Complex::new(0.0, -1.0), // -π/2
321+
Complex::new(1.0, 1.0), // π/4
322+
Complex::new(-1.0, -1.0), // -3π/4
323+
]);
324+
let a = arr.angle();
325+
326+
assert_approx_eq!(a[0], 0.0, 1e-10);
327+
assert_approx_eq!(a[1], PI / 2.0, 1e-10);
328+
assert_approx_eq!(a[2], PI, 1e-10);
329+
assert_approx_eq!(a[3], -PI / 2.0, 1e-10);
330+
assert_approx_eq!(a[4], PI / 4.0, 1e-10);
331+
assert_approx_eq!(a[5], -3.0 * PI / 4.0, 1e-10);
332+
}
333+
334+
#[test]
335+
fn test_complex_numbers_degrees()
336+
{
337+
let arr = Array::from_vec(vec![
338+
Complex::new(1.0f64, 0.0),
339+
Complex::new(0.0, 1.0),
340+
Complex::new(-1.0, 0.0),
341+
Complex::new(1.0, 1.0),
342+
]);
343+
let a = arr.angle().to_degrees();
344+
345+
assert_approx_eq!(a[0], 0.0, 1e-10);
346+
assert_approx_eq!(a[1], 90.0, 1e-10);
347+
assert_approx_eq!(a[2], 180.0, 1e-10);
348+
assert_approx_eq!(a[3], 45.0, 1e-10);
349+
}
350+
351+
#[test]
352+
fn test_signed_zeros()
353+
{
354+
let arr = Array::from_vec(vec![
355+
Complex::new(0.0f64, 0.0),
356+
Complex::new(-0.0, 0.0),
357+
Complex::new(0.0, -0.0),
358+
Complex::new(-0.0, -0.0),
359+
]);
360+
let a = arr.angle();
361+
362+
assert!(a[0] >= 0.0 && a[0].abs() < 1e-10);
363+
assert_approx_eq!(a[1], PI, 1e-10);
364+
assert!(a[2] <= 0.0 && a[2].abs() < 1e-10);
365+
assert_approx_eq!(a[3], -PI, 1e-10);
366+
}
367+
368+
#[test]
369+
fn test_edge_cases()
370+
{
371+
let arr = Array::from_vec(vec![
372+
Complex::new(f64::INFINITY, 0.0),
373+
Complex::new(0.0, f64::INFINITY),
374+
Complex::new(f64::NEG_INFINITY, 0.0),
375+
Complex::new(0.0, f64::NEG_INFINITY),
376+
]);
377+
let a = arr.angle();
378+
379+
assert_approx_eq!(a[0], 0.0, 1e-10);
380+
assert_approx_eq!(a[1], PI / 2.0, 1e-10);
381+
assert_approx_eq!(a[2], PI, 1e-10);
382+
assert_approx_eq!(a[3], -PI / 2.0, 1e-10);
383+
}
384+
385+
#[test]
386+
fn test_range_validation()
387+
{
388+
let n = 16;
389+
let complex_arr: Vec<_> = (0..n)
390+
.map(|i| {
391+
let theta = 2.0 * PI * (i as f64) / (n as f64);
392+
Complex::new(theta.cos(), theta.sin())
393+
})
394+
.collect();
395+
396+
let a = Array::from_vec(complex_arr).angle();
397+
398+
for &x in &a {
399+
assert!(x > -PI && x <= PI, "Angle {} outside (-π, π]", x);
400+
}
401+
}
402+
}

0 commit comments

Comments
 (0)