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,15 +45,15 @@ 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)

where_fi/processing/aoa.py

@@ -118,7 +118,7 @@ class AoA:
         theta: torch.Tensor,
         tof: torch.Tensor,
         visualiser: None
-        | Callable[[npt.NDArray[Any], visualise.DataType], None] = None,
+        | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
     ) -> torch.Tensor:
         R = torch.mean(self.historical_autocorr, dim=0)
 
@@ -127,7 +127,7 @@ class AoA:
         logger.debug(f"Calculating eigenvectors of R: {R.shape}")
         eigvals, eigvecs = torch.linalg.eig(R)
         if visualiser:
-            visualiser(eigvals.numpy(), visualise.DataType.EIGENVALUES)
+            visualiser(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
         assert isinstance(eigvals, torch.Tensor)
         assert isinstance(eigvecs, torch.Tensor)
         logger.debug(f"Eigenvalues: {eigvals}")
@@ -148,7 +148,7 @@ class AoA:
 
     def heatmap(
         self,
-        visualiser: None | Callable[[npt.NDArray[Any], visualise.DataType], None],
+        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
@@ -175,7 +175,7 @@ class AoA:
         ).numpy(force=True)
 
         if visualiser:
-            visualiser(heatmap, visualise.DataType.HEATMAP)
+            visualiser(heatmap, visualise.figures.Figure.AOA_HEATMAP)
         return heatmap
 
 

where_fi/visualise/server/__init__.py

@@ -2,19 +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 ...config import config
+from . import figures
 from .generated import figure_pb2, figure_pb2_grpc
-from .generated.figure_type import heatmap_pb2, histogram_pb2, line_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -22,72 +19,18 @@ 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
-    EIGENVALUES = 5
-
-
 @dataclass
 class VisualiserData:
     data: npt.NDArray[Any]
-    dtype: DataType
+    dtype: figures.Figure
-
-
-figures = {
-    "raw_phase": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="Raw CSI Phase",
-        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",
-        x_label="Subcarrier",
-        y_label="Unwrapped phase",
-    ),
-    "aoa_eig": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="AoA Eigenvalues",
-        x_label="Eigenvalue",
-        y_label="Frequency",
-        logx=True,
-    ),
-    "processed_phase": figure_pb2.Figure(
-        uuid=str(uuid.uuid4()),
-        title="Preprocessed CSI Phase",
-        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",
-        x_label="AoA",
-        y_label="ToF",
-    ),
-}
 
 
 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():
+        for figure_group in figures.Figure:
+            for figure in figures.all_figures[figure_group].figures:
                 yield figure
 
     def GetFigureUpdate(
@@ -114,62 +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:
+def add_data(dtype: figures.Figure, new_data: npt.NDArray[np.complex128]) -> None:
-    match dtype:
+    updates = figures.all_figures[dtype].update(new_data)
-        case DataType.RAW_CSI | DataType.PROCESSED_CSI:
+    for fig_id, update in updates.items():
-            for component, func in [("phase", np.angle), ("magn", np.abs)]:
+        for client in clients.get(fig_id, []):
-                uuid = (
+            client.put(update)
-                    figures[f"raw_{component}"].uuid
-                    if dtype == DataType.RAW_CSI
-                    else figures[f"processed_{component}"].uuid
-                )
-                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(
-                uuid=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,
-            )
-            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))
-        case DataType.EIGENVALUES:
-            uuid = figures["aoa_eig"].uuid
-            magn = np.abs(new_data)
-            bins = np.logspace(np.log10(magn.min()), np.log10(magn.max()), 10)
-            hist, _ = np.histogram(magn, bins=bins)
-            logger.info(f"bins: {bins}, hist: {hist}")
-            logger.info(
-                f"Updating histogram with {len(bins)} bins, and {len(hist)} bars"
-            )
-            logger.info(bins)
-            series = [histogram_pb2.HistogramSeries(data=hist)]
-            histogram = histogram_pb2.HistogramData(data=series, bins=bins)
-            for q in clients.get(uuid, []):
-                q.put(figure_pb2.FigureData(uuid=uuid, histogram=histogram))
 
 
 def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None:
@@ -179,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")
@@ -189,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()
+    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(),
+}