1use 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
13pub 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
66pub 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}