1use std::sync::Arc;
2
3use ad_core_rs::ndarray::{NDArray, NDDataBuffer, NDDataType, NDDimension};
4use ad_core_rs::ndarray_pool::NDArrayPool;
5use ad_core_rs::plugin::runtime::{NDPluginProcess, ProcessResult};
6use rustfft::FftPlanner;
7use rustfft::num_complex::Complex;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum FFTMode {
12 Rows1D,
13 Full2D,
14}
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum FFTDirection {
19 Forward,
20 Inverse,
21}
22
23pub struct FFTConfig {
25 pub mode: FFTMode,
26 pub direction: FFTDirection,
27 pub suppress_dc: bool,
29 pub num_average: usize,
31}
32
33impl Default for FFTConfig {
34 fn default() -> Self {
35 Self {
36 mode: FFTMode::Rows1D,
37 direction: FFTDirection::Forward,
38 suppress_dc: false,
39 num_average: 0,
40 }
41 }
42}
43
44pub fn next_pow2(n: usize) -> usize {
48 if n <= 1 {
49 return 1;
50 }
51 let mut p = 1usize;
52 while p < n {
53 p <<= 1;
54 }
55 p
56}
57
58pub fn fft_1d_rows(src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
64 if src.dims.is_empty() {
65 return None;
66 }
67
68 let width = src.dims[0].size;
69 let height = if src.dims.len() >= 2 {
70 src.dims[1].size
71 } else {
72 1
73 };
74
75 if width == 0 {
76 return None;
77 }
78
79 let padded = next_pow2(width);
81
82 let mut planner = FftPlanner::<f64>::new();
83 let fft = planner.plan_fft_forward(padded);
84
85 let n_freq = padded / 2;
87 if n_freq == 0 {
88 return None;
89 }
90 let scale = 1.0 / padded as f64;
91
92 let mut magnitudes = vec![0.0f64; n_freq * height];
93 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
94
95 for row in 0..height {
96 for c in row_buf.iter_mut() {
98 *c = Complex::new(0.0, 0.0);
99 }
100 for i in 0..width {
101 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
102 }
103
104 fft.process(&mut row_buf);
105
106 for i in 0..n_freq {
108 magnitudes[row * n_freq + i] = row_buf[i].norm() * scale;
109 }
110
111 if suppress_dc {
112 magnitudes[row * n_freq] = 0.0;
113 }
114 }
115
116 let dims = if height > 1 {
117 vec![NDDimension::new(n_freq), NDDimension::new(height)]
118 } else {
119 vec![NDDimension::new(n_freq)]
120 };
121 let mut arr = NDArray::new(dims, NDDataType::Float64);
122 arr.data = NDDataBuffer::F64(magnitudes);
123 arr.unique_id = src.unique_id;
124 arr.timestamp = src.timestamp;
125 arr.attributes = src.attributes.clone();
126 Some(arr)
127}
128
129pub fn fft_2d(src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
131 if src.dims.len() < 2 {
132 return None;
133 }
134
135 let src_w = src.dims[0].size;
136 let src_h = src.dims[1].size;
137
138 if src_w == 0 || src_h == 0 {
139 return None;
140 }
141
142 let w = next_pow2(src_w);
144 let h = next_pow2(src_h);
145
146 let mut planner = FftPlanner::<f64>::new();
147 let fft_row = planner.plan_fft_forward(w);
148 let fft_col = planner.plan_fft_forward(h);
149
150 let mut data = vec![Complex::new(0.0, 0.0); w * h];
152 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
153
154 for row in 0..src_h {
155 for c in row_buf.iter_mut() {
156 *c = Complex::new(0.0, 0.0);
157 }
158 for i in 0..src_w {
159 row_buf[i] = Complex::new(src.data.get_as_f64(row * src_w + i).unwrap_or(0.0), 0.0);
160 }
161 fft_row.process(&mut row_buf);
162 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
163 }
164
165 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
167
168 for col in 0..w {
169 for row in 0..h {
171 col_buf[row] = data[row * w + col];
172 }
173 fft_col.process(&mut col_buf);
174 for row in 0..h {
176 data[row * w + col] = col_buf[row];
177 }
178 }
179
180 let n_freq_x = w / 2;
182 let n_freq_y = h / 2;
183 if n_freq_x == 0 || n_freq_y == 0 {
184 return None;
185 }
186 let scale = 1.0 / (w * h) as f64;
187
188 let mut magnitudes = vec![0.0f64; n_freq_x * n_freq_y];
189 for fy in 0..n_freq_y {
190 for fx in 0..n_freq_x {
191 magnitudes[fy * n_freq_x + fx] = data[fy * w + fx].norm() * scale;
192 }
193 }
194
195 if suppress_dc {
196 magnitudes[0] = 0.0;
197 }
198
199 let dims = vec![NDDimension::new(n_freq_x), NDDimension::new(n_freq_y)];
200 let mut arr = NDArray::new(dims, NDDataType::Float64);
201 arr.data = NDDataBuffer::F64(magnitudes);
202 arr.unique_id = src.unique_id;
203 arr.timestamp = src.timestamp;
204 arr.attributes = src.attributes.clone();
205 Some(arr)
206}
207
208#[derive(Default)]
210struct FFTParamIndices {
211 direction: Option<usize>,
212 suppress_dc: Option<usize>,
213 num_average: Option<usize>,
214 num_averaged: Option<usize>,
215 reset_average: Option<usize>,
216 time_per_point: Option<usize>,
217 time_series: Option<usize>,
219 real: Option<usize>,
221 imaginary: Option<usize>,
223 abs_value: Option<usize>,
225 time_axis: Option<usize>,
227 freq_axis: Option<usize>,
229}
230
231pub struct FFTProcessor {
232 config: FFTConfig,
233 planner: FftPlanner<f64>,
234 avg_buffer: Option<Vec<f64>>,
236 avg_count: usize,
238 cached_dims: Vec<usize>,
240 time_per_point: f64,
243 params: FFTParamIndices,
244}
245
246impl FFTProcessor {
247 pub fn new(mode: FFTMode) -> Self {
248 Self {
249 config: FFTConfig {
250 mode,
251 direction: FFTDirection::Forward,
252 suppress_dc: false,
253 num_average: 0,
254 },
255 planner: FftPlanner::new(),
256 avg_buffer: None,
257 avg_count: 0,
258 cached_dims: Vec::new(),
259 time_per_point: 1.0,
260 params: FFTParamIndices::default(),
261 }
262 }
263
264 pub fn with_config(config: FFTConfig) -> Self {
265 Self {
266 config,
267 planner: FftPlanner::new(),
268 avg_buffer: None,
269 avg_count: 0,
270 cached_dims: Vec::new(),
271 time_per_point: 1.0,
272 params: FFTParamIndices::default(),
273 }
274 }
275
276 fn check_dims_changed(&mut self, dims: &[NDDimension]) {
278 let current: Vec<usize> = dims.iter().map(|d| d.size).collect();
279 if current != self.cached_dims {
280 self.cached_dims = current;
281 self.avg_buffer = None;
282 self.avg_count = 0;
283 }
284 }
285
286 fn compute_fft(&mut self, src: &NDArray) -> Option<NDArray> {
288 let suppress_dc = self.config.suppress_dc;
289
290 match (self.config.mode, self.config.direction) {
291 (FFTMode::Rows1D, FFTDirection::Forward) => {
292 self.compute_fft_1d_rows_forward(src, suppress_dc)
293 }
294 (FFTMode::Rows1D, FFTDirection::Inverse) => {
295 self.compute_fft_1d_rows_inverse(src, suppress_dc)
296 }
297 (FFTMode::Full2D, FFTDirection::Forward) => {
298 self.compute_fft_2d_forward(src, suppress_dc)
299 }
300 (FFTMode::Full2D, FFTDirection::Inverse) => {
301 self.compute_fft_2d_inverse(src, suppress_dc)
302 }
303 }
304 }
305
306 fn compute_row_spectrum(
315 &mut self,
316 src: &NDArray,
317 suppress_dc: bool,
318 ) -> Option<(Vec<f64>, Vec<f64>, Vec<f64>)> {
319 if src.dims.is_empty() {
320 return None;
321 }
322 let width = src.dims[0].size;
323 if width == 0 {
324 return None;
325 }
326 let padded = next_pow2(width);
327 let n_freq = padded / 2;
328 if n_freq == 0 {
329 return None;
330 }
331 let fft = self.planner.plan_fft_forward(padded);
332
333 let time_series: Vec<f64> = (0..width)
335 .map(|i| src.data.get_as_f64(i).unwrap_or(0.0))
336 .collect();
337
338 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
339 for (i, &v) in time_series.iter().enumerate() {
340 row_buf[i] = Complex::new(v, 0.0);
341 }
342 fft.process(&mut row_buf);
343
344 let mut real = vec![0.0f64; n_freq];
345 let mut imag = vec![0.0f64; n_freq];
346 for i in 0..n_freq {
347 real[i] = row_buf[i].re;
348 imag[i] = row_buf[i].im;
349 }
350 if suppress_dc {
351 real[0] = 0.0;
352 imag[0] = 0.0;
353 }
354 Some((time_series, real, imag))
355 }
356
357 fn freq_axis(&self, n_freq: usize) -> Vec<f64> {
360 if n_freq <= 1 {
361 return vec![0.0; n_freq];
362 }
363 let tpp = if self.time_per_point > 0.0 {
364 self.time_per_point
365 } else {
366 1.0
367 };
368 let step = 0.5 / tpp / (n_freq - 1) as f64;
369 (0..n_freq).map(|i| i as f64 * step).collect()
370 }
371
372 fn time_axis(&self, n_time: usize) -> Vec<f64> {
374 let tpp = if self.time_per_point > 0.0 {
375 self.time_per_point
376 } else {
377 1.0
378 };
379 (0..n_time).map(|i| i as f64 * tpp).collect()
380 }
381
382 fn compute_fft_1d_rows_forward(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
383 if src.dims.is_empty() {
384 return None;
385 }
386
387 let width = src.dims[0].size;
388 let height = if src.dims.len() >= 2 {
389 src.dims[1].size
390 } else {
391 1
392 };
393
394 if width == 0 {
395 return None;
396 }
397
398 let padded = next_pow2(width);
400 let fft = self.planner.plan_fft_forward(padded);
401
402 let n_freq = padded / 2;
404 if n_freq == 0 {
405 return None;
406 }
407 let scale = 1.0 / padded as f64;
408
409 let mut magnitudes = vec![0.0f64; n_freq * height];
410 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
411
412 for row in 0..height {
413 for c in row_buf.iter_mut() {
414 *c = Complex::new(0.0, 0.0);
415 }
416 for i in 0..width {
417 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
418 }
419 fft.process(&mut row_buf);
420 for i in 0..n_freq {
421 magnitudes[row * n_freq + i] = row_buf[i].norm() * scale;
422 }
423 if suppress_dc {
424 magnitudes[row * n_freq] = 0.0;
425 }
426 }
427
428 let dims = if height > 1 {
429 vec![NDDimension::new(n_freq), NDDimension::new(height)]
430 } else {
431 vec![NDDimension::new(n_freq)]
432 };
433 let mut arr = NDArray::new(dims, NDDataType::Float64);
434 arr.data = NDDataBuffer::F64(magnitudes);
435 arr.unique_id = src.unique_id;
436 arr.timestamp = src.timestamp;
437 arr.attributes = src.attributes.clone();
438 Some(arr)
439 }
440
441 fn compute_fft_1d_rows_inverse(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
442 if src.dims.is_empty() {
443 return None;
444 }
445
446 let width = src.dims[0].size;
447 let height = if src.dims.len() >= 2 {
448 src.dims[1].size
449 } else {
450 1
451 };
452
453 if width == 0 {
454 return None;
455 }
456
457 let fft = self.planner.plan_fft_inverse(width);
458 let scale = 1.0 / width as f64;
459
460 let mut samples = vec![0.0f64; width * height];
464 let mut row_buf = vec![Complex::new(0.0, 0.0); width];
465
466 for row in 0..height {
467 for i in 0..width {
468 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
469 }
470 if suppress_dc {
471 row_buf[0] = Complex::new(0.0, 0.0);
472 }
473 fft.process(&mut row_buf);
474 for (i, c) in row_buf.iter().enumerate() {
475 samples[row * width + i] = c.re * scale;
476 }
477 }
478
479 let dims = src.dims.clone();
480 let mut arr = NDArray::new(dims, NDDataType::Float64);
481 arr.data = NDDataBuffer::F64(samples);
482 arr.unique_id = src.unique_id;
483 arr.timestamp = src.timestamp;
484 arr.attributes = src.attributes.clone();
485 Some(arr)
486 }
487
488 fn compute_fft_2d_forward(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
489 if src.dims.len() < 2 {
490 return None;
491 }
492
493 let src_w = src.dims[0].size;
494 let src_h = src.dims[1].size;
495
496 if src_w == 0 || src_h == 0 {
497 return None;
498 }
499
500 let w = next_pow2(src_w);
502 let h = next_pow2(src_h);
503
504 let fft_row = self.planner.plan_fft_forward(w);
505 let fft_col = self.planner.plan_fft_forward(h);
506
507 let mut data = vec![Complex::new(0.0, 0.0); w * h];
508 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
509
510 for row in 0..src_h {
511 for c in row_buf.iter_mut() {
512 *c = Complex::new(0.0, 0.0);
513 }
514 for i in 0..src_w {
515 row_buf[i] = Complex::new(src.data.get_as_f64(row * src_w + i).unwrap_or(0.0), 0.0);
516 }
517 fft_row.process(&mut row_buf);
518 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
519 }
520
521 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
522 for col in 0..w {
523 for row in 0..h {
524 col_buf[row] = data[row * w + col];
525 }
526 fft_col.process(&mut col_buf);
527 for row in 0..h {
528 data[row * w + col] = col_buf[row];
529 }
530 }
531
532 let n_freq_x = w / 2;
534 let n_freq_y = h / 2;
535 if n_freq_x == 0 || n_freq_y == 0 {
536 return None;
537 }
538 let scale = 1.0 / (w * h) as f64;
539
540 let mut magnitudes = vec![0.0f64; n_freq_x * n_freq_y];
541 for fy in 0..n_freq_y {
542 for fx in 0..n_freq_x {
543 magnitudes[fy * n_freq_x + fx] = data[fy * w + fx].norm() * scale;
544 }
545 }
546
547 if suppress_dc {
548 magnitudes[0] = 0.0;
549 }
550
551 let dims = vec![NDDimension::new(n_freq_x), NDDimension::new(n_freq_y)];
552 let mut arr = NDArray::new(dims, NDDataType::Float64);
553 arr.data = NDDataBuffer::F64(magnitudes);
554 arr.unique_id = src.unique_id;
555 arr.timestamp = src.timestamp;
556 arr.attributes = src.attributes.clone();
557 Some(arr)
558 }
559
560 fn compute_fft_2d_inverse(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
561 if src.dims.len() < 2 {
562 return None;
563 }
564
565 let w = src.dims[0].size;
566 let h = src.dims[1].size;
567
568 if w == 0 || h == 0 {
569 return None;
570 }
571
572 let fft_row = self.planner.plan_fft_inverse(w);
573 let fft_col = self.planner.plan_fft_inverse(h);
574 let scale = 1.0 / (w * h) as f64;
575
576 let mut data = vec![Complex::new(0.0, 0.0); w * h];
577 for i in 0..w * h {
578 data[i] = Complex::new(src.data.get_as_f64(i).unwrap_or(0.0), 0.0);
579 }
580
581 if suppress_dc {
582 data[0] = Complex::new(0.0, 0.0);
583 }
584
585 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
586 for col in 0..w {
587 for row in 0..h {
588 col_buf[row] = data[row * w + col];
589 }
590 fft_col.process(&mut col_buf);
591 for row in 0..h {
592 data[row * w + col] = col_buf[row];
593 }
594 }
595
596 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
597 for row in 0..h {
598 row_buf.copy_from_slice(&data[row * w..(row * w + w)]);
599 fft_row.process(&mut row_buf);
600 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
601 }
602
603 let samples: Vec<f64> = data.iter().map(|c| c.re * scale).collect();
605
606 let dims = vec![NDDimension::new(w), NDDimension::new(h)];
607 let mut arr = NDArray::new(dims, NDDataType::Float64);
608 arr.data = NDDataBuffer::F64(samples);
609 arr.unique_id = src.unique_id;
610 arr.timestamp = src.timestamp;
611 arr.attributes = src.attributes.clone();
612 Some(arr)
613 }
614
615 fn apply_averaging(&mut self, magnitudes: &[f64]) -> Vec<f64> {
620 let num_avg = self.config.num_average;
621 if num_avg <= 1 {
622 return magnitudes.to_vec();
623 }
624
625 let buf = self
626 .avg_buffer
627 .get_or_insert_with(|| vec![0.0; magnitudes.len()]);
628
629 if buf.len() != magnitudes.len() {
631 *buf = vec![0.0; magnitudes.len()];
632 self.avg_count = 0;
633 }
634
635 self.avg_count += 1;
636 let n = self.avg_count.min(num_avg) as f64;
638 let new_fraction = 1.0 / n;
639 let old_fraction = 1.0 - new_fraction;
640
641 for (b, &m) in buf.iter_mut().zip(magnitudes.iter()) {
643 *b = *b * old_fraction + m * new_fraction;
644 }
645
646 buf.clone()
647 }
648}
649
650impl NDPluginProcess for FFTProcessor {
651 fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
652 use ad_core_rs::plugin::runtime::ParamUpdate;
653
654 self.check_dims_changed(&array.dims);
655
656 let result = self.compute_fft(array);
657 let mut updates = Vec::new();
658 if let Some(idx) = self.params.num_averaged {
659 updates.push(ParamUpdate::int32(idx, self.avg_count as i32));
660 }
661
662 let mut averaged_mags: Option<Vec<f64>> = None;
671 if self.config.direction == FFTDirection::Forward {
672 let suppress_dc = self.config.suppress_dc;
673 if let Some((time_series, real, imag)) = self.compute_row_spectrum(array, suppress_dc) {
674 let n_time = time_series.len();
675 let n_freq = real.len();
676 if let Some(idx) = self.params.time_series {
677 updates.push(ParamUpdate::float64_array(idx, time_series));
678 }
679 if let Some(idx) = self.params.real {
680 updates.push(ParamUpdate::float64_array(idx, real));
681 }
682 if let Some(idx) = self.params.imaginary {
683 updates.push(ParamUpdate::float64_array(idx, imag));
684 }
685 if let Some(idx) = self.params.time_axis {
686 updates.push(ParamUpdate::float64_array(idx, self.time_axis(n_time)));
687 }
688 if let Some(idx) = self.params.freq_axis {
689 updates.push(ParamUpdate::float64_array(idx, self.freq_axis(n_freq)));
690 }
691 }
692 }
693
694 match result {
695 Some(mut out) => {
696 if self.config.num_average > 1 {
697 if let NDDataBuffer::F64(ref mags) = out.data {
698 let averaged = self.apply_averaging(mags);
699 averaged_mags = Some(averaged.clone());
700 out.data = NDDataBuffer::F64(averaged);
701 }
702 }
703 if self.config.direction == FFTDirection::Forward {
707 if let Some(idx) = self.params.abs_value {
708 let abs = match (&averaged_mags, &out.data) {
709 (Some(avg), _) => avg.clone(),
710 (None, NDDataBuffer::F64(mags)) => mags.clone(),
711 _ => Vec::new(),
712 };
713 if !abs.is_empty() {
714 updates.push(ParamUpdate::float64_array(idx, abs));
715 }
716 }
717 }
718 let mut r = ProcessResult::arrays(vec![Arc::new(out)]);
719 r.param_updates = updates;
720 r
721 }
722 None => ProcessResult::sink(updates),
723 }
724 }
725
726 fn plugin_type(&self) -> &str {
727 "NDPluginFFT"
728 }
729
730 fn register_params(
731 &mut self,
732 base: &mut asyn_rs::port::PortDriverBase,
733 ) -> asyn_rs::error::AsynResult<()> {
734 use asyn_rs::param::ParamType;
735 base.create_param("FFT_TIME_PER_POINT", ParamType::Float64)?;
736 base.create_param("FFT_TIME_AXIS", ParamType::Float64Array)?;
737 base.create_param("FFT_FREQ_AXIS", ParamType::Float64Array)?;
738 base.create_param("FFT_DIRECTION", ParamType::Int32)?;
739 base.create_param("FFT_SUPPRESS_DC", ParamType::Int32)?;
740 base.create_param("FFT_NUM_AVERAGE", ParamType::Int32)?;
741 base.create_param("FFT_NUM_AVERAGED", ParamType::Int32)?;
742 base.create_param("FFT_RESET_AVERAGE", ParamType::Int32)?;
743 base.create_param("FFT_TIME_SERIES", ParamType::Float64Array)?;
744 base.create_param("FFT_REAL", ParamType::Float64Array)?;
745 base.create_param("FFT_IMAGINARY", ParamType::Float64Array)?;
746 base.create_param("FFT_ABS_VALUE", ParamType::Float64Array)?;
747
748 self.params.direction = base.find_param("FFT_DIRECTION");
749 self.params.suppress_dc = base.find_param("FFT_SUPPRESS_DC");
750 self.params.num_average = base.find_param("FFT_NUM_AVERAGE");
751 self.params.num_averaged = base.find_param("FFT_NUM_AVERAGED");
752 self.params.reset_average = base.find_param("FFT_RESET_AVERAGE");
753 self.params.time_per_point = base.find_param("FFT_TIME_PER_POINT");
754 self.params.time_series = base.find_param("FFT_TIME_SERIES");
755 self.params.real = base.find_param("FFT_REAL");
756 self.params.imaginary = base.find_param("FFT_IMAGINARY");
757 self.params.abs_value = base.find_param("FFT_ABS_VALUE");
758 self.params.time_axis = base.find_param("FFT_TIME_AXIS");
759 self.params.freq_axis = base.find_param("FFT_FREQ_AXIS");
760 Ok(())
761 }
762
763 fn on_param_change(
764 &mut self,
765 reason: usize,
766 params: &ad_core_rs::plugin::runtime::PluginParamSnapshot,
767 ) -> ad_core_rs::plugin::runtime::ParamChangeResult {
768 if Some(reason) == self.params.direction {
769 self.config.direction = if params.value.as_i32() == 0 {
770 FFTDirection::Forward
771 } else {
772 FFTDirection::Inverse
773 };
774 } else if Some(reason) == self.params.suppress_dc {
775 self.config.suppress_dc = params.value.as_i32() != 0;
776 } else if Some(reason) == self.params.num_average {
777 self.config.num_average = params.value.as_i32().max(0) as usize;
778 } else if Some(reason) == self.params.reset_average {
779 if params.value.as_i32() != 0 {
780 self.avg_buffer = None;
781 self.avg_count = 0;
782 }
783 } else if Some(reason) == self.params.time_per_point {
784 let v = params.value.as_f64();
786 if v > 0.0 {
787 self.time_per_point = v;
788 }
789 }
790 ad_core_rs::plugin::runtime::ParamChangeResult::updates(vec![])
791 }
792}
793
794#[cfg(test)]
795mod tests {
796 use super::*;
797
798 #[test]
799 fn test_fft_1d_dc() {
800 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
802 if let NDDataBuffer::F64(ref mut v) = arr.data {
803 for i in 0..8 {
804 v[i] = 1.0;
805 }
806 }
807
808 let result = fft_1d_rows(&arr, false).unwrap();
809 assert_eq!(result.dims[0].size, 4);
811 if let NDDataBuffer::F64(ref v) = result.data {
812 assert!((v[0] - 1.0).abs() < 1e-10);
814 assert!(v[1].abs() < 1e-10);
816 }
817 }
818
819 #[test]
820 fn test_fft_1d_sine() {
821 let n = 16;
823 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
824 if let NDDataBuffer::F64(ref mut v) = arr.data {
825 for i in 0..n {
826 v[i] = (2.0 * std::f64::consts::PI * i as f64 / n as f64).sin();
827 }
828 }
829
830 let result = fft_1d_rows(&arr, false).unwrap();
831 assert_eq!(result.dims[0].size, 8);
833 if let NDDataBuffer::F64(ref v) = result.data {
834 assert!(v[0].abs() < 1e-10);
836 assert!((v[1] - 0.5).abs() < 1e-10);
838 assert!(v[2].abs() < 1e-10);
840 }
841 }
842
843 #[test]
844 fn test_fft_2d_dimensions() {
845 let arr = NDArray::new(
846 vec![NDDimension::new(4), NDDimension::new(4)],
847 NDDataType::UInt8,
848 );
849 let result = fft_2d(&arr, false).unwrap();
850 assert_eq!(result.dims[0].size, 2);
852 assert_eq!(result.dims[1].size, 2);
853 assert_eq!(result.data.data_type(), NDDataType::Float64);
854 }
855
856 #[test]
857 fn test_fft_1d_suppress_dc() {
858 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
860 if let NDDataBuffer::F64(ref mut v) = arr.data {
861 for i in 0..8 {
862 v[i] = 1.0;
863 }
864 }
865
866 let result = fft_1d_rows(&arr, true).unwrap();
867 if let NDDataBuffer::F64(ref v) = result.data {
868 assert!((v[0]).abs() < 1e-15);
870 assert!(v[1].abs() < 1e-10);
872 } else {
873 panic!("expected F64 data");
874 }
875 }
876
877 #[test]
878 fn test_fft_2d_suppress_dc() {
879 let mut arr = NDArray::new(
881 vec![NDDimension::new(4), NDDimension::new(4)],
882 NDDataType::Float64,
883 );
884 if let NDDataBuffer::F64(ref mut v) = arr.data {
885 for val in v.iter_mut() {
886 *val = 3.0;
887 }
888 }
889
890 let result = fft_2d(&arr, true).unwrap();
891 if let NDDataBuffer::F64(ref v) = result.data {
892 assert!((v[0]).abs() < 1e-15);
894 } else {
895 panic!("expected F64 data");
896 }
897 }
898
899 #[test]
900 fn test_fft_2d_known_dc() {
901 let mut arr = NDArray::new(
903 vec![NDDimension::new(4), NDDimension::new(4)],
904 NDDataType::Float64,
905 );
906 if let NDDataBuffer::F64(ref mut v) = arr.data {
907 for val in v.iter_mut() {
908 *val = 2.0;
909 }
910 }
911
912 let result = fft_2d(&arr, false).unwrap();
913 assert_eq!(result.dims[0].size, 2);
915 assert_eq!(result.dims[1].size, 2);
916 if let NDDataBuffer::F64(ref v) = result.data {
917 assert!((v[0] - 2.0).abs() < 1e-10, "DC = {}, expected 2", v[0]);
919 for i in 1..v.len() {
921 assert!(v[i].abs() < 1e-10, "bin {} = {}, expected ~0", i, v[i]);
922 }
923 } else {
924 panic!("expected F64 data");
925 }
926 }
927
928 #[test]
929 fn test_fft_1d_known_cosine_peaks() {
930 let n = 16;
932 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
933 if let NDDataBuffer::F64(ref mut v) = arr.data {
934 for i in 0..n {
935 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
936 }
937 }
938
939 let result = fft_1d_rows(&arr, false).unwrap();
940 assert_eq!(result.dims[0].size, 8);
942 if let NDDataBuffer::F64(ref v) = result.data {
943 assert!(v[0].abs() < 1e-10);
945 assert!(
947 (v[3] - 0.5).abs() < 1e-10,
948 "k=3 magnitude = {}, expected 0.5",
949 v[3]
950 );
951 for k in [1, 2, 4, 5, 6, 7] {
953 assert!(
954 v[k].abs() < 1e-10,
955 "k={} magnitude = {}, expected ~0",
956 k,
957 v[k]
958 );
959 }
960 } else {
961 panic!("expected F64 data");
962 }
963 }
964
965 #[test]
966 fn test_processor_with_config() {
967 let config = FFTConfig {
968 mode: FFTMode::Rows1D,
969 direction: FFTDirection::Forward,
970 suppress_dc: true,
971 num_average: 0,
972 };
973 let mut proc = FFTProcessor::with_config(config);
974 let pool = NDArrayPool::new(0);
975
976 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
977 if let NDDataBuffer::F64(ref mut v) = arr.data {
978 for i in 0..8 {
979 v[i] = 5.0;
980 }
981 }
982
983 let result = proc.process_array(&arr, &pool);
984 assert_eq!(result.output_arrays.len(), 1);
985 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
986 assert!(v[0].abs() < 1e-15);
988 } else {
989 panic!("expected F64 data");
990 }
991 }
992
993 #[test]
994 fn test_processor_averaging() {
995 let config = FFTConfig {
996 mode: FFTMode::Rows1D,
997 direction: FFTDirection::Forward,
998 suppress_dc: false,
999 num_average: 2,
1000 };
1001 let mut proc = FFTProcessor::with_config(config);
1002 let pool = NDArrayPool::new(0);
1003
1004 let mut arr1 = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1006 if let NDDataBuffer::F64(ref mut v) = arr1.data {
1007 for i in 0..8 {
1008 v[i] = 2.0;
1009 }
1010 }
1011
1012 let mut arr2 = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1014 if let NDDataBuffer::F64(ref mut v) = arr2.data {
1015 for i in 0..8 {
1016 v[i] = 4.0;
1017 }
1018 }
1019
1020 let r1 = proc.process_array(&arr1, &pool);
1021 assert_eq!(r1.output_arrays.len(), 1);
1022 if let NDDataBuffer::F64(ref v) = r1.output_arrays[0].data {
1024 assert!((v[0] - 2.0).abs() < 1e-10, "partial avg DC = {}", v[0]);
1025 }
1026
1027 let r2 = proc.process_array(&arr2, &pool);
1028 assert_eq!(r2.output_arrays.len(), 1);
1029 if let NDDataBuffer::F64(ref v) = r2.output_arrays[0].data {
1031 assert!((v[0] - 3.0).abs() < 1e-10, "averaged DC = {}", v[0]);
1032 }
1033 }
1034
1035 #[test]
1036 fn test_processor_averaging_dimension_change_resets() {
1037 let config = FFTConfig {
1038 mode: FFTMode::Rows1D,
1039 direction: FFTDirection::Forward,
1040 suppress_dc: false,
1041 num_average: 3,
1042 };
1043 let mut proc = FFTProcessor::with_config(config);
1044 let pool = NDArrayPool::new(0);
1045
1046 let mut arr1 = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1048 if let NDDataBuffer::F64(ref mut v) = arr1.data {
1049 for i in 0..8 {
1050 v[i] = 1.0;
1051 }
1052 }
1053 let _ = proc.process_array(&arr1, &pool);
1054 assert_eq!(proc.avg_count, 1);
1055
1056 let mut arr2 = NDArray::new(vec![NDDimension::new(4)], NDDataType::Float64);
1058 if let NDDataBuffer::F64(ref mut v) = arr2.data {
1059 for i in 0..4 {
1060 v[i] = 1.0;
1061 }
1062 }
1063 let _ = proc.process_array(&arr2, &pool);
1064 assert_eq!(proc.avg_count, 1);
1066 }
1067
1068 #[test]
1069 fn test_fft_1d_multirow() {
1070 let w = 4;
1072 let h = 2;
1073 let mut arr = NDArray::new(
1074 vec![NDDimension::new(w), NDDimension::new(h)],
1075 NDDataType::Float64,
1076 );
1077 if let NDDataBuffer::F64(ref mut v) = arr.data {
1078 for i in 0..w {
1080 v[i] = 1.0;
1081 }
1082 for i in w..2 * w {
1084 v[i] = 3.0;
1085 }
1086 }
1087
1088 let result = fft_1d_rows(&arr, false).unwrap();
1089 let n_freq = w / 2; assert_eq!(result.dims[0].size, n_freq);
1091 if let NDDataBuffer::F64(ref v) = result.data {
1092 assert!((v[0] - 1.0).abs() < 1e-10);
1094 assert!((v[n_freq] - 3.0).abs() < 1e-10);
1096 } else {
1097 panic!("expected F64 data");
1098 }
1099 }
1100
1101 #[test]
1102 fn test_inverse_fft_1d() {
1103 let n = 8;
1107 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1108 if let NDDataBuffer::F64(ref mut v) = arr.data {
1109 v[0] = 8.0; }
1112
1113 let config = FFTConfig {
1114 mode: FFTMode::Rows1D,
1115 direction: FFTDirection::Inverse,
1116 suppress_dc: false,
1117 num_average: 0,
1118 };
1119 let mut proc = FFTProcessor::with_config(config);
1120 let pool = NDArrayPool::new(0);
1121
1122 let result = proc.process_array(&arr, &pool);
1123 assert_eq!(result.output_arrays.len(), 1);
1124 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
1125 for i in 0..n {
1127 assert!(
1128 (v[i] - 1.0).abs() < 1e-10,
1129 "sample {} = {}, expected 1.0",
1130 i,
1131 v[i]
1132 );
1133 }
1134 } else {
1135 panic!("expected F64 data");
1136 }
1137 }
1138
1139 #[test]
1140 fn test_fft_preserves_metadata() {
1141 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::Float64);
1142 arr.unique_id = 42;
1143 if let NDDataBuffer::F64(ref mut v) = arr.data {
1144 v[0] = 1.0;
1145 }
1146
1147 let result = fft_1d_rows(&arr, false).unwrap();
1148 assert_eq!(result.unique_id, 42);
1149 assert_eq!(result.timestamp, arr.timestamp);
1150 }
1151
1152 #[test]
1153 fn test_next_pow2() {
1154 assert_eq!(next_pow2(0), 1);
1155 assert_eq!(next_pow2(1), 1);
1156 assert_eq!(next_pow2(2), 2);
1157 assert_eq!(next_pow2(3), 4);
1158 assert_eq!(next_pow2(5), 8);
1159 assert_eq!(next_pow2(8), 8);
1160 assert_eq!(next_pow2(100), 128);
1161 }
1162
1163 #[test]
1164 fn test_fft_1d_pads_to_power_of_two() {
1165 let n = 5; let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1169 if let NDDataBuffer::F64(ref mut v) = arr.data {
1170 for i in 0..n {
1171 v[i] = 1.0;
1172 }
1173 }
1174 let result = fft_1d_rows(&arr, false).unwrap();
1175 assert_eq!(result.dims[0].size, 4); }
1177
1178 #[test]
1179 fn test_fft_2d_pads_to_power_of_two() {
1180 let arr = NDArray::new(
1182 vec![NDDimension::new(6), NDDimension::new(3)],
1183 NDDataType::Float64,
1184 );
1185 let result = fft_2d(&arr, false).unwrap();
1186 assert_eq!(result.dims[0].size, 4); assert_eq!(result.dims[1].size, 2); }
1189
1190 use ad_core_rs::plugin::runtime::ParamUpdate;
1193
1194 fn fft_proc_with_params(config: FFTConfig) -> FFTProcessor {
1196 let mut proc = FFTProcessor::with_config(config);
1197 let mut base =
1198 asyn_rs::port::PortDriverBase::new("FFT_TEST", 1, asyn_rs::port::PortFlags::default());
1199 proc.register_params(&mut base).unwrap();
1200 proc
1201 }
1202
1203 fn find_array_update(updates: &[ParamUpdate], reason: usize) -> Option<&[f64]> {
1205 updates.iter().find_map(|u| match u {
1206 ParamUpdate::Float64Array {
1207 reason: r, value, ..
1208 } if *r == reason => Some(value.as_slice()),
1209 _ => None,
1210 })
1211 }
1212
1213 #[test]
1214 fn test_fft_emits_all_waveforms() {
1215 let mut proc = fft_proc_with_params(FFTConfig::default());
1218 let pool = NDArrayPool::new(0);
1219
1220 let n = 16;
1221 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1222 if let NDDataBuffer::F64(ref mut v) = arr.data {
1223 for i in 0..n {
1224 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
1225 }
1226 }
1227 let result = proc.process_array(&arr, &pool);
1228 let u = &result.param_updates;
1229
1230 for reason in [
1233 proc.params.time_series.unwrap(),
1234 proc.params.real.unwrap(),
1235 proc.params.imaginary.unwrap(),
1236 proc.params.abs_value.unwrap(),
1237 proc.params.time_axis.unwrap(),
1238 proc.params.freq_axis.unwrap(),
1239 ] {
1240 let wf = find_array_update(u, reason)
1241 .unwrap_or_else(|| panic!("missing waveform for reason {reason}"));
1242 assert!(!wf.is_empty(), "waveform {reason} is empty");
1243 }
1244 let array_updates = u
1245 .iter()
1246 .filter(|x| matches!(x, ParamUpdate::Float64Array { .. }))
1247 .count();
1248 assert_eq!(
1249 array_updates, 6,
1250 "expected 6 waveform updates, got {array_updates}"
1251 );
1252 }
1253
1254 #[test]
1255 fn test_fft_real_imaginary_match_spectrum() {
1256 let mut proc = fft_proc_with_params(FFTConfig::default());
1259 let real_reason = proc.params.real.unwrap();
1260 let imag_reason = proc.params.imaginary.unwrap();
1261 let abs_reason = proc.params.abs_value.unwrap();
1262 let ts_reason = proc.params.time_series.unwrap();
1263 let pool = NDArrayPool::new(0);
1264
1265 let n = 16;
1266 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1267 if let NDDataBuffer::F64(ref mut v) = arr.data {
1268 for i in 0..n {
1269 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
1270 }
1271 }
1272 let result = proc.process_array(&arr, &pool);
1273 let u = &result.param_updates;
1274
1275 let real = find_array_update(u, real_reason).unwrap();
1276 let imag = find_array_update(u, imag_reason).unwrap();
1277 let abs = find_array_update(u, abs_reason).unwrap();
1278 let ts = find_array_update(u, ts_reason).unwrap();
1279
1280 assert_eq!(real.len(), 8);
1282 assert_eq!(imag.len(), 8);
1283 assert!((real[3] - 8.0).abs() < 1e-9, "real[3] = {}", real[3]);
1285 for k in [0usize, 1, 2, 4, 5, 6, 7] {
1286 assert!(real[k].abs() < 1e-9, "real[{k}] = {}", real[k]);
1287 assert!(imag[k].abs() < 1e-9, "imag[{k}] = {}", imag[k]);
1288 }
1289 assert!(imag[3].abs() < 1e-9, "imag[3] = {}", imag[3]);
1291 assert!((abs[3] - 0.5).abs() < 1e-9, "abs[3] = {}", abs[3]);
1293 assert_eq!(ts.len(), n);
1295 assert!((ts[0] - 1.0).abs() < 1e-9);
1296 }
1297
1298 #[test]
1299 fn test_fft_axes_scale_with_time_per_point() {
1300 let mut proc = fft_proc_with_params(FFTConfig::default());
1302 let time_axis_reason = proc.params.time_axis.unwrap();
1303 let freq_axis_reason = proc.params.freq_axis.unwrap();
1304 let tpp_reason = proc.params.time_per_point.unwrap();
1305 let pool = NDArrayPool::new(0);
1306
1307 use ad_core_rs::plugin::runtime::{ParamChangeValue, PluginParamSnapshot};
1309 proc.on_param_change(
1310 tpp_reason,
1311 &PluginParamSnapshot {
1312 enable_callbacks: true,
1313 reason: tpp_reason,
1314 addr: 0,
1315 value: ParamChangeValue::Float64(0.5),
1316 },
1317 );
1318
1319 let n = 8;
1320 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1321 if let NDDataBuffer::F64(ref mut v) = arr.data {
1322 v[0] = 1.0;
1323 }
1324 let result = proc.process_array(&arr, &pool);
1325 let u = &result.param_updates;
1326
1327 let time_axis = find_array_update(u, time_axis_reason).unwrap();
1328 let freq_axis = find_array_update(u, freq_axis_reason).unwrap();
1329
1330 assert_eq!(time_axis.len(), 8);
1332 assert!((time_axis[1] - 0.5).abs() < 1e-12);
1333 assert!((time_axis[7] - 3.5).abs() < 1e-12);
1334 assert_eq!(freq_axis.len(), 4);
1336 let step = 0.5 / 0.5 / 3.0;
1337 assert!((freq_axis[1] - step).abs() < 1e-12);
1338 assert!((freq_axis[3] - 3.0 * step).abs() < 1e-12);
1339 }
1340
1341 #[test]
1342 fn test_fft_inverse_emits_no_spectrum_waveforms() {
1343 let config = FFTConfig {
1345 mode: FFTMode::Rows1D,
1346 direction: FFTDirection::Inverse,
1347 suppress_dc: false,
1348 num_average: 0,
1349 };
1350 let mut proc = fft_proc_with_params(config);
1351 let pool = NDArrayPool::new(0);
1352 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1353 if let NDDataBuffer::F64(ref mut v) = arr.data {
1354 v[0] = 8.0;
1355 }
1356 let result = proc.process_array(&arr, &pool);
1357 let array_updates = result
1358 .param_updates
1359 .iter()
1360 .filter(|x| matches!(x, ParamUpdate::Float64Array { .. }))
1361 .count();
1362 assert_eq!(
1363 array_updates, 0,
1364 "inverse FFT must not emit spectrum waveforms"
1365 );
1366 }
1367
1368 #[test]
1369 fn test_inverse_fft_preserves_sign() {
1370 let n = 8;
1374 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1375 if let NDDataBuffer::F64(ref mut v) = arr.data {
1376 v[1] = 4.0;
1379 v[n - 1] = 4.0;
1380 }
1381 let config = FFTConfig {
1382 mode: FFTMode::Rows1D,
1383 direction: FFTDirection::Inverse,
1384 suppress_dc: false,
1385 num_average: 0,
1386 };
1387 let mut proc = FFTProcessor::with_config(config);
1388 let pool = NDArrayPool::new(0);
1389 let result = proc.process_array(&arr, &pool);
1390 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
1391 let has_negative = v.iter().any(|&x| x < -1e-6);
1392 assert!(
1393 has_negative,
1394 "inverse FFT must keep negative samples: {v:?}"
1395 );
1396 } else {
1397 panic!("expected F64 data");
1398 }
1399 }
1400}