refactor figure visualisation

Split off figure-serving from the data transformations needed to make
the figures themselves (e.g. binning)

where_fi/cli/__init__.py

@@ -45,20 +45,19 @@ def heatmap() -> None:
     webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
     webapp.start()
 
-    def visualise_data(data: npt.NDArray[Any], dtype: visualise.DataType) -> None:
+    def visualise_data(data: npt.NDArray[Any], dtype: visualise.figures.Figure) -> None:
         if not webapp_queue.full():
             webapp_queue.put(visualise.VisualiserData(data, dtype))
 
     def callback(antenna_data: npt.NDArray[np.complex64]) -> None:
         logger.info(f"Got final CSI data with shape {antenna_data.shape}")
-        visualise_data(antenna_data, visualise.DataType.RAW_CSI)
+        visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI)
         processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
-        visualise_data(processed, visualise.DataType.PROCESSED_CSI)
+        visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
         logger.info(f"Processed CSI data with shape {processed.shape}")
         processed_tensor = torch.tensor(processed, device=device)
         aoa.update(processed_tensor)
-        heatmap = aoa.heatmap()
-        visualise_data(heatmap, visualise.DataType.HEATMAP)
+        aoa.heatmap(visualiser=visualise_data)
 
     globals.csi_producer(csi_callback=callback)
     logger.info("Finished processing CSI data")

where_fi/processing/aoa.py

@@ -1,11 +1,13 @@
 import logging
 from datetime import datetime
+from typing import Any, Callable
 
 import numpy as np
 import numpy.typing as npt
 import torch
 
 from ..config import config
+from ..visualise import server as visualise
 
 logger = logging.getLogger(__name__)
 
@@ -111,16 +113,24 @@ class AoA:
         assert steering.shape == (N, self.N_rx // 2, self.N_subcarriers // 2)
         return steering.reshape(N, -1)
 
-    def evaluate(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
+    def evaluate(
+        self,
+        theta: torch.Tensor,
+        tof: torch.Tensor,
+        visualiser: None
+        | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
+    ) -> torch.Tensor:
         R = torch.mean(self.historical_autocorr, dim=0)
 
         # The smallest eigenvectors span the noise subspace,
         # and the largest span the signal subspace.
         logger.debug(f"Calculating eigenvectors of R: {R.shape}")
         eigvals, eigvecs = torch.linalg.eig(R)
+        if visualiser:
+            visualiser(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
         assert isinstance(eigvals, torch.Tensor)
         assert isinstance(eigvecs, torch.Tensor)
-        logger.info(f"Eigenvalues: {eigvals}")
+        logger.debug(f"Eigenvalues: {eigvals}")
         E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]
 
         logger.debug(f"Signal subspace: {E_n.shape}")
@@ -136,7 +146,10 @@ class AoA:
         c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
         return torch.abs(c)[:, 0, 0]
 
-    def heatmap(self) -> npt.NDArray[np.float32]:
+    def heatmap(
+        self,
+        visualiser: None | Callable[[npt.NDArray[Any], visualise.figures.Figure], None],
+    ) -> npt.NDArray[np.float32]:
         thetas = np.linspace(
             0, np.pi, config.music.heatmap.theta_resolution, dtype=np.float32
         )
@@ -153,12 +166,16 @@ class AoA:
         evaluated = self.evaluate(
             torch.tensor(thetas_mesh.reshape(-1)),
             torch.tensor(tofs_mesh.reshape(-1)),
+            visualiser=visualiser,
         )
         logger.debug(f"Evaluated heatmap: {evaluated.shape}")
         heatmap: npt.NDArray[np.float32] = evaluated.reshape(
             config.music.heatmap.tof_resolution,
             config.music.heatmap.theta_resolution,
         ).numpy(force=True)
+
+        if visualiser:
+            visualiser(heatmap, visualise.figures.Figure.AOA_HEATMAP)
         return heatmap
 
 

where_fi/visualise/frontend/.gitignore

@@ -13,6 +13,7 @@ dist
 dist-ssr
 coverage
 *.local
+.vite/
 
 /cypress/videos/
 /cypress/screenshots/

where_fi/visualise/frontend/src/components/Figure.vue

@@ -36,12 +36,44 @@ 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':
+        const bin_start = data.value.figure.histogram.bins.slice(0, -1)
+        const bin_end = data.value.figure.histogram.bins.slice(1)
+        const bin_center = bin_start.map((start, i) => (start + bin_end[i]) / 2)
+        const bin_width = bin_start.map((start, i) => bin_end[i] - start)
+        console.log('Histogram data:', bin_center, bin_width)
+        console.log(
+          'Sizes:',
+          bin_center.length,
+          bin_width.length,
+          data.value.figure.histogram.data.length,
+        )
+        return data.value.figure.histogram.data.map((series) => ({
+          x: bin_center,
+          y: series.data,
+          width: bin_width,
+          type: 'bar' as const,
+        }))
     }
   }
   return []
