Skip to main content

fpm_rs/
complex.rs

1//! Allocation-aware helpers for two-dimensional complex sample fields.
2//!
3//! Use [`crate::complex::amplitude`] and [`crate::complex::phase`] to derive real arrays,
4//! or [`crate::complex::from_amplitude_phase`] to construct a complex field. Inputs are borrowed
5//! [`ndarray`] views and results are newly allocated in standard row-major order.
6
7use ndarray::{Array2, ArrayView2, Zip};
8use num_complex::Complex64;
9
10use crate::{Error, Result};
11
12/// Double-precision complex scalar used for fields, spectra, and pupils.
13pub type Complex = Complex64;
14
15/// Computes amplitude from any logical two-dimensional layout.
16///
17/// The returned array is newly allocated in standard row-major order.
18pub fn amplitude(field: ArrayView2<'_, Complex64>) -> Array2<f64> {
19    field.mapv(|value| value.norm())
20}
21
22/// Computes phase from any logical two-dimensional layout.
23///
24/// The returned array is newly allocated in standard row-major order.
25pub fn phase(field: ArrayView2<'_, Complex64>) -> Array2<f64> {
26    field.mapv(|value| value.arg())
27}
28
29/// Combines amplitude and phase from arbitrary, matching ndarray layouts.
30///
31/// The inputs are borrowed without copying. The output allocation is standard
32/// row-major and follows their logical traversal order.
33pub fn from_amplitude_phase(
34    amplitude: ArrayView2<'_, f64>,
35    phase: ArrayView2<'_, f64>,
36) -> Result<Array2<Complex64>> {
37    if amplitude.dim() != phase.dim() {
38        return Err(Error::InvalidShape(format!(
39            "amplitude {:?} and phase {:?} differ",
40            amplitude.dim(),
41            phase.dim()
42        )));
43    }
44    let mut output = Array2::from_elem(amplitude.dim(), Complex64::default());
45    Zip::from(&mut output)
46        .and(amplitude)
47        .and(phase)
48        .for_each(|destination, &amplitude, &phase| {
49            *destination = Complex64::from_polar(amplitude, phase);
50        });
51    Ok(output)
52}