From 376985b19a25ec53c0f5ab7a8514bce5c4749ca5 Mon Sep 17 00:00:00 2001
From: AsiiaPine <asiiapine@gmail.com>
Date: Tue, 8 Oct 2024 23:19:53 +0300
Subject: [PATCH] simplified dominant search logic, add mag multiplier

---
 Src/common/FFT.cpp           | 19 ++++++++-----------
 Src/common/FFT.hpp           |  4 ++--
 Src/modules/imu/imu.cpp      |  4 ++--
 Src/platform/ubuntu/rfft.hpp |  1 +
 4 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/Src/common/FFT.cpp b/Src/common/FFT.cpp
index 7d77c80..d4cb39d 100644
--- a/Src/common/FFT.cpp
+++ b/Src/common/FFT.cpp
@@ -50,7 +50,7 @@ void FFT::update(float *input) {
                 sizeof(real_t) * overlap_start * 3);
         _fft_buffer_index[axis] = overlap_start * 3;
     }
-    find_dominant();
+    _find_dominant();
 }
 
 void FFT::find_peaks(uint8_t axis) {
@@ -61,12 +61,9 @@ void FFT::find_peaks(uint8_t axis) {
     float bin_mag_sum = 0;
     // calculate magnitudes for each fft bin
     for (uint16_t fft_index = 0; fft_index < size/2; fft_index ++) {
-        real_t real_imag[2] = {0, 0};
-        fft::get_real_imag_by_index(_fft_output_buffer.data(), real_imag, size, fft_index);
-        float real_f, imag_f;
-        fft::convert_real_t_to_float(&real_imag[0], &real_f, 1);
-        fft::convert_real_t_to_float(&real_imag[1], &imag_f, 1);
-        const float fft_magnitude = sqrtf(real_f * real_f + imag_f * imag_f);
+        auto real = fft::get_imag_by_index(fft_output_buffer_float, fft_index);
+        auto imag = fft::get_imag_by_index(fft_output_buffer_float, fft_index);
+        const float fft_magnitude = sqrtf(real * real + imag * imag);
         _peak_magnitudes_all[fft_index] = fft_magnitude;
         bin_mag_sum += fft_magnitude;
     }
@@ -88,7 +85,7 @@ void FFT::find_peaks(uint8_t axis) {
     }
 }
 
-void FFT::find_dominant() {
+void FFT::_find_dominant() {
     if (!is_updated()) {
         return;
     }
@@ -126,7 +123,7 @@ void FFT::_identify_peaks_bins(float peak_magnitude[MAX_NUM_PEAKS],
 
         if (largest_peak_index > 0) {
             raw_peak_index[i] = largest_peak_index;
-            peak_magnitude[i] = _peak_magnitudes_all[largest_peak_index];
+            peak_magnitude[i] = largest_peak;
             // remove peak + sides (included in frequency estimate later)
             _peak_magnitudes_all[largest_peak_index - 1] = 0;
             _peak_magnitudes_all[largest_peak_index]     = 0;
@@ -145,7 +142,7 @@ uint16_t FFT::_estimate_peaks(float* peak_magnitude,
             continue;
         }
         float adjusted_bin = 0.5f *
-                        estimate_peak_freq(fft, 2 * raw_peak_index[peak_new]);
+                        _estimate_peak_freq(fft, 2 * raw_peak_index[peak_new]);
         if (adjusted_bin > size || adjusted_bin < 0) {
             continue;
         }
@@ -196,7 +193,7 @@ static constexpr float tau(float x) {
     return addend_1 - multiplier_2 * addend_2;
 }
 
-float FFT::estimate_peak_freq(float fft[], int peak_index) {
+float FFT::_estimate_peak_freq(float fft[], int peak_index) {
     if (peak_index < 2 || peak_index >= size) {
         return -1;
     }
diff --git a/Src/common/FFT.hpp b/Src/common/FFT.hpp
index 536d574..f6efbf2 100644
--- a/Src/common/FFT.hpp
+++ b/Src/common/FFT.hpp
@@ -58,9 +58,9 @@ class FFT {
     uint8_t n_axes;
     float _sample_rate_hz;
 
-    float estimate_peak_freq(float fft[], int peak_index);
+    float _estimate_peak_freq(float fft[], int peak_index);
     void find_peaks(uint8_t axis);
-    void find_dominant();
+    void _find_dominant();
     // void identify_bin_peaks(uint8_t axis);
     void _identify_peaks_bins(float peak_magnitude[MAX_NUM_PEAKS],
                              uint16_t raw_peak_index[MAX_NUM_PEAKS]);
diff --git a/Src/modules/imu/imu.cpp b/Src/modules/imu/imu.cpp
index 8ea4e4a..6bf818c 100644
--- a/Src/modules/imu/imu.cpp
+++ b/Src/modules/imu/imu.cpp
@@ -102,7 +102,7 @@ void ImuModule::update_accel_fft() {
     }
     fft_accel.update(accel.data());
     pub.msg.accelerometer_integral[0] = fft_accel.dominant_frequency;
-    pub.msg.accelerometer_integral[1] = fft_accel.dominant_mag;
+    pub.msg.accelerometer_integral[1] = fft_accel.dominant_mag * 1000;
     pub.msg.accelerometer_integral[2] = fft_accel.dominant_snr;
 }
 
@@ -112,6 +112,6 @@ void ImuModule::update_gyro_fft() {
     }
     fft_gyro.update(gyro.data());
     pub.msg.rate_gyro_integral[0] = fft_gyro.dominant_frequency;
-    pub.msg.rate_gyro_integral[1] = fft_gyro.dominant_mag;
+    pub.msg.rate_gyro_integral[1] = fft_gyro.dominant_mag * 1000;
     pub.msg.rate_gyro_integral[2] = fft_gyro.dominant_snr;
 }
diff --git a/Src/platform/ubuntu/rfft.hpp b/Src/platform/ubuntu/rfft.hpp
index 78d17b5..3ba897f 100644
--- a/Src/platform/ubuntu/rfft.hpp
+++ b/Src/platform/ubuntu/rfft.hpp
@@ -56,6 +56,7 @@ namespace fft {
     inline T get_imag_by_index(T* in, uint16_t index) {
         return 0;
     }
+
     /*
     The function written based on fftw3 library.
     @param plan: The plan for the r2c transform.