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 FFTDirection {
12 Forward,
13 Inverse,
14}
15
16pub struct FFTConfig {
23 pub direction: FFTDirection,
24 pub suppress_dc: bool,
26 pub num_average: usize,
28}
29
30impl Default for FFTConfig {
31 fn default() -> Self {
32 Self {
33 direction: FFTDirection::Forward,
34 suppress_dc: false,
35 num_average: 0,
36 }
37 }
38}
39
40pub fn next_pow2(n: usize) -> usize {
44 if n <= 1 {
45 return 1;
46 }
47 let mut p = 1usize;
48 while p < n {
49 p <<= 1;
50 }
51 p
52}
53
54pub fn fft_1d_rows(src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
60 if src.dims.is_empty() {
61 return None;
62 }
63
64 let width = src.dims[0].size;
65 let height = if src.dims.len() >= 2 {
66 src.dims[1].size
67 } else {
68 1
69 };
70
71 if width == 0 {
72 return None;
73 }
74
75 let padded = next_pow2(width);
77
78 let mut planner = FftPlanner::<f64>::new();
79 let fft = planner.plan_fft_forward(padded);
80
81 let n_freq = padded / 2;
83 if n_freq == 0 {
84 return None;
85 }
86 let scale = 1.0 / padded as f64;
87
88 let mut magnitudes = vec![0.0f64; n_freq * height];
89 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
90
91 for row in 0..height {
92 for c in row_buf.iter_mut() {
94 *c = Complex::new(0.0, 0.0);
95 }
96 for i in 0..width {
97 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
98 }
99
100 fft.process(&mut row_buf);
101
102 for i in 0..n_freq {
104 magnitudes[row * n_freq + i] = row_buf[i].norm() * scale;
105 }
106
107 if suppress_dc {
108 magnitudes[row * n_freq] = 0.0;
109 }
110 }
111
112 let dims = if height > 1 {
113 vec![NDDimension::new(n_freq), NDDimension::new(height)]
114 } else {
115 vec![NDDimension::new(n_freq)]
116 };
117 let mut arr = NDArray::new(dims, NDDataType::Float64);
118 arr.data = NDDataBuffer::F64(magnitudes);
119 arr.unique_id = src.unique_id;
120 arr.timestamp = src.timestamp;
121 arr.attributes = src.attributes.clone();
122 Some(arr)
123}
124
125pub fn fft_2d(src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
127 if src.dims.len() < 2 {
128 return None;
129 }
130
131 let src_w = src.dims[0].size;
132 let src_h = src.dims[1].size;
133
134 if src_w == 0 || src_h == 0 {
135 return None;
136 }
137
138 let w = next_pow2(src_w);
140 let h = next_pow2(src_h);
141
142 let mut planner = FftPlanner::<f64>::new();
143 let fft_row = planner.plan_fft_forward(w);
144 let fft_col = planner.plan_fft_forward(h);
145
146 let mut data = vec![Complex::new(0.0, 0.0); w * h];
148 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
149
150 for row in 0..src_h {
151 for c in row_buf.iter_mut() {
152 *c = Complex::new(0.0, 0.0);
153 }
154 for i in 0..src_w {
155 row_buf[i] = Complex::new(src.data.get_as_f64(row * src_w + i).unwrap_or(0.0), 0.0);
156 }
157 fft_row.process(&mut row_buf);
158 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
159 }
160
161 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
163
164 for col in 0..w {
165 for row in 0..h {
167 col_buf[row] = data[row * w + col];
168 }
169 fft_col.process(&mut col_buf);
170 for row in 0..h {
172 data[row * w + col] = col_buf[row];
173 }
174 }
175
176 let n_freq_x = w / 2;
178 let n_freq_y = h / 2;
179 if n_freq_x == 0 || n_freq_y == 0 {
180 return None;
181 }
182 let scale = 1.0 / (w * h) as f64;
183
184 let mut magnitudes = vec![0.0f64; n_freq_x * n_freq_y];
185 for fy in 0..n_freq_y {
186 for fx in 0..n_freq_x {
187 magnitudes[fy * n_freq_x + fx] = data[fy * w + fx].norm() * scale;
188 }
189 }
190
191 if suppress_dc {
192 magnitudes[0] = 0.0;
193 }
194
195 let dims = vec![NDDimension::new(n_freq_x), NDDimension::new(n_freq_y)];
196 let mut arr = NDArray::new(dims, NDDataType::Float64);
197 arr.data = NDDataBuffer::F64(magnitudes);
198 arr.unique_id = src.unique_id;
199 arr.timestamp = src.timestamp;
200 arr.attributes = src.attributes.clone();
201 Some(arr)
202}
203
204#[derive(Default)]
206struct FFTParamIndices {
207 direction: Option<usize>,
208 suppress_dc: Option<usize>,
209 num_average: Option<usize>,
210 num_averaged: Option<usize>,
211 reset_average: Option<usize>,
212 time_per_point: Option<usize>,
213 time_series: Option<usize>,
215 real: Option<usize>,
217 imaginary: Option<usize>,
219 abs_value: Option<usize>,
221 time_axis: Option<usize>,
223 freq_axis: Option<usize>,
225}
226
227pub struct FFTProcessor {
228 config: FFTConfig,
229 planner: FftPlanner<f64>,
230 avg_buffer: Option<Vec<f64>>,
232 avg_count: usize,
234 cached_dims: Vec<usize>,
236 time_per_point: f64,
239 params: FFTParamIndices,
240}
241
242impl FFTProcessor {
243 pub fn new() -> Self {
244 Self::with_config(FFTConfig::default())
245 }
246
247 pub fn with_config(config: FFTConfig) -> Self {
248 Self {
249 config,
250 planner: FftPlanner::new(),
251 avg_buffer: None,
252 avg_count: 0,
253 cached_dims: Vec::new(),
254 time_per_point: 1.0,
255 params: FFTParamIndices::default(),
256 }
257 }
258
259 fn check_dims_changed(&mut self, dims: &[NDDimension]) {
261 let current: Vec<usize> = dims.iter().map(|d| d.size).collect();
262 if current != self.cached_dims {
263 self.cached_dims = current;
264 self.avg_buffer = None;
265 self.avg_count = 0;
266 }
267 }
268
269 fn compute_fft(&mut self, src: &NDArray) -> Option<NDArray> {
276 let suppress_dc = self.config.suppress_dc;
277
278 match (src.dims.len(), self.config.direction) {
279 (1, FFTDirection::Forward) => self.compute_fft_1d_rows_forward(src, suppress_dc),
280 (1, FFTDirection::Inverse) => self.compute_fft_1d_rows_inverse(src, suppress_dc),
281 (2, FFTDirection::Forward) => self.compute_fft_2d_forward(src, suppress_dc),
282 (2, FFTDirection::Inverse) => self.compute_fft_2d_inverse(src, suppress_dc),
283 _ => None,
284 }
285 }
286
287 fn compute_row_spectrum(
299 &mut self,
300 src: &NDArray,
301 suppress_dc: bool,
302 ) -> Option<(Vec<f64>, Vec<f64>, Vec<f64>)> {
303 if src.dims.is_empty() {
304 return None;
305 }
306 let width = src.dims[0].size;
307 if width == 0 {
308 return None;
309 }
310 let padded = next_pow2(width);
311 let n_freq = padded / 2;
312 if n_freq == 0 {
313 return None;
314 }
315 let fft = self.planner.plan_fft_forward(padded);
316
317 let mut time_series = vec![0.0f64; padded];
321 for (i, slot) in time_series.iter_mut().enumerate().take(width) {
322 *slot = src.data.get_as_f64(i).unwrap_or(0.0);
323 }
324
325 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
326 for (i, &v) in time_series.iter().enumerate() {
327 row_buf[i] = Complex::new(v, 0.0);
328 }
329 fft.process(&mut row_buf);
330
331 let mut real = vec![0.0f64; n_freq];
332 let mut imag = vec![0.0f64; n_freq];
333 for i in 0..n_freq {
334 real[i] = row_buf[i].re;
335 imag[i] = row_buf[i].im;
336 }
337 if suppress_dc {
338 real[0] = 0.0;
339 imag[0] = 0.0;
340 }
341 Some((time_series, real, imag))
342 }
343
344 fn freq_axis(&self, n_freq: usize) -> Vec<f64> {
347 if n_freq <= 1 {
348 return vec![0.0; n_freq];
349 }
350 let tpp = if self.time_per_point > 0.0 {
351 self.time_per_point
352 } else {
353 1.0
354 };
355 let step = 0.5 / tpp / (n_freq - 1) as f64;
356 (0..n_freq).map(|i| i as f64 * step).collect()
357 }
358
359 fn time_axis(&self, n_time: usize) -> Vec<f64> {
361 let tpp = if self.time_per_point > 0.0 {
362 self.time_per_point
363 } else {
364 1.0
365 };
366 (0..n_time).map(|i| i as f64 * tpp).collect()
367 }
368
369 fn compute_fft_1d_rows_forward(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
370 if src.dims.is_empty() {
371 return None;
372 }
373
374 let width = src.dims[0].size;
375 let height = if src.dims.len() >= 2 {
376 src.dims[1].size
377 } else {
378 1
379 };
380
381 if width == 0 {
382 return None;
383 }
384
385 let padded = next_pow2(width);
387 let fft = self.planner.plan_fft_forward(padded);
388
389 let n_freq = padded / 2;
391 if n_freq == 0 {
392 return None;
393 }
394 let scale = 1.0 / padded as f64;
395
396 let mut magnitudes = vec![0.0f64; n_freq * height];
397 let mut row_buf = vec![Complex::new(0.0, 0.0); padded];
398
399 for row in 0..height {
400 for c in row_buf.iter_mut() {
401 *c = Complex::new(0.0, 0.0);
402 }
403 for i in 0..width {
404 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
405 }
406 fft.process(&mut row_buf);
407 for i in 0..n_freq {
408 magnitudes[row * n_freq + i] = row_buf[i].norm() * scale;
409 }
410 if suppress_dc {
411 magnitudes[row * n_freq] = 0.0;
412 }
413 }
414
415 let dims = if height > 1 {
416 vec![NDDimension::new(n_freq), NDDimension::new(height)]
417 } else {
418 vec![NDDimension::new(n_freq)]
419 };
420 let mut arr = NDArray::new(dims, NDDataType::Float64);
421 arr.data = NDDataBuffer::F64(magnitudes);
422 arr.unique_id = src.unique_id;
423 arr.timestamp = src.timestamp;
424 arr.attributes = src.attributes.clone();
425 Some(arr)
426 }
427
428 fn compute_fft_1d_rows_inverse(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
429 if src.dims.is_empty() {
430 return None;
431 }
432
433 let width = src.dims[0].size;
434 let height = if src.dims.len() >= 2 {
435 src.dims[1].size
436 } else {
437 1
438 };
439
440 if width == 0 {
441 return None;
442 }
443
444 let fft = self.planner.plan_fft_inverse(width);
445 let scale = 1.0 / width as f64;
446
447 let mut samples = vec![0.0f64; width * height];
451 let mut row_buf = vec![Complex::new(0.0, 0.0); width];
452
453 for row in 0..height {
454 for i in 0..width {
455 row_buf[i] = Complex::new(src.data.get_as_f64(row * width + i).unwrap_or(0.0), 0.0);
456 }
457 if suppress_dc {
458 row_buf[0] = Complex::new(0.0, 0.0);
459 }
460 fft.process(&mut row_buf);
461 for (i, c) in row_buf.iter().enumerate() {
462 samples[row * width + i] = c.re * scale;
463 }
464 }
465
466 let dims = src.dims.clone();
467 let mut arr = NDArray::new(dims, NDDataType::Float64);
468 arr.data = NDDataBuffer::F64(samples);
469 arr.unique_id = src.unique_id;
470 arr.timestamp = src.timestamp;
471 arr.attributes = src.attributes.clone();
472 Some(arr)
473 }
474
475 fn compute_fft_2d_forward(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
476 if src.dims.len() < 2 {
477 return None;
478 }
479
480 let src_w = src.dims[0].size;
481 let src_h = src.dims[1].size;
482
483 if src_w == 0 || src_h == 0 {
484 return None;
485 }
486
487 let w = next_pow2(src_w);
489 let h = next_pow2(src_h);
490
491 let fft_row = self.planner.plan_fft_forward(w);
492 let fft_col = self.planner.plan_fft_forward(h);
493
494 let mut data = vec![Complex::new(0.0, 0.0); w * h];
495 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
496
497 for row in 0..src_h {
498 for c in row_buf.iter_mut() {
499 *c = Complex::new(0.0, 0.0);
500 }
501 for i in 0..src_w {
502 row_buf[i] = Complex::new(src.data.get_as_f64(row * src_w + i).unwrap_or(0.0), 0.0);
503 }
504 fft_row.process(&mut row_buf);
505 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
506 }
507
508 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
509 for col in 0..w {
510 for row in 0..h {
511 col_buf[row] = data[row * w + col];
512 }
513 fft_col.process(&mut col_buf);
514 for row in 0..h {
515 data[row * w + col] = col_buf[row];
516 }
517 }
518
519 let n_freq_x = w / 2;
521 let n_freq_y = h / 2;
522 if n_freq_x == 0 || n_freq_y == 0 {
523 return None;
524 }
525 let scale = 1.0 / (w * h) as f64;
526
527 let mut magnitudes = vec![0.0f64; n_freq_x * n_freq_y];
528 for fy in 0..n_freq_y {
529 for fx in 0..n_freq_x {
530 magnitudes[fy * n_freq_x + fx] = data[fy * w + fx].norm() * scale;
531 }
532 }
533
534 if suppress_dc {
535 magnitudes[0] = 0.0;
536 }
537
538 let dims = vec![NDDimension::new(n_freq_x), NDDimension::new(n_freq_y)];
539 let mut arr = NDArray::new(dims, NDDataType::Float64);
540 arr.data = NDDataBuffer::F64(magnitudes);
541 arr.unique_id = src.unique_id;
542 arr.timestamp = src.timestamp;
543 arr.attributes = src.attributes.clone();
544 Some(arr)
545 }
546
547 fn compute_fft_2d_inverse(&mut self, src: &NDArray, suppress_dc: bool) -> Option<NDArray> {
548 if src.dims.len() < 2 {
549 return None;
550 }
551
552 let w = src.dims[0].size;
553 let h = src.dims[1].size;
554
555 if w == 0 || h == 0 {
556 return None;
557 }
558
559 let fft_row = self.planner.plan_fft_inverse(w);
560 let fft_col = self.planner.plan_fft_inverse(h);
561 let scale = 1.0 / (w * h) as f64;
562
563 let mut data = vec![Complex::new(0.0, 0.0); w * h];
564 for i in 0..w * h {
565 data[i] = Complex::new(src.data.get_as_f64(i).unwrap_or(0.0), 0.0);
566 }
567
568 if suppress_dc {
569 data[0] = Complex::new(0.0, 0.0);
570 }
571
572 let mut col_buf = vec![Complex::new(0.0, 0.0); h];
573 for col in 0..w {
574 for row in 0..h {
575 col_buf[row] = data[row * w + col];
576 }
577 fft_col.process(&mut col_buf);
578 for row in 0..h {
579 data[row * w + col] = col_buf[row];
580 }
581 }
582
583 let mut row_buf = vec![Complex::new(0.0, 0.0); w];
584 for row in 0..h {
585 row_buf.copy_from_slice(&data[row * w..(row * w + w)]);
586 fft_row.process(&mut row_buf);
587 data[row * w..(row * w + w)].copy_from_slice(&row_buf);
588 }
589
590 let samples: Vec<f64> = data.iter().map(|c| c.re * scale).collect();
592
593 let dims = vec![NDDimension::new(w), NDDimension::new(h)];
594 let mut arr = NDArray::new(dims, NDDataType::Float64);
595 arr.data = NDDataBuffer::F64(samples);
596 arr.unique_id = src.unique_id;
597 arr.timestamp = src.timestamp;
598 arr.attributes = src.attributes.clone();
599 Some(arr)
600 }
601
602 fn apply_averaging(&mut self, magnitudes: &[f64]) -> Vec<f64> {
607 let num_avg = self.config.num_average;
608 if num_avg <= 1 {
609 return magnitudes.to_vec();
610 }
611
612 let buf = self
613 .avg_buffer
614 .get_or_insert_with(|| vec![0.0; magnitudes.len()]);
615
616 if buf.len() != magnitudes.len() {
618 *buf = vec![0.0; magnitudes.len()];
619 self.avg_count = 0;
620 }
621
622 self.avg_count += 1;
623 let n = self.avg_count.min(num_avg) as f64;
625 let new_fraction = 1.0 / n;
626 let old_fraction = 1.0 - new_fraction;
627
628 for (b, &m) in buf.iter_mut().zip(magnitudes.iter()) {
630 *b = *b * old_fraction + m * new_fraction;
631 }
632
633 buf.clone()
634 }
635}
636
637impl Default for FFTProcessor {
638 fn default() -> Self {
639 Self::new()
640 }
641}
642
643impl NDPluginProcess for FFTProcessor {
644 fn process_array(&mut self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
645 use ad_core_rs::plugin::runtime::ParamUpdate;
646
647 let rank = array.dims.len();
652 if rank != 1 && rank != 2 {
653 return ProcessResult::sink(Vec::new());
654 }
655
656 self.check_dims_changed(&array.dims);
657
658 let result = self.compute_fft(array);
659 let mut updates = Vec::new();
660 if let Some(idx) = self.params.num_averaged {
661 updates.push(ParamUpdate::int32(idx, self.avg_count as i32));
662 }
663
664 let mut averaged_mags: Option<Vec<f64>> = None;
673 if self.config.direction == FFTDirection::Forward {
674 let suppress_dc = self.config.suppress_dc;
675 if let Some((time_series, real, imag)) = self.compute_row_spectrum(array, suppress_dc) {
676 let n_time = time_series.len();
677 let n_freq = real.len();
678 if let Some(idx) = self.params.time_series {
679 updates.push(ParamUpdate::float64_array(idx, time_series));
680 }
681 if let Some(idx) = self.params.real {
682 updates.push(ParamUpdate::float64_array(idx, real));
683 }
684 if let Some(idx) = self.params.imaginary {
685 updates.push(ParamUpdate::float64_array(idx, imag));
686 }
687 if let Some(idx) = self.params.time_axis {
688 updates.push(ParamUpdate::float64_array(idx, self.time_axis(n_time)));
689 }
690 if let Some(idx) = self.params.freq_axis {
691 updates.push(ParamUpdate::float64_array(idx, self.freq_axis(n_freq)));
692 }
693 }
694 }
695
696 match result {
697 Some(mut out) => {
698 if self.config.num_average > 1 {
699 if let NDDataBuffer::F64(ref mags) = out.data {
700 let averaged = self.apply_averaging(mags);
701 averaged_mags = Some(averaged.clone());
702 out.data = NDDataBuffer::F64(averaged);
703 }
704 }
705 if self.config.direction == FFTDirection::Forward {
709 if let Some(idx) = self.params.abs_value {
710 let abs = match (&averaged_mags, &out.data) {
711 (Some(avg), _) => avg.clone(),
712 (None, NDDataBuffer::F64(mags)) => mags.clone(),
713 _ => Vec::new(),
714 };
715 if !abs.is_empty() {
716 updates.push(ParamUpdate::float64_array(idx, abs));
717 }
718 }
719 }
720 let mut r = ProcessResult::arrays(vec![Arc::new(out)]);
721 r.param_updates = updates;
722 r
723 }
724 None => ProcessResult::sink(updates),
725 }
726 }
727
728 fn plugin_type(&self) -> &str {
729 "NDPluginFFT"
730 }
731
732 fn register_params(
733 &mut self,
734 base: &mut asyn_rs::port::PortDriverBase,
735 ) -> asyn_rs::error::AsynResult<()> {
736 use asyn_rs::param::ParamType;
737 base.create_param("FFT_TIME_PER_POINT", ParamType::Float64)?;
738 base.create_param("FFT_TIME_AXIS", ParamType::Float64Array)?;
739 base.create_param("FFT_FREQ_AXIS", ParamType::Float64Array)?;
740 base.create_param("FFT_DIRECTION", ParamType::Int32)?;
741 base.create_param("FFT_SUPPRESS_DC", ParamType::Int32)?;
742 base.create_param("FFT_NUM_AVERAGE", ParamType::Int32)?;
743 base.create_param("FFT_NUM_AVERAGED", ParamType::Int32)?;
744 base.create_param("FFT_RESET_AVERAGE", ParamType::Int32)?;
745 base.create_param("FFT_TIME_SERIES", ParamType::Float64Array)?;
746 base.create_param("FFT_REAL", ParamType::Float64Array)?;
747 base.create_param("FFT_IMAGINARY", ParamType::Float64Array)?;
748 base.create_param("FFT_ABS_VALUE", ParamType::Float64Array)?;
749
750 self.params.direction = base.find_param("FFT_DIRECTION");
751 self.params.suppress_dc = base.find_param("FFT_SUPPRESS_DC");
752 self.params.num_average = base.find_param("FFT_NUM_AVERAGE");
753 self.params.num_averaged = base.find_param("FFT_NUM_AVERAGED");
754 self.params.reset_average = base.find_param("FFT_RESET_AVERAGE");
755 self.params.time_per_point = base.find_param("FFT_TIME_PER_POINT");
756 self.params.time_series = base.find_param("FFT_TIME_SERIES");
757 self.params.real = base.find_param("FFT_REAL");
758 self.params.imaginary = base.find_param("FFT_IMAGINARY");
759 self.params.abs_value = base.find_param("FFT_ABS_VALUE");
760 self.params.time_axis = base.find_param("FFT_TIME_AXIS");
761 self.params.freq_axis = base.find_param("FFT_FREQ_AXIS");
762 Ok(())
763 }
764
765 fn on_param_change(
766 &mut self,
767 reason: usize,
768 params: &ad_core_rs::plugin::runtime::PluginParamSnapshot,
769 ) -> ad_core_rs::plugin::runtime::ParamChangeResult {
770 if Some(reason) == self.params.direction {
771 self.config.direction = if params.value.as_i32() == 0 {
772 FFTDirection::Forward
773 } else {
774 FFTDirection::Inverse
775 };
776 } else if Some(reason) == self.params.suppress_dc {
777 self.config.suppress_dc = params.value.as_i32() != 0;
778 } else if Some(reason) == self.params.num_average {
779 self.config.num_average = params.value.as_i32().max(0) as usize;
780 } else if Some(reason) == self.params.reset_average {
781 if params.value.as_i32() != 0 {
782 self.avg_buffer = None;
783 self.avg_count = 0;
784 }
785 } else if Some(reason) == self.params.time_per_point {
786 let v = params.value.as_f64();
788 if v > 0.0 {
789 self.time_per_point = v;
790 }
791 }
792 ad_core_rs::plugin::runtime::ParamChangeResult::updates(vec![])
793 }
794}
795
796#[cfg(test)]
797mod tests {
798 use super::*;
799
800 #[test]
801 fn test_fft_1d_dc() {
802 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
804 if let NDDataBuffer::F64(ref mut v) = arr.data {
805 for i in 0..8 {
806 v[i] = 1.0;
807 }
808 }
809
810 let result = fft_1d_rows(&arr, false).unwrap();
811 assert_eq!(result.dims[0].size, 4);
813 if let NDDataBuffer::F64(ref v) = result.data {
814 assert!((v[0] - 1.0).abs() < 1e-10);
816 assert!(v[1].abs() < 1e-10);
818 }
819 }
820
821 #[test]
822 fn test_fft_1d_sine() {
823 let n = 16;
825 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
826 if let NDDataBuffer::F64(ref mut v) = arr.data {
827 for i in 0..n {
828 v[i] = (2.0 * std::f64::consts::PI * i as f64 / n as f64).sin();
829 }
830 }
831
832 let result = fft_1d_rows(&arr, false).unwrap();
833 assert_eq!(result.dims[0].size, 8);
835 if let NDDataBuffer::F64(ref v) = result.data {
836 assert!(v[0].abs() < 1e-10);
838 assert!((v[1] - 0.5).abs() < 1e-10);
840 assert!(v[2].abs() < 1e-10);
842 }
843 }
844
845 #[test]
846 fn test_fft_2d_dimensions() {
847 let arr = NDArray::new(
848 vec![NDDimension::new(4), NDDimension::new(4)],
849 NDDataType::UInt8,
850 );
851 let result = fft_2d(&arr, false).unwrap();
852 assert_eq!(result.dims[0].size, 2);
854 assert_eq!(result.dims[1].size, 2);
855 assert_eq!(result.data.data_type(), NDDataType::Float64);
856 }
857
858 #[test]
859 fn test_fft_1d_suppress_dc() {
860 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
862 if let NDDataBuffer::F64(ref mut v) = arr.data {
863 for i in 0..8 {
864 v[i] = 1.0;
865 }
866 }
867
868 let result = fft_1d_rows(&arr, true).unwrap();
869 if let NDDataBuffer::F64(ref v) = result.data {
870 assert!((v[0]).abs() < 1e-15);
872 assert!(v[1].abs() < 1e-10);
874 } else {
875 panic!("expected F64 data");
876 }
877 }
878
879 #[test]
880 fn test_fft_2d_suppress_dc() {
881 let mut arr = NDArray::new(
883 vec![NDDimension::new(4), NDDimension::new(4)],
884 NDDataType::Float64,
885 );
886 if let NDDataBuffer::F64(ref mut v) = arr.data {
887 for val in v.iter_mut() {
888 *val = 3.0;
889 }
890 }
891
892 let result = fft_2d(&arr, true).unwrap();
893 if let NDDataBuffer::F64(ref v) = result.data {
894 assert!((v[0]).abs() < 1e-15);
896 } else {
897 panic!("expected F64 data");
898 }
899 }
900
901 #[test]
902 fn test_fft_2d_known_dc() {
903 let mut arr = NDArray::new(
905 vec![NDDimension::new(4), NDDimension::new(4)],
906 NDDataType::Float64,
907 );
908 if let NDDataBuffer::F64(ref mut v) = arr.data {
909 for val in v.iter_mut() {
910 *val = 2.0;
911 }
912 }
913
914 let result = fft_2d(&arr, false).unwrap();
915 assert_eq!(result.dims[0].size, 2);
917 assert_eq!(result.dims[1].size, 2);
918 if let NDDataBuffer::F64(ref v) = result.data {
919 assert!((v[0] - 2.0).abs() < 1e-10, "DC = {}, expected 2", v[0]);
921 for i in 1..v.len() {
923 assert!(v[i].abs() < 1e-10, "bin {} = {}, expected ~0", i, v[i]);
924 }
925 } else {
926 panic!("expected F64 data");
927 }
928 }
929
930 #[test]
931 fn test_fft_1d_known_cosine_peaks() {
932 let n = 16;
934 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
935 if let NDDataBuffer::F64(ref mut v) = arr.data {
936 for i in 0..n {
937 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
938 }
939 }
940
941 let result = fft_1d_rows(&arr, false).unwrap();
942 assert_eq!(result.dims[0].size, 8);
944 if let NDDataBuffer::F64(ref v) = result.data {
945 assert!(v[0].abs() < 1e-10);
947 assert!(
949 (v[3] - 0.5).abs() < 1e-10,
950 "k=3 magnitude = {}, expected 0.5",
951 v[3]
952 );
953 for k in [1, 2, 4, 5, 6, 7] {
955 assert!(
956 v[k].abs() < 1e-10,
957 "k={} magnitude = {}, expected ~0",
958 k,
959 v[k]
960 );
961 }
962 } else {
963 panic!("expected F64 data");
964 }
965 }
966
967 #[test]
968 fn test_processor_with_config() {
969 let config = FFTConfig {
970 direction: FFTDirection::Forward,
971 suppress_dc: true,
972 num_average: 0,
973 };
974 let mut proc = FFTProcessor::with_config(config);
975 let pool = NDArrayPool::new(0);
976
977 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
978 if let NDDataBuffer::F64(ref mut v) = arr.data {
979 for i in 0..8 {
980 v[i] = 5.0;
981 }
982 }
983
984 let result = proc.process_array(&arr, &pool);
985 assert_eq!(result.output_arrays.len(), 1);
986 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
987 assert!(v[0].abs() < 1e-15);
989 } else {
990 panic!("expected F64 data");
991 }
992 }
993
994 #[test]
995 fn test_processor_averaging() {
996 let config = FFTConfig {
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 direction: FFTDirection::Forward,
1039 suppress_dc: false,
1040 num_average: 3,
1041 };
1042 let mut proc = FFTProcessor::with_config(config);
1043 let pool = NDArrayPool::new(0);
1044
1045 let mut arr1 = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1047 if let NDDataBuffer::F64(ref mut v) = arr1.data {
1048 for i in 0..8 {
1049 v[i] = 1.0;
1050 }
1051 }
1052 let _ = proc.process_array(&arr1, &pool);
1053 assert_eq!(proc.avg_count, 1);
1054
1055 let mut arr2 = NDArray::new(vec![NDDimension::new(4)], NDDataType::Float64);
1057 if let NDDataBuffer::F64(ref mut v) = arr2.data {
1058 for i in 0..4 {
1059 v[i] = 1.0;
1060 }
1061 }
1062 let _ = proc.process_array(&arr2, &pool);
1063 assert_eq!(proc.avg_count, 1);
1065 }
1066
1067 #[test]
1068 fn test_fft_1d_multirow() {
1069 let w = 4;
1071 let h = 2;
1072 let mut arr = NDArray::new(
1073 vec![NDDimension::new(w), NDDimension::new(h)],
1074 NDDataType::Float64,
1075 );
1076 if let NDDataBuffer::F64(ref mut v) = arr.data {
1077 for i in 0..w {
1079 v[i] = 1.0;
1080 }
1081 for i in w..2 * w {
1083 v[i] = 3.0;
1084 }
1085 }
1086
1087 let result = fft_1d_rows(&arr, false).unwrap();
1088 let n_freq = w / 2; assert_eq!(result.dims[0].size, n_freq);
1090 if let NDDataBuffer::F64(ref v) = result.data {
1091 assert!((v[0] - 1.0).abs() < 1e-10);
1093 assert!((v[n_freq] - 3.0).abs() < 1e-10);
1095 } else {
1096 panic!("expected F64 data");
1097 }
1098 }
1099
1100 #[test]
1101 fn test_inverse_fft_1d() {
1102 let n = 8;
1106 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1107 if let NDDataBuffer::F64(ref mut v) = arr.data {
1108 v[0] = 8.0; }
1111
1112 let config = FFTConfig {
1113 direction: FFTDirection::Inverse,
1114 suppress_dc: false,
1115 num_average: 0,
1116 };
1117 let mut proc = FFTProcessor::with_config(config);
1118 let pool = NDArrayPool::new(0);
1119
1120 let result = proc.process_array(&arr, &pool);
1121 assert_eq!(result.output_arrays.len(), 1);
1122 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
1123 for i in 0..n {
1125 assert!(
1126 (v[i] - 1.0).abs() < 1e-10,
1127 "sample {} = {}, expected 1.0",
1128 i,
1129 v[i]
1130 );
1131 }
1132 } else {
1133 panic!("expected F64 data");
1134 }
1135 }
1136
1137 #[test]
1138 fn test_fft_preserves_metadata() {
1139 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::Float64);
1140 arr.unique_id = 42;
1141 if let NDDataBuffer::F64(ref mut v) = arr.data {
1142 v[0] = 1.0;
1143 }
1144
1145 let result = fft_1d_rows(&arr, false).unwrap();
1146 assert_eq!(result.unique_id, 42);
1147 assert_eq!(result.timestamp, arr.timestamp);
1148 }
1149
1150 #[test]
1151 fn test_next_pow2() {
1152 assert_eq!(next_pow2(0), 1);
1153 assert_eq!(next_pow2(1), 1);
1154 assert_eq!(next_pow2(2), 2);
1155 assert_eq!(next_pow2(3), 4);
1156 assert_eq!(next_pow2(5), 8);
1157 assert_eq!(next_pow2(8), 8);
1158 assert_eq!(next_pow2(100), 128);
1159 }
1160
1161 #[test]
1162 fn test_fft_1d_pads_to_power_of_two() {
1163 let n = 5; let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1167 if let NDDataBuffer::F64(ref mut v) = arr.data {
1168 for i in 0..n {
1169 v[i] = 1.0;
1170 }
1171 }
1172 let result = fft_1d_rows(&arr, false).unwrap();
1173 assert_eq!(result.dims[0].size, 4); }
1175
1176 #[test]
1177 fn test_fft_2d_pads_to_power_of_two() {
1178 let arr = NDArray::new(
1180 vec![NDDimension::new(6), NDDimension::new(3)],
1181 NDDataType::Float64,
1182 );
1183 let result = fft_2d(&arr, false).unwrap();
1184 assert_eq!(result.dims[0].size, 4); assert_eq!(result.dims[1].size, 2); }
1187
1188 #[test]
1189 fn test_adp9_processor_selects_2d_fft_from_input_rank() {
1190 let mut proc = FFTProcessor::new();
1194 let pool = NDArrayPool::new(0);
1195
1196 let mut arr = NDArray::new(
1197 vec![NDDimension::new(4), NDDimension::new(4)],
1198 NDDataType::Float64,
1199 );
1200 if let NDDataBuffer::F64(ref mut v) = arr.data {
1201 v.iter_mut().for_each(|x| *x = 2.0);
1202 }
1203 let result = proc.process_array(&arr, &pool);
1204 assert_eq!(result.output_arrays.len(), 1);
1205 let out = &result.output_arrays[0];
1206 assert_eq!(out.dims.len(), 2);
1207 assert_eq!(out.dims[0].size, 2); assert_eq!(out.dims[1].size, 2); if let NDDataBuffer::F64(ref v) = out.data {
1210 assert!((v[0] - 2.0).abs() < 1e-10, "DC = {}", v[0]);
1212 }
1213 }
1214
1215 #[test]
1216 fn test_adp9_processor_keeps_1d_input_1d() {
1217 let mut proc = FFTProcessor::new();
1219 let pool = NDArrayPool::new(0);
1220 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1221 if let NDDataBuffer::F64(ref mut v) = arr.data {
1222 v.iter_mut().for_each(|x| *x = 1.0);
1223 }
1224 let result = proc.process_array(&arr, &pool);
1225 let out = &result.output_arrays[0];
1226 assert_eq!(out.dims.len(), 1);
1227 assert_eq!(out.dims[0].size, 4); }
1229
1230 #[test]
1231 fn test_adp9_processor_rejects_rank_above_2() {
1232 let mut proc = fft_proc_with_params(FFTConfig::default());
1235 let pool = NDArrayPool::new(0);
1236 let arr = NDArray::new(
1237 vec![
1238 NDDimension::new(3),
1239 NDDimension::new(4),
1240 NDDimension::new(4),
1241 ],
1242 NDDataType::Float64,
1243 );
1244 let result = proc.process_array(&arr, &pool);
1245 assert_eq!(result.output_arrays.len(), 0);
1246 assert!(
1247 result.param_updates.is_empty(),
1248 "rank>2 must emit no waveforms, got {} updates",
1249 result.param_updates.len()
1250 );
1251 }
1252
1253 use ad_core_rs::plugin::runtime::ParamUpdate;
1256
1257 fn fft_proc_with_params(config: FFTConfig) -> FFTProcessor {
1259 let mut proc = FFTProcessor::with_config(config);
1260 let mut base =
1261 asyn_rs::port::PortDriverBase::new("FFT_TEST", 1, asyn_rs::port::PortFlags::default());
1262 proc.register_params(&mut base).unwrap();
1263 proc
1264 }
1265
1266 fn find_array_update(updates: &[ParamUpdate], reason: usize) -> Option<&[f64]> {
1268 updates.iter().find_map(|u| match u {
1269 ParamUpdate::Float64Array {
1270 reason: r, value, ..
1271 } if *r == reason => Some(value.as_slice()),
1272 _ => None,
1273 })
1274 }
1275
1276 #[test]
1277 fn test_fft_emits_all_waveforms() {
1278 let mut proc = fft_proc_with_params(FFTConfig::default());
1281 let pool = NDArrayPool::new(0);
1282
1283 let n = 16;
1284 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1285 if let NDDataBuffer::F64(ref mut v) = arr.data {
1286 for i in 0..n {
1287 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
1288 }
1289 }
1290 let result = proc.process_array(&arr, &pool);
1291 let u = &result.param_updates;
1292
1293 for reason in [
1296 proc.params.time_series.unwrap(),
1297 proc.params.real.unwrap(),
1298 proc.params.imaginary.unwrap(),
1299 proc.params.abs_value.unwrap(),
1300 proc.params.time_axis.unwrap(),
1301 proc.params.freq_axis.unwrap(),
1302 ] {
1303 let wf = find_array_update(u, reason)
1304 .unwrap_or_else(|| panic!("missing waveform for reason {reason}"));
1305 assert!(!wf.is_empty(), "waveform {reason} is empty");
1306 }
1307 let array_updates = u
1308 .iter()
1309 .filter(|x| matches!(x, ParamUpdate::Float64Array { .. }))
1310 .count();
1311 assert_eq!(
1312 array_updates, 6,
1313 "expected 6 waveform updates, got {array_updates}"
1314 );
1315 }
1316
1317 #[test]
1318 fn test_fft_real_imaginary_match_spectrum() {
1319 let mut proc = fft_proc_with_params(FFTConfig::default());
1322 let real_reason = proc.params.real.unwrap();
1323 let imag_reason = proc.params.imaginary.unwrap();
1324 let abs_reason = proc.params.abs_value.unwrap();
1325 let ts_reason = proc.params.time_series.unwrap();
1326 let pool = NDArrayPool::new(0);
1327
1328 let n = 16;
1329 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1330 if let NDDataBuffer::F64(ref mut v) = arr.data {
1331 for i in 0..n {
1332 v[i] = (2.0 * std::f64::consts::PI * 3.0 * i as f64 / n as f64).cos();
1333 }
1334 }
1335 let result = proc.process_array(&arr, &pool);
1336 let u = &result.param_updates;
1337
1338 let real = find_array_update(u, real_reason).unwrap();
1339 let imag = find_array_update(u, imag_reason).unwrap();
1340 let abs = find_array_update(u, abs_reason).unwrap();
1341 let ts = find_array_update(u, ts_reason).unwrap();
1342
1343 assert_eq!(real.len(), 8);
1345 assert_eq!(imag.len(), 8);
1346 assert!((real[3] - 8.0).abs() < 1e-9, "real[3] = {}", real[3]);
1348 for k in [0usize, 1, 2, 4, 5, 6, 7] {
1349 assert!(real[k].abs() < 1e-9, "real[{k}] = {}", real[k]);
1350 assert!(imag[k].abs() < 1e-9, "imag[{k}] = {}", imag[k]);
1351 }
1352 assert!(imag[3].abs() < 1e-9, "imag[3] = {}", imag[3]);
1354 assert!((abs[3] - 0.5).abs() < 1e-9, "abs[3] = {}", abs[3]);
1356 assert_eq!(ts.len(), n);
1358 assert!((ts[0] - 1.0).abs() < 1e-9);
1359 }
1360
1361 #[test]
1362 fn test_fft_axes_scale_with_time_per_point() {
1363 let mut proc = fft_proc_with_params(FFTConfig::default());
1365 let time_axis_reason = proc.params.time_axis.unwrap();
1366 let freq_axis_reason = proc.params.freq_axis.unwrap();
1367 let tpp_reason = proc.params.time_per_point.unwrap();
1368 let pool = NDArrayPool::new(0);
1369
1370 use ad_core_rs::plugin::runtime::{ParamChangeValue, PluginParamSnapshot};
1372 proc.on_param_change(
1373 tpp_reason,
1374 &PluginParamSnapshot {
1375 enable_callbacks: true,
1376 reason: tpp_reason,
1377 addr: 0,
1378 value: ParamChangeValue::Float64(0.5),
1379 },
1380 );
1381
1382 let n = 8;
1383 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1384 if let NDDataBuffer::F64(ref mut v) = arr.data {
1385 v[0] = 1.0;
1386 }
1387 let result = proc.process_array(&arr, &pool);
1388 let u = &result.param_updates;
1389
1390 let time_axis = find_array_update(u, time_axis_reason).unwrap();
1391 let freq_axis = find_array_update(u, freq_axis_reason).unwrap();
1392
1393 assert_eq!(time_axis.len(), 8);
1395 assert!((time_axis[1] - 0.5).abs() < 1e-12);
1396 assert!((time_axis[7] - 3.5).abs() < 1e-12);
1397 assert_eq!(freq_axis.len(), 4);
1399 let step = 0.5 / 0.5 / 3.0;
1400 assert!((freq_axis[1] - step).abs() < 1e-12);
1401 assert!((freq_axis[3] - 3.0 * step).abs() < 1e-12);
1402 }
1403
1404 #[test]
1405 fn test_adp25_timeseries_and_timeaxis_use_padded_length() {
1406 let mut proc = fft_proc_with_params(FFTConfig::default());
1410 let ts_reason = proc.params.time_series.unwrap();
1411 let time_axis_reason = proc.params.time_axis.unwrap();
1412 let real_reason = proc.params.real.unwrap();
1413 let freq_axis_reason = proc.params.freq_axis.unwrap();
1414 let pool = NDArrayPool::new(0);
1415
1416 let n = 5;
1417 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1418 if let NDDataBuffer::F64(ref mut v) = arr.data {
1419 for (i, x) in v.iter_mut().enumerate() {
1420 *x = (i + 1) as f64; }
1422 }
1423 let result = proc.process_array(&arr, &pool);
1424 let u = &result.param_updates;
1425
1426 let ts = find_array_update(u, ts_reason).unwrap();
1427 let time_axis = find_array_update(u, time_axis_reason).unwrap();
1428 let real = find_array_update(u, real_reason).unwrap();
1429 let freq_axis = find_array_update(u, freq_axis_reason).unwrap();
1430
1431 assert_eq!(ts.len(), 8);
1433 assert_eq!(&ts[..5], &[1.0, 2.0, 3.0, 4.0, 5.0]);
1434 assert_eq!(&ts[5..], &[0.0, 0.0, 0.0]);
1435 assert_eq!(time_axis.len(), 8);
1437 assert_eq!(real.len(), 4);
1439 assert_eq!(freq_axis.len(), 4);
1440 }
1441
1442 #[test]
1443 fn test_fft_inverse_emits_no_spectrum_waveforms() {
1444 let config = FFTConfig {
1446 direction: FFTDirection::Inverse,
1447 suppress_dc: false,
1448 num_average: 0,
1449 };
1450 let mut proc = fft_proc_with_params(config);
1451 let pool = NDArrayPool::new(0);
1452 let mut arr = NDArray::new(vec![NDDimension::new(8)], NDDataType::Float64);
1453 if let NDDataBuffer::F64(ref mut v) = arr.data {
1454 v[0] = 8.0;
1455 }
1456 let result = proc.process_array(&arr, &pool);
1457 let array_updates = result
1458 .param_updates
1459 .iter()
1460 .filter(|x| matches!(x, ParamUpdate::Float64Array { .. }))
1461 .count();
1462 assert_eq!(
1463 array_updates, 0,
1464 "inverse FFT must not emit spectrum waveforms"
1465 );
1466 }
1467
1468 #[test]
1469 fn test_inverse_fft_preserves_sign() {
1470 let n = 8;
1474 let mut arr = NDArray::new(vec![NDDimension::new(n)], NDDataType::Float64);
1475 if let NDDataBuffer::F64(ref mut v) = arr.data {
1476 v[1] = 4.0;
1479 v[n - 1] = 4.0;
1480 }
1481 let config = FFTConfig {
1482 direction: FFTDirection::Inverse,
1483 suppress_dc: false,
1484 num_average: 0,
1485 };
1486 let mut proc = FFTProcessor::with_config(config);
1487 let pool = NDArrayPool::new(0);
1488 let result = proc.process_array(&arr, &pool);
1489 if let NDDataBuffer::F64(ref v) = result.output_arrays[0].data {
1490 let has_negative = v.iter().any(|&x| x < -1e-6);
1491 assert!(
1492 has_negative,
1493 "inverse FFT must keep negative samples: {v:?}"
1494 );
1495 } else {
1496 panic!("expected F64 data");
1497 }
1498 }
1499}