@@ -49,8 +81,8 @@ const plotData = computed(() => {
 
 const layout = {
   title: { text: figure.title },
-  xaxis: { title: { text: figure.xLabel } },
-  yaxis: { title: { text: figure.yLabel } },
+  xaxis: { title: { text: figure.xLabel }, type: figure.logx ? 'log' : undefined },
+  yaxis: { title: { text: figure.yLabel }, type: figure.logy ? 'log' : undefined },
   height: 700,
 }
 

where_fi/visualise/protos/figure.proto

@@ -2,12 +2,16 @@ syntax = "proto3";
 
 import "figure_type/line.proto";
 import "figure_type/heatmap.proto";
+import "figure_type/histogram.proto";
 
 message Figure {
     string uuid = 1;
     string title = 2;
     string x_label = 3;
     string y_label = 4;
+
+    bool logx = 5;
+    bool logy = 6;
 }
 
 message FigureData {
@@ -15,6 +19,7 @@ message FigureData {
     oneof figure {
         LineChartData line = 2;
         HeatmapData heatmap = 3;
+        HistogramData histogram = 4;
     }
 }
 

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;
 }

where_fi/visualise/protos/figure_type/histogram.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+message HistogramSeries { repeated float data = 1; }
+
+message HistogramData {
+    string uuid = 1;
+    repeated HistogramSeries data = 2;
+    repeated float bins = 3;
+}

where_fi/visualise/server/__init__.py

@@ -2,18 +2,16 @@ import logging
 import multiprocessing as mp
 import queue
 import threading
-import uuid
 from concurrent import futures
 from dataclasses import dataclass
-from enum import Enum
 from typing import Any, Generator
 
 import grpc
 import numpy as np
 import numpy.typing as npt
 
+from . import figures
 from .generated import figure_pb2, figure_pb2_grpc
-from .generated.figure_type import heatmap_pb2, line_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -21,53 +19,19 @@ clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {}
 clients_lock = threading.Lock()
 
 
-class DataType(Enum):
-    RAW_CSI = 1
-    UNWRAPPED_PHASE = 2
-    PROCESSED_CSI = 3
-    HEATMAP = 4
-
-
 @dataclass
 class VisualiserData:
     data: npt.NDArray[Any]
-    dtype: DataType
-
-
-figures = {
-    "raw_phase": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="Raw CSI Phase",
-        x_label="Subcarrier",
-        y_label="Phase",
-    ),
-    "unwrapped_phase": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="Unwrapped CSI Phase",
-        x_label="Subcarrier",
-        y_label="Unwrapped phase",
-    ),
-    "processed_phase": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="Preprocessed CSI Phase",
-        x_label="Subcarrier",
-        y_label="Phase",
-    ),
-    "aoa_heatmap": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="AoA Heatmap",
-        x_label="ToF",
-        y_label="AoA",
-    ),
-}
+    dtype: figures.Figure
 
 
 class FigureServer(figure_pb2_grpc.FigureServiceServicer):
     def GetFigure(
         self, request: figure_pb2.FigureRequest, context: grpc.ServicerContext
     ) -> Generator[figure_pb2.Figure, None, None]:
