refactor figure visualisation

Split off figure-serving from the data transformations needed to make
the figures themselves (e.g. binning)
This commit is contained in:
Christos Falas 2025-03-04 17:30:26 +00:00
parent f6e10bedec
commit 8fe5995c5f
No known key found for this signature in database
4 changed files with 211 additions and 128 deletions

View File

@ -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)

View File

@ -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

View File

@ -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
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():
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:
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()
try:
server.wait_for_termination()
except KeyboardInterrupt:
logger.info("Exiting visualisation server")

View File

@ -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(),
}