From f6e10bedecb2768b2751f633d672d5ddad8cb15f Mon Sep 17 00:00:00 2001
From: Christos Falas <raspberrymaker@gmail.com>
Date: Tue, 4 Mar 2025 14:07:21 +0000
Subject: [PATCH] Add figure for magnitude

---
 where_fi/visualise/server/__init__.py | 39 ++++++++++++++++++---------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/where_fi/visualise/server/__init__.py b/where_fi/visualise/server/__init__.py
index 9dfaf90..c8f98de 100644
--- a/where_fi/visualise/server/__init__.py
+++ b/where_fi/visualise/server/__init__.py
@@ -43,6 +43,12 @@ figures = {
         x_label="Subcarrier",
         y_label="Phase",
     ),
+    "raw_magn": figure_pb2.Figure(
+        uuid=str(uuid.uuid4()),
+        title="Raw CSI Amplitude",
+        x_label="Subcarrier",
+        y_label="Amplitude",
+    ),
     "unwrapped_phase": figure_pb2.Figure(
         uuid=str(uuid.uuid4()),
         title="Unwrapped CSI Phase",
@@ -62,6 +68,12 @@ figures = {
         x_label="Subcarrier",
         y_label="Phase",
     ),
+    "processed_magn": figure_pb2.Figure(
+        uuid=str(uuid.uuid4()),
+        title="Preprocessed CSI Amplitude",
+        x_label="Subcarrier",
+        y_label="Amplitude",
+    ),
     "aoa_heatmap": figure_pb2.Figure(
         uuid=str(uuid.uuid4()),
         title="AoA Heatmap",
@@ -105,20 +117,21 @@ class FigureServer(figure_pb2_grpc.FigureServiceServicer):
 def add_data(dtype: DataType, new_data: npt.NDArray[np.complex128]) -> None:
     match dtype:
         case DataType.RAW_CSI | DataType.PROCESSED_CSI:
-            uuid = (
-                figures["raw_phase"].uuid
-                if dtype == DataType.RAW_CSI
-                else figures["processed_phase"].uuid
-            )
-            lines = [
-                line_pb2.LineChartData.Line(
-                    y=np.angle(new_data)[:, i, 0], label=f"Antenna {i}"
+            for component, func in [("phase", np.angle), ("magn", np.abs)]:
+                uuid = (
+                    figures[f"raw_{component}"].uuid
+                    if dtype == DataType.RAW_CSI
+                    else figures[f"processed_{component}"].uuid
                 )
-                for i in range(new_data.shape[1])
-            ]
-            linechart = line_pb2.LineChartData(lines=lines)
-            for q in clients.get(uuid, []):
-                q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
+                lines = [
+                    line_pb2.LineChartData.Line(
+                        y=func(new_data)[:, i, 0], label=f"Antenna {i}"
+                    )
+                    for i in range(new_data.shape[1])
+                ]
+                linechart = line_pb2.LineChartData(lines=lines)
+                for q in clients.get(uuid, []):
+                    q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
         case DataType.HEATMAP:
             uuid = figures["aoa_heatmap"].uuid
             heatmap = heatmap_pb2.HeatmapData(