-        for figure in figures.values():
-            yield figure
+        for figure_group in figures.Figure:
+            for figure in figures.all_figures[figure_group].figures:
+                yield figure
 
     def GetFigureUpdate(
         self, request: figure_pb2.FigureDataRequest, context: grpc.ServicerContext
@@ -93,43 +57,11 @@ class FigureServer(figure_pb2_grpc.FigureServiceServicer):
                 clients[request.uuid].remove(q)
 
 
-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 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(
-                uuid=uuid,
-                data=new_data.flatten(),
-                width=new_data.shape[1],
-                height=new_data.shape[0],
-            )
-            for q in clients.get(uuid, []):
-                q.put(figure_pb2.FigureData(uuid=uuid, heatmap=heatmap))
-            pass
-        case DataType.UNWRAPPED_PHASE:
-            uuid = figures["unwrapped_phase"].uuid
-            lines = [
-                line_pb2.LineChartData.Line(y=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))
+def add_data(dtype: figures.Figure, new_data: npt.NDArray[np.complex128]) -> None:
+    updates = figures.all_figures[dtype].update(new_data)
+    for fig_id, update in updates.items():
+        for client in clients.get(fig_id, []):
+            client.put(update)
 
 
 def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None:
@@ -139,8 +71,6 @@ def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None:
 
 
 def start(data_queue: "mp.Queue[VisualiserData]") -> None:
-    logger.info("Visualisation server shut down")
-
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
     figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server)
     server.add_insecure_port("[::]:50051")
@@ -149,4 +79,7 @@ def start(data_queue: "mp.Queue[VisualiserData]") -> None:
     server.start()
     logger.info("Server started")
     threading.Thread(target=listen_for_data, args=(data_queue,), daemon=True).start()
-    server.wait_for_termination()
+    try:
+        server.wait_for_termination()
+    except KeyboardInterrupt:
+        logger.info("Exiting visualisation server")

where_fi/visualise/server/figures.py

