From 8fe5995c5f8c3a833ddc4a2848f3a29958284fae Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Tue, 4 Mar 2025 17:30:26 +0000 Subject: [PATCH] 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 | 6 +- where_fi/processing/aoa.py | 8 +- where_fi/visualise/server/__init__.py | 135 ++---------------- where_fi/visualise/server/figures.py | 190 ++++++++++++++++++++++++++ 4 files changed, 211 insertions(+), 128 deletions(-) create mode 100644 where_fi/visualise/server/figures.py diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py index eb65adb..131a52d 100644 --- a/where_fi/cli/__init__.py +++ b/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) diff --git a/where_fi/processing/aoa.py b/where_fi/processing/aoa.py index 841bd2a..a70a5c1 100644 --- a/where_fi/processing/aoa.py +++ b/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 diff --git a/where_fi/visualise/server/__init__.py b/where_fi/visualise/server/__init__.py index c8f98de..bf63216 100644 --- a/where_fi/visualise/server/__init__.py +++ b/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,73 +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 - EIGENVALUES = 5 - - @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", - ), - "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", - ), -} + 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 @@ -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: - match dtype: - case DataType.RAW_CSI | DataType.PROCESSED_CSI: - 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 - ) - 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 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: @@ -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() - server.wait_for_termination() + try: + server.wait_for_termination() + except KeyboardInterrupt: + logger.info("Exiting visualisation server") diff --git a/where_fi/visualise/server/figures.py b/where_fi/visualise/server/figures.py new file mode 100644 index 0000000..64f2a60 --- /dev/null +++ b/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(), +}