dissertation/where_fi/visualise/server/figures.py

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