@@ -0,0 +1,190 @@
+import uuid
+from abc import ABC, abstractmethod
+from enum import Enum
+from functools import reduce
+from typing import Any, Callable, Sequence, cast
+
+import numpy as np
+import numpy.typing as npt
+
+from ...config import config
+from .generated import figure_pb2
+from .generated.figure_type import heatmap_pb2, histogram_pb2, line_pb2
+
+FigureUpdate = dict[str, figure_pb2.FigureData]
+
+
+class SpecificFigure(ABC):
+    """
+    This is a base class for all figures that can be visualised through the
+    visualisation server.
+    """
+
+    figures: Sequence[figure_pb2.Figure]
+
+    @abstractmethod
+    def __init__(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def update(self, new_data: npt.NDArray[np.complex128]) -> FigureUpdate:
+        """
+        Update the figure with new data.
+
+        The exact shape and format of the data passed as an argument will differ
+        depending on the exact figure being plotted.
+
+        The return value should be a dictionary with the UUID of the figure as the key
+        and the new data as the value.
+
+        This allows one class to update multiple figures at once (e.g. a figure plotting
+        the phase and amplitude of a signal).
+        """
+        raise NotImplementedError
+
+
+class SimpleLineChart:
+    """Helper class to create a simple line chart with one or more lines.
+
+    This allows generalising the creation of line charts, e.g. as in PerAntennaFigure.
+    """
+
+    def __init__(self, title: str, x_label: str, y_label: str) -> None:
+        self.figure = figure_pb2.Figure(
+            uuid=str(uuid.uuid4()),
+            title=title,
+            x_label=x_label,
+            y_label=y_label,
+        )
+
+    def update(
+        self, new_data: npt.NDArray[np.complex128], labels: Sequence[str]
+    ) -> FigureUpdate:
+        lines = [
+            line_pb2.LineChartData.Line(y=new_data[i], label=labels[i])
+            for i in range(new_data.shape[0])
+        ]
+        return {
+            self.figure.uuid: figure_pb2.FigureData(
+                uuid=self.figure.uuid,
+                line=line_pb2.LineChartData(lines=lines),
+            )
+        }
+
+
+class PerAntennaFigure(SpecificFigure):
+    """A figure that plots data for each antenna separately.
+
+    This allows creating multiple figures, each having one line per antenna.
+
+    Each figure can have a different function that is used to transform the data before
+    plotting. For example, can be used to generate plots for the phase and amplitude of
+    a signal.
+    """
+
+    def __init__(
+        self,
+        figures: Sequence[SimpleLineChart],
+        funcs: list[Callable[[npt.NDArray[Any]], npt.NDArray[Any]]],
+    ) -> None:
+        self.charts = figures
+        self.figures = [figure.figure for figure in figures]
+        self.funcs = funcs
+
+    def update(self, new_data: npt.NDArray[np.complex128]) -> FigureUpdate:
+        data_by_antenna = new_data[:, :, 0].T
+        antenna_labels = [f"Antenna {i}" for i in range(data_by_antenna.shape[0])]
+        updates = [
+            figure.update(func(data_by_antenna), labels=antenna_labels)
+            for func, figure in zip(self.funcs, self.charts, strict=True)
+        ]
+        return reduce((lambda a, b: a | b), updates)
+
+
+class MusicEigenvalueHistogram(SpecificFigure):
+    def __init__(self) -> None:
+        self.figure = figure_pb2.Figure(
+            uuid=str(uuid.uuid4()),
+            title="AoA Eigenvalues",
+            x_label="Eigenvalue",
+            y_label="Frequency of occurrence",
+            logx=True,
+        )
+
+        self.figures = [self.figure]
+
+    def update(self, new_data: npt.NDArray[np.complex128]) -> FigureUpdate:
+        magn = np.abs(new_data)
+
+        # The frontend library used for plotting doesn't support logarithmic
+        # binning[1], so we have to do it manually.
+        # Using base 10 log for the bins to make the plots easier to comprehend.
+        # [1] - https://github.com/plotly/plotly.js/issues/1844
+        bins = cast(
+            npt.NDArray[np.float32],
+            np.logspace(np.log10(magn.min()), np.log10(magn.max()), 10),
+        )
+        hist, _ = np.histogram(magn, bins=bins)
+
+        series = [histogram_pb2.HistogramSeries(data=hist)]
+        histogram = histogram_pb2.HistogramData(data=series, bins=bins)
+        return {self.figure.uuid: figure_pb2.FigureData(histogram=histogram)}
+
+
+class HeatmapFigure(SpecificFigure):
+    def __init__(self) -> None:
+        self.figure = figure_pb2.Figure(
+            uuid=str(uuid.uuid4()),
+            title="Angle of arrival Heatmap",
+            x_label="Angle of arrival",
+            y_label="Time of Flight",
+        )
+
+        self.figures = [self.figure]
+
+    def update(self, new_data: npt.NDArray[np.complex128]) -> FigureUpdate:
+        heatmap = heatmap_pb2.HeatmapData(
+            uuid=self.figure.uuid,
+            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,
+        )
+        return {self.figure.uuid: figure_pb2.FigureData(heatmap=heatmap)}
+
+
+class Figure(Enum):
+    RAW_CSI = 0
+    UNWRAPPED_PHASE = 1
+    PROCESSED_CSI = 2
+    MUSIC_EIGENVALUES = 3
+    AOA_HEATMAP = 4
+
+    def figure_class(self) -> SpecificFigure:
+        return all_figures[self]
+
+
+all_figures = {
+    Figure.RAW_CSI: PerAntennaFigure(
+        [
+            SimpleLineChart("Raw CSI Phase", "Subcarrier", "Phase"),
+            SimpleLineChart("Raw CSI Amplitude", "Subcarrier", "Amplitude"),
+        ],
+        [np.angle, np.abs],
+    ),
+    Figure.UNWRAPPED_PHASE: PerAntennaFigure(
+        [SimpleLineChart("Unwrapped CSI Phase", "Subcarrier", "Phase")], [lambda x: x]
+    ),
+    Figure.PROCESSED_CSI: PerAntennaFigure(
+        [
+            SimpleLineChart("Processed CSI Phase", "Subcarrier", "Phase"),
+            SimpleLineChart("Processed CSI Amplitude", "Subcarrier", "Amplitude"),
+        ],
+        [np.angle, np.abs],
+    ),
+    Figure.MUSIC_EIGENVALUES: MusicEigenvalueHistogram(),
+    Figure.AOA_HEATMAP: HeatmapFigure(),
+}