Skip to main content

sim_lib_numbers_signal/
transform.rs

1//! Plan execution over strided views and canonical tensor buffers.
2
3use sim_lib_numbers_tensor_cmplxf::ComplexFTensor;
4use sim_lib_numbers_tensor_f64::F64Tensor;
5
6use crate::{
7    Direction, LengthPolicy, Normalization, PlacementPolicy, SignalBuffer, SignalError, SignalView,
8    SignalViewMut, SpectrumPacking, TransformKind, TransformPlan,
9    fft::{Complex, fft},
10    reference::{reference_dct, reference_dft, reference_dst},
11};
12
13/// Executes an out-of-place transform into canonical tensor storage.
14pub fn transform(plan: &TransformPlan, input: SignalView<'_>) -> Result<SignalBuffer, SignalError> {
15    plan.validate()?;
16    if plan.placement != PlacementPolicy::OutOfPlace {
17        return Err(SignalError::InvalidPolicy {
18            policy: "placement",
19            reason: "transform requires OutOfPlace; use transform_in_place for InPlace",
20        });
21    }
22    match plan.kind {
23        TransformKind::Dft | TransformKind::Fft => {
24            let values = expect_complex(input)?;
25            let selected = select_complex(plan, values, plan.len)?;
26            let mut output = if plan.kind == TransformKind::Dft {
27                reference_dft(&selected, plan.direction, plan.sign)?
28            } else {
29                fft(
30                    &selected
31                        .iter()
32                        .copied()
33                        .map(Complex::from)
34                        .collect::<Vec<_>>(),
35                    plan.sign.angle_sign(plan.direction),
36                )?
37                .into_iter()
38                .map(Into::into)
39                .collect()
40            };
41            apply_fourier_normalization(&mut output, plan);
42            complex_buffer(output)
43        }
44        TransformKind::RealFft => transform_real_fft(plan, input),
45        TransformKind::Dct(kind) => {
46            let selected = select_real(plan, expect_real(input)?, plan.len)?;
47            real_buffer(reference_dct(
48                &selected,
49                kind,
50                plan.direction,
51                plan.normalization,
52            )?)
53        }
54        TransformKind::Dst(kind) => {
55            let selected = select_real(plan, expect_real(input)?, plan.len)?;
56            real_buffer(reference_dst(
57                &selected,
58                kind,
59                plan.direction,
60                plan.normalization,
61            )?)
62        }
63    }
64}
65
66/// Executes a same-representation transform into mutable caller storage.
67///
68/// Complex DFT/FFT and real DCT/DST plans are supported. Real FFT is rejected
69/// because its real/complex representations and packed lengths differ.
70pub fn transform_in_place(
71    plan: &TransformPlan,
72    input: SignalViewMut<'_>,
73) -> Result<(), SignalError> {
74    plan.validate()?;
75    if plan.placement != PlacementPolicy::InPlace {
76        return Err(SignalError::InvalidPolicy {
77            policy: "placement",
78            reason: "transform_in_place requires InPlace",
79        });
80    }
81    if plan.kind == TransformKind::RealFft {
82        return Err(SignalError::InvalidPolicy {
83            policy: "placement",
84            reason: "real FFT changes representation and may change packed length",
85        });
86    }
87    if plan.length == LengthPolicy::Pad {
88        return Err(SignalError::InvalidPolicy {
89            policy: "length",
90            reason: "in-place transforms cannot extend caller storage",
91        });
92    }
93    let mut out_of_place = plan.clone();
94    out_of_place.placement = PlacementPolicy::OutOfPlace;
95    let stride = plan.stride;
96    match input {
97        SignalViewMut::Complex(values) => {
98            let SignalBuffer::Complex(output) =
99                transform(&out_of_place, SignalView::Complex(values))?
100            else {
101                return Err(SignalError::InputKind {
102                    expected: "real",
103                    actual: "complex",
104                });
105            };
106            for (index, value) in output.as_slice().iter().copied().enumerate() {
107                values[stride.physical_index(index)?] = value;
108            }
109        }
110        SignalViewMut::Real(values) => {
111            let SignalBuffer::Real(output) = transform(&out_of_place, SignalView::Real(values))?
112            else {
113                return Err(SignalError::InputKind {
114                    expected: "complex",
115                    actual: "real",
116                });
117            };
118            for (index, value) in output.as_slice().iter().copied().enumerate() {
119                values[stride.physical_index(index)?] = value;
120            }
121        }
122    }
123    Ok(())
124}
125
126fn transform_real_fft(
127    plan: &TransformPlan,
128    input: SignalView<'_>,
129) -> Result<SignalBuffer, SignalError> {
130    match plan.direction {
131        Direction::Forward => {
132            let selected = select_real(plan, expect_real(input)?, plan.len)?;
133            let mut output = fft(
134                &selected
135                    .into_iter()
136                    .map(|value| Complex::new(value, 0.0))
137                    .collect::<Vec<_>>(),
138                plan.sign.angle_sign(Direction::Forward),
139            )?
140            .into_iter()
141            .map(Into::into)
142            .collect::<Vec<_>>();
143            apply_fourier_normalization(&mut output, plan);
144            if plan.packing == SpectrumPacking::HermitianHalf {
145                output.truncate(plan.len / 2 + 1);
146            }
147            complex_buffer(output)
148        }
149        Direction::Inverse => {
150            let packed_len = match plan.packing {
151                SpectrumPacking::Full => plan.len,
152                SpectrumPacking::HermitianHalf => plan.len / 2 + 1,
153            };
154            let selected = select_complex(plan, expect_complex(input)?, packed_len)?;
155            let spectrum = match plan.packing {
156                SpectrumPacking::Full => selected,
157                SpectrumPacking::HermitianHalf => unpack_hermitian(&selected, plan.len),
158            };
159            let mut output = fft(
160                &spectrum.into_iter().map(Complex::from).collect::<Vec<_>>(),
161                plan.sign.angle_sign(Direction::Inverse),
162            )?
163            .into_iter()
164            .map(Into::into)
165            .collect::<Vec<_>>();
166            apply_fourier_normalization(&mut output, plan);
167            let tolerance = 64.0 * f64::EPSILON * plan.len.max(1) as f64;
168            let real = output
169                .into_iter()
170                .enumerate()
171                .map(|(index, (real, imag))| {
172                    if imag.abs() > tolerance * real.abs().max(1.0) {
173                        Err(SignalError::InvalidPolicy {
174                            policy: "packing",
175                            reason: "inverse real FFT spectrum is not Hermitian",
176                        })
177                    } else if !real.is_finite() {
178                        Err(SignalError::NonFinite {
179                            index,
180                            component: "value",
181                        })
182                    } else {
183                        Ok(real)
184                    }
185                })
186                .collect::<Result<Vec<_>, _>>()?;
187            real_buffer(real)
188        }
189    }
190}
191
192fn unpack_hermitian(packed: &[(f64, f64)], len: usize) -> Vec<(f64, f64)> {
193    let mut output = vec![(0.0, 0.0); len];
194    output[..packed.len()].copy_from_slice(packed);
195    for (frequency, value) in output.iter_mut().enumerate().skip(packed.len()) {
196        let mirror = len - frequency;
197        *value = (packed[mirror].0, -packed[mirror].1);
198    }
199    output
200}
201
202fn apply_fourier_normalization(output: &mut [(f64, f64)], plan: &TransformPlan) {
203    let scale = match (plan.normalization, plan.direction) {
204        (Normalization::None, _) => 1.0,
205        (Normalization::Forward, Direction::Forward)
206        | (Normalization::Inverse, Direction::Inverse) => 1.0 / plan.len as f64,
207        (Normalization::Forward | Normalization::Inverse, _) => 1.0,
208        (Normalization::Orthonormal, _) => 1.0 / (plan.len as f64).sqrt(),
209    };
210    for (real, imag) in output {
211        *real *= scale;
212        *imag *= scale;
213    }
214}
215
216fn select_real(
217    plan: &TransformPlan,
218    values: &[f64],
219    expected: usize,
220) -> Result<Vec<f64>, SignalError> {
221    let available = plan.stride.available(values.len());
222    let selected_len = admitted_len(plan, expected, available)?;
223    let mut selected = Vec::with_capacity(expected);
224    for index in 0..selected_len {
225        let value = values[plan.stride.physical_index(index)?];
226        if !value.is_finite() {
227            return Err(SignalError::NonFinite {
228                index,
229                component: "value",
230            });
231        }
232        selected.push(value);
233    }
234    selected.resize(expected, 0.0);
235    Ok(selected)
236}
237
238fn select_complex(
239    plan: &TransformPlan,
240    values: &[(f64, f64)],
241    expected: usize,
242) -> Result<Vec<(f64, f64)>, SignalError> {
243    let available = plan.stride.available(values.len());
244    let selected_len = admitted_len(plan, expected, available)?;
245    let mut selected = Vec::with_capacity(expected);
246    for index in 0..selected_len {
247        let (real, imag) = values[plan.stride.physical_index(index)?];
248        if !real.is_finite() {
249            return Err(SignalError::NonFinite {
250                index,
251                component: "real",
252            });
253        }
254        if !imag.is_finite() {
255            return Err(SignalError::NonFinite {
256                index,
257                component: "imag",
258            });
259        }
260        selected.push((real, imag));
261    }
262    selected.resize(expected, (0.0, 0.0));
263    Ok(selected)
264}
265
266fn admitted_len(
267    plan: &TransformPlan,
268    expected: usize,
269    available: usize,
270) -> Result<usize, SignalError> {
271    match plan.length {
272        LengthPolicy::Exact if available == expected => Ok(expected),
273        LengthPolicy::Exact => Err(SignalError::LengthMismatch {
274            expected,
275            actual: available,
276        }),
277        LengthPolicy::Pad if available <= expected => Ok(available),
278        LengthPolicy::Pad => Err(SignalError::LengthMismatch {
279            expected,
280            actual: available,
281        }),
282        LengthPolicy::Truncate if available >= expected => Ok(expected),
283        LengthPolicy::Truncate => Err(SignalError::LengthMismatch {
284            expected,
285            actual: available,
286        }),
287    }
288}
289
290fn expect_complex(input: SignalView<'_>) -> Result<&[(f64, f64)], SignalError> {
291    match input {
292        SignalView::Complex(values) => Ok(values),
293        SignalView::Real(_) => Err(SignalError::InputKind {
294            expected: "complex",
295            actual: "real",
296        }),
297    }
298}
299
300fn expect_real(input: SignalView<'_>) -> Result<&[f64], SignalError> {
301    match input {
302        SignalView::Real(values) => Ok(values),
303        SignalView::Complex(_) => Err(SignalError::InputKind {
304            expected: "real",
305            actual: "complex",
306        }),
307    }
308}
309
310fn complex_buffer(values: Vec<(f64, f64)>) -> Result<SignalBuffer, SignalError> {
311    let len = values.len();
312    for (index, (real, imag)) in values.iter().copied().enumerate() {
313        if !real.is_finite() {
314            return Err(SignalError::NonFinite {
315                index,
316                component: "real",
317            });
318        }
319        if !imag.is_finite() {
320            return Err(SignalError::NonFinite {
321                index,
322                component: "imag",
323            });
324        }
325    }
326    ComplexFTensor::new(vec![len], values)
327        .map(SignalBuffer::Complex)
328        .ok_or(SignalError::InvalidLength {
329            len,
330            reason: "complex tensor shape overflowed",
331        })
332}
333
334fn real_buffer(values: Vec<f64>) -> Result<SignalBuffer, SignalError> {
335    let len = values.len();
336    for (index, value) in values.iter().enumerate() {
337        if !value.is_finite() {
338            return Err(SignalError::NonFinite {
339                index,
340                component: "value",
341            });
342        }
343    }
344    F64Tensor::new(vec![len], values)
345        .map(SignalBuffer::Real)
346        .ok_or(SignalError::InvalidLength {
347            len,
348            reason: "real tensor shape overflowed",
349        })
350}