From dd370499f8fcc5ce5a41a86a08b1e2058af194e7 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Mon, 11 Sep 2023 20:00:48 +0800
Subject: [PATCH] TEC control loop

---
 main.cpp | 42 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/main.cpp b/main.cpp
index 55764d3..bfeb3af 100644
--- a/main.cpp
+++ b/main.cpp
@@ -33,6 +33,11 @@ static int waterfall_width = 1000;
 static int waterfall_height = 2000;
 static unsigned int waterfall_data[1000*2000];
 
+static std::atomic<float> freq_setpoint;
+static std::atomic<float> freq_peak;
+static std::atomic<float> tec_bias;
+static std::atomic<float> tec_p;
+
 static std::atomic<int> fps;
 
 static void dsp_thread() {
@@ -47,10 +52,11 @@ static void dsp_thread() {
     std::vector<std::complex<float>> frames_f(len);
     pocketfft::shape_t shape{len};
     pocketfft::stride_t stride(1);
-    stride[0] = 2*sizeof(float);
+    stride[0] = sizeof(std::complex<float>);
     pocketfft::shape_t axes;
     axes.push_back(0);
-    std::vector<std::complex<float>> frames_ft(len);
+    std::vector<std::complex<float>> frames_ft_rev(len);
+    std::vector<float> frames_mag(len);
 
     int iterations = 0;
     auto last_second = std::chrono::steady_clock::now();
@@ -65,12 +71,22 @@ static void dsp_thread() {
         // FFT
         for(size_t i=0;i<len;i++)
             frames_f[i] = frames[i];
-        pocketfft::c2c(shape, stride, stride, axes, pocketfft::FORWARD, frames_f.data(), frames_ft.data(), 1.0f);
+        pocketfft::c2c(shape, stride, stride, axes, pocketfft::FORWARD, frames_f.data(), frames_ft_rev.data(), 1.0f);
+        // negative frequencies first
+        for(size_t i=0;i<len/2;i++)
+            frames_mag[i] = abs(frames_ft_rev[i + len/2]);
+        for(size_t i=0;i<len/2;i++)
+            frames_mag[i + len/2] = abs(frames_ft_rev[i]);
 
         // stabilize laser
         if((std::chrono::steady_clock::now() - last_tec) >= 100ms) {
+            float freq_peak_local;
+            freq_peak_local = 40.0f*float(distance(frames_mag.begin(), max_element(frames_mag.begin(), frames_mag.end())))/float(len);
+            freq_peak = freq_peak_local;
             // FIXME: net::write seems unimplemented as of libstdc++ 13
-            tec_socket.write_some(net::buffer(std::format(":SOUR2:VOLT {:.3f}\n", 0.1)));
+            float freq_error = freq_peak_local - freq_setpoint;
+            tec_socket.write_some(net::buffer(std::format(":SOUR2:VOLT {:.3f}\n",
+                std::min(tec_bias+tec_p*freq_error, 0.0f))));
             last_tec += 100ms;
         }
 
@@ -80,9 +96,9 @@ static void dsp_thread() {
             std::memmove(&waterfall_data[waterfall_width], &waterfall_data[0], sizeof(int)*waterfall_width*(waterfall_height - 1));
             for(int i=0;i<waterfall_width;i++) {
                 int j = i*len/waterfall_width;
-                j = (j + len/2) % len; // display negative frequencies first
-                waterfall_data[i] = 0xff000000 | 0x010101*std::min(int(abs(frames_ft[j])/900.), 255);
+                waterfall_data[i] = 0xff000000 | 0x010101*std::min(int(frames_mag[j]/900.0f), 255);
             }
+            waterfall_data[int(freq_setpoint*waterfall_width)/40] = 0xff0000ff;
         }
 
         // FPS counter
@@ -192,6 +208,9 @@ int main(int argc, char* argv[])
 
     bool exit = false;
     bool update = true;
+    float freq_setpoint_local = 0.0f;
+    float tec_bias_local = 0.02f;
+    float tec_p_local = 0.001f;
     while(!exit && !glfwWindowShouldClose(window)) {
         glfwPollEvents();
         glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
@@ -216,10 +235,17 @@ int main(int argc, char* argv[])
         ImGui::Begin("fastsa", nullptr, ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoResize);
 
         ImGui::BeginTable("fastsa", 2, ImGuiTableFlags_SizingStretchSame | ImGuiTableFlags_Resizable);
-        ImGui::TableSetupColumn("", 0, 280.0);
-        ImGui::TableSetupColumn("", 0, 1000.0);
+        ImGui::TableSetupColumn("", 0, 280.0f);
+        ImGui::TableSetupColumn("", 0, 1000.0f);
         ImGui::TableNextColumn();
         ImGui::Checkbox("Update", &update);
+        ImGui::Text("Peak: %.3f", (float)freq_peak);
+        ImGui::SliderFloat("Setpoint", &freq_setpoint_local, 0.0f, 40.0f);
+        ImGui::SliderFloat("TEC bias", &tec_bias_local, 0.0f, 0.5f);
+        ImGui::SliderFloat("TEC P", &tec_p_local, -0.05f, 0.05f);
+        freq_setpoint = freq_setpoint_local;
+        tec_bias = tec_bias_local;
+        tec_p = tec_p_local;
         if(ImGui::Button("Exit"))
             exit = true;
         ImGui::Text("FPS: %d", (int)fps);