From 6ef4fe3dc1f77d4896e8e7972db7e6394ea9008e Mon Sep 17 00:00:00 2001
From: Christos Falas <raspberrymaker@gmail.com>
Date: Tue, 4 Mar 2025 14:01:52 +0000
Subject: [PATCH] Improve heatmap visualisation

- Correct axis titles
- Add axis ticks to indicate angle of arrival/time of flight instead of
  arbitrary numbering
---
 .../visualise/frontend/src/components/Figure.vue   | 14 ++++++++++++++
 .../visualise/protos/figure_type/heatmap.proto     |  4 ++++
 where_fi/visualise/server/__init__.py              |  9 +++++++--
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/where_fi/visualise/frontend/src/components/Figure.vue b/where_fi/visualise/frontend/src/components/Figure.vue
index f4064d7..7a5b08b 100644
--- a/where_fi/visualise/frontend/src/components/Figure.vue
+++ b/where_fi/visualise/frontend/src/components/Figure.vue
@@ -36,10 +36,24 @@ const plotData = computed(() => {
           color: line.color ? line.color : undefined,
         }))
       case 'heatmap':
+        const xmin = data.value.figure.heatmap.xMin ?? 0
+        const xmax = data.value.figure.heatmap.xMax ?? data.value.figure.heatmap.width
+        const ymin = data.value.figure.heatmap.yMin ?? 0
+        const ymax = data.value.figure.heatmap.yMax ?? data.value.figure.heatmap.height
+        const x = Array(data.value.figure.heatmap.width)
+          .fill(0)
+          .map((_, i) => xmin + ((xmax - xmin) * i) / data.value.figure.heatmap.width)
+        const y = Array(data.value.figure.heatmap.height)
+          .fill(0)
+          .map((_, i) => ymin + ((ymax - ymin) * i) / data.value.figure.heatmap.height)
         return [
           {
             z: reshape(data.value.figure.heatmap.data, data.value.figure.heatmap.width),
+            x: x,
+            y: y,
             type: 'heatmap' as const,
+            colorscale: 'Blues',
+            reversescale: true,
           },
         ]
       case 'histogram':
diff --git a/where_fi/visualise/protos/figure_type/heatmap.proto b/where_fi/visualise/protos/figure_type/heatmap.proto
index 9a8be5a..7590e4d 100644
--- a/where_fi/visualise/protos/figure_type/heatmap.proto
+++ b/where_fi/visualise/protos/figure_type/heatmap.proto
@@ -6,4 +6,8 @@ message HeatmapData {
     uint32 width = 3;
     uint32 height = 4;
     string cmap = 5;
+    float x_min = 6;
+    float x_max = 7;
+    float y_min = 8;
+    float y_max = 9;
 }
diff --git a/where_fi/visualise/server/__init__.py b/where_fi/visualise/server/__init__.py
index 83a2f4b..9dfaf90 100644
--- a/where_fi/visualise/server/__init__.py
+++ b/where_fi/visualise/server/__init__.py
@@ -12,6 +12,7 @@ import grpc
 import numpy as np
 import numpy.typing as npt
 
+from ...config import config
 from .generated import figure_pb2, figure_pb2_grpc
 from .generated.figure_type import heatmap_pb2, histogram_pb2, line_pb2
 
@@ -64,8 +65,8 @@ figures = {
     "aoa_heatmap": figure_pb2.Figure(
         uuid=str(uuid.uuid4()),
         title="AoA Heatmap",
-        x_label="ToF",
-        y_label="AoA",
+        x_label="AoA",
+        y_label="ToF",
     ),
 }
 
@@ -125,6 +126,10 @@ def add_data(dtype: DataType, new_data: npt.NDArray[np.complex128]) -> None:
                 data=new_data.flatten(),
                 width=new_data.shape[1],
                 height=new_data.shape[0],
+                x_min=0,
+                x_max=np.pi,
+                y_min=0,
+                y_max=config.music.heatmap.tof_max,
             )
             for q in clients.get(uuid, []):
                 q.put(figure_pb2.FigureData(uuid=uuid, heatmap=heatmap))