improve motion detection examples

examples/measurements.py

@@ -0,0 +1,110 @@
+from itertools import product
+from queue import Queue
+from typing import cast
+
+import numpy as np
+
+from where_fi.application import CSIApplication
+from where_fi.collection import CSIMatrix
+from where_fi.collection.ingest import RealtimeCSIProducer
+from where_fi.config import config
+from where_fi.visualise import server as visualise
+
+app = CSIApplication(visualise_raw=True)
+visualise.figures.all_figures["median"] = visualise.figures.RandomVariable(
+    "Median Phase"
+)
+visualise.figures.all_figures["median_magn"] = visualise.figures.RandomVariable(
+    "Median Magnitude"
+)
+visualise.figures.all_figures["denoised"] = visualise.figures.PerAntennaFigure(
+    [
+        visualise.figures.SimpleLineChart("Denoised CSI Phase", "Subcarrier", "Phase"),
+        visualise.figures.SimpleLineChart(
+            "Denoised CSI Amplitude", "Subcarrier", "Amplitude"
+        ),
+    ],
+    [np.angle, np.abs],
+)
+measurements: list[np.complex64] = []
+
+
+subcarriers = [0, 1]
+rx_antenna = [0, 1]
+tx_antenna = [0]
+
+subcarrier_phase: dict[tuple[int, int, int], Queue[float]] = {
+    x: Queue(config.collection_sample_rate)
+    for x in product(subcarriers, rx_antenna, tx_antenna)
+}
+subcarrier_magn: dict[tuple[int, int, int], Queue[float]] = {
+    x: Queue(config.collection_sample_rate)
+    for x in product(subcarriers, rx_antenna, tx_antenna)
+}
+
+proc_phase: dict[tuple[int, int, int], Queue[float]] = {
+    x: Queue(config.collection_sample_rate)
+    for x in product(subcarriers, rx_antenna, tx_antenna)
+}
+
+proc_magn: dict[tuple[int, int, int], Queue[float]] = {
+    x: Queue(config.collection_sample_rate)
+    for x in product(subcarriers, rx_antenna, tx_antenna)
+}
+
+
+@app.on_sample
+def _(sample: CSIMatrix) -> None:
+    for i in product(subcarriers, rx_antenna, tx_antenna):
+        phase = cast(float, np.angle(sample[i]))
+        magn = cast(float, np.abs(sample[i]))
+        if subcarrier_phase[i].full():
+            subcarrier_phase[i].get()
+        subcarrier_phase[i].put(phase)
+        if subcarrier_magn[i].full():
+            subcarrier_magn[i].get()
+        subcarrier_magn[i].put(magn)
+
+
+cnt = 0
+
+
+@app.on_process
+def _(proc: CSIMatrix) -> None:
+    global cnt
+    global proc_phase
+
+    for i in product(subcarriers, rx_antenna, tx_antenna):
+        phase = cast(float, np.angle(proc[i]))
+        magn = cast(float, np.abs(proc[i]))
+        if proc_phase[i].full():
+            proc_phase[i].get()
+        proc_phase[i].put(phase)
+        if proc_magn[i].full():
+            proc_magn[i].get()
+        proc_magn[i].put(magn)
+    print(sum(proc_phase[0, 0, 0].queue))
+    app.visualise_data(
+        np.array([x.queue for x in proc_phase.values()]),
+        "median",
+    )
+    app.visualise_data(
+        np.array([x.queue for x in proc_magn.values()]),
+        "median_magn",
+    )
+    app.visualise_data(
+        np.array([x.queue for x in subcarrier_phase.values()]),
+        visualise.figures.Figure.PHASE_ANALYSIS,
+    )
+    app.visualise_data(
+        np.array([x.queue for x in subcarrier_magn.values()]),
+        visualise.figures.Figure.MAGN_ANALYSIS,
+    )
+    app.visualise_data(proc, "denoised")
+    cnt += 1
+
+
+if __name__ == "__main__":
+    print("Starting app")
+    app.set_producer(RealtimeCSIProducer())
+    app.start()

examples/motion_detector.py

@@ -3,25 +3,48 @@ This example shows how to use the CSI framework to connect to a FeitCSI host and
 detect changes in the environment
 """
 
-from queue import Queue
-
 import numpy as np
 import numpy.typing as npt
 
 from where_fi.application import CSIApplication
 from where_fi.collection.ingest import RealtimeCSIProducer
+from where_fi.config import config
+from where_fi.visualise import server as visualise
 
 # Connect to a FeitCSI host
-producer = RealtimeCSIProducer()
-app = CSIApplication(producer)
+app = CSIApplication(visualise_raw=True)
 
+# Register visualisations
+app.register_figure(
+    "magn-diff",
+    visualise.figures.PerAntennaFigure(
+        [
+            visualise.figures.SimpleLineChart(
+                "Magnitude diff", "Subcarrier", "Magnitude"
+            ),
+        ],
+        [lambda x: x],
+    ),
+)
+app.register_figure(
+    "phase-diff",
+    visualise.figures.PerAntennaFigure(
+        [
+            visualise.figures.SimpleLineChart("Phase diff", "Subcarrier", "Phase"),
+        ],
+        [lambda x: x],
+    ),
+)
+
+MAGN_THRESHOLD = 0.2
+PHASE_THRESHOLD = 0.02
+QUEUE_SIZE = 20
 
 # Stores historical data for each receiving antenna, for each subcarrier
-historical: dict[tuple[int, int], Queue[np.complex64]] = {}
-
-MAGN_THRESHOLD = 20
-PHASE_THRESHOLD = 0.5
-QUEUE_SIZE = 2000
+historical = np.zeros(
+    (QUEUE_SIZE, config.subcarriers, config.antennas.count, 1), dtype=np.complex64
+)
+sample_position = 0
 
 
 @app.on_process
@@ -35,38 +58,27 @@ def _(sample: npt.NDArray[np.complex64]) -> None:
     It is used to detect changes in the environment caused by motion, by comparing each
     entry in the matrix with a moving average
     """
-    change = False
-    for antenna in range(sample.shape[1]):
-        for subcarrier in range(sample.shape[0]):
-            # Get the current subcarrier data
-            current = sample[subcarrier, antenna, 0]
+    global historical
+    global sample_position
 
-            # Get the historical data for this antenna and subcarrier
-            if (antenna, subcarrier) not in historical:
-                historical[(antenna, subcarrier)] = Queue(maxsize=QUEUE_SIZE)
-            historical_data = historical[(antenna, subcarrier)]
+    historical[sample_position] = sample
+    sample_position = (sample_position + 1) % QUEUE_SIZE
 
-            # Calculate the average of the historical data for this antenna and
-            # subcarrier
-            mean = np.mean(historical_data.queue)
+    mean = np.mean(historical, axis=0)
+    magn_diff = np.abs(mean - sample)
+    phase_diff = np.abs(np.angle(mean) - np.angle(sample))
 
-            # If we have enough historical data, compare it with the current data
-            if historical_data.full():
-                historical_data.get()
+    app.visualise_data(magn_diff, "magn-diff")
+    app.visualise_data(phase_diff, "phase-diff")
 
-            # Add the current sample to the historical data
-            historical_data.put(current)
-
-            # Compare the current data with the historical data
-            if (
-                np.abs(mean - current) > MAGN_THRESHOLD
-                or np.abs(np.angle(mean) - np.angle(current)) > PHASE_THRESHOLD
-            ):
-                change = True
-    if change:
+    if (magn_diff > MAGN_THRESHOLD).any() or (phase_diff > PHASE_THRESHOLD).any():
         print("Motion detected!")
     else:
         print("No motion detected!")
 
 
-app.start()
+if __name__ == "__main__":
+    # Start the application
+    print("Starting app")
+    app.set_producer(RealtimeCSIProducer())
+    app.start()

examples/motion_detector_home_assistant.py

@@ -1,24 +1,51 @@
 """Motion Detection example
 This example shows how to use the CSI framework to connect to a FeitCSI host and
 detect changes in the environment
-
-Once motion is detected, the application will log the change to Home Assistant
-through an HTTP request.
 """
 
-import os
-from queue import Queue
-
 import numpy as np
 import numpy.typing as npt
 import requests
 
 from where_fi.application import CSIApplication
 from where_fi.collection.ingest import RealtimeCSIProducer
+from where_fi.config import config
+from where_fi.visualise import server as visualise
 
 # Connect to a FeitCSI host
-producer = RealtimeCSIProducer()
-app = CSIApplication(producer)
+app = CSIApplication(visualise_raw=True)
+
+# Register visualisations
+app.register_figure(
+    "magn-diff",
+    visualise.figures.PerAntennaFigure(
+        [
+            visualise.figures.SimpleLineChart(
+                "Magnitude diff", "Subcarrier", "Magnitude"
+            ),
+        ],
+        [lambda x: x],
+    ),
+)
+app.register_figure(
+    "phase-diff",
+    visualise.figures.PerAntennaFigure(
+        [
+            visualise.figures.SimpleLineChart("Phase diff", "Subcarrier", "Phase"),
+        ],
+        [lambda x: x],
+    ),
+)
+
+MAGN_THRESHOLD = 0.2
+PHASE_THRESHOLD = 0.02
+QUEUE_SIZE = 20
+
+# Stores historical data for each receiving antenna, for each subcarrier
+historical = np.zeros(
+    (QUEUE_SIZE, config.subcarriers, config.antennas.count, 1), dtype=np.complex64
+)
+sample_position = 0
 
 
 class HomeAssistantBinarySensor:
@@ -49,18 +76,13 @@ class HomeAssistantBinarySensor:
         self.state = state
         data = {
             "state": "on" if state else "off",
-            "attributes": {"friendly_name": self.name},
+            "attributes": {"friendly_name": self.name, "device_class": "motion"},
         }
+        print(f"Updating sensor {self.name} to {data}")
         requests.post(self.url, json=data, headers=self.headers)
 
 
-# Stores historical data for each receiving antenna, for each subcarrier
-historical: dict[tuple[int, int], Queue[np.complex64]] = {}
-sensor = HomeAssistantBinarySensor("motion_detector", "Motion Detector")
-
-MAGN_THRESHOLD = 20
-PHASE_THRESHOLD = 0.5
-QUEUE_SIZE = 2000
+sensor = HomeAssistantBinarySensor("motion_detector", "Room Motion Detector")
 
 
 @app.on_process
@@ -74,36 +96,28 @@ def _(sample: npt.NDArray[np.complex64]) -> None:
     It is used to detect changes in the environment caused by motion, by comparing each
     entry in the matrix with a moving average
     """
-    change = False
-    for antenna in range(sample.shape[1]):
-        for subcarrier in range(sample.shape[0]):
-            # Get the current subcarrier data
-            current = sample[subcarrier, antenna, 0]
+    global sample_position
 
-            # Get the historical data for this antenna and subcarrier
-            if (antenna, subcarrier) not in historical:
-                historical[(antenna, subcarrier)] = Queue(maxsize=QUEUE_SIZE)
-            historical_data = historical[(antenna, subcarrier)]
+    historical[sample_position] = sample
+    sample_position = (sample_position + 1) % QUEUE_SIZE
 
-            # Calculate the average of the historical data for this antenna and
-            # subcarrier
-            mean = np.mean(historical_data.queue)
+    mean = np.mean(historical, axis=0)
+    magn_diff = np.abs(mean - sample)
+    phase_diff = np.abs(np.angle(mean) - np.angle(sample))
 
-            # If we have enough historical data, compare it with the current data
-            if historical_data.full():
-                historical_data.get()
+    app.visualise_data(magn_diff, "magn-diff")
+    app.visualise_data(phase_diff, "phase-diff")
 
-            # Add the current sample to the historical data
-            historical_data.put(current)
-
-            # Compare the current data with the historical data
-            if (
-                np.abs(mean - current) > MAGN_THRESHOLD
-                or np.abs(np.angle(mean) - np.angle(current)) > PHASE_THRESHOLD
-            ):
-                change = True
-
-    sensor.update(change)
+    if (magn_diff > MAGN_THRESHOLD).any() or (phase_diff > PHASE_THRESHOLD).any():
+        print("Motion detected!")
+        sensor.update(True)
+    else:
+        print("No motion detected!")
+        sensor.update(False)
 
 
-app.start()
+if __name__ == "__main__":
+    # Start the application
+    print("Starting app")
+    app.set_producer(RealtimeCSIProducer())
+    app.start()

examples/music.py

@@ -0,0 +1,63 @@
+import logging
+
+import numpy as np
+import torch
+
+from where_fi.application import CSIApplication
+from where_fi.collection import CSIMatrix
+from where_fi.collection.ingest import RealtimeCSIProducer
+from where_fi.config import config
+from where_fi.processing import aoa
+
+app = CSIApplication(visualise_raw=True)
+
+logging.basicConfig(level=logging.INFO)
+
+T = 500
+
+N_sub1 = config.antennas.count // 2
+N_sub2 = config.subcarriers // 2
+L2 = config.subcarriers - N_sub2 + 1
+L1 = config.antennas.count - N_sub1 + 1
+N_sensors = config.subcarriers * config.antennas.count
+historical = torch.zeros(T, N_sensors, N_sensors, dtype=torch.complex64)
+
+cnt = 0
+
+
+@app.on_sample
+def _(sample: CSIMatrix) -> None:
+    global cnt
+    sample = sample.T.reshape(-1, 1)
+    sample_tensor = torch.tensor(sample)
+    historical[cnt] = sample_tensor @ torch.conj(sample_tensor).T
+    cnt = (cnt + 1) % T
+
+
+aoa = aoa.AoA()
+
+
+@app.on_process
+def _(_: CSIMatrix) -> None:
+    R: torch.Tensor = torch.mean(historical, axis=0)
+    Rss = torch.zeros(N_sub1 * N_sub2, N_sub1 * N_sub2, dtype=torch.complex64)
+    for i in range(L1):
+        for j in range(L2):
+            Rss += R[i : i + N_sub1 * N_sub2, j : j + N_sub1 * N_sub2]
+    Rss /= L1 * L2
+
+    aoa.historical_autocorr = torch.unsqueeze(Rss, 0)
+    aoa.heatmap(app.visualise_data)
+    # eigvals, eigvecs = torch.linalg.eig(Rss)
+    # app.visualise_data(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
+    # E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]
+
+    # print(E_n)
+    # c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
+    # return torch.abs(c)[:, 0, 0]
+
+
+if __name__ == "__main__":
+    producer = RealtimeCSIProducer()
+    app.set_producer(producer)
+    app.start()

where_fi/application.py

@@ -6,12 +6,13 @@ from typing import Any, Callable, NamedTuple
 
 import numpy.typing as npt
 
-from where_fi.collection import CSIMatrix, ingest
+from where_fi.collection import CSIMatrix, NoopCSIProducer, ingest
 from where_fi.collection.csi_frame import CSI
 from where_fi.collection.protocols import CSIProducer, MergedCSI
 from where_fi.config import config
 from where_fi.processing.preprocess import Preprocessor
 from where_fi.visualise import server as visualise
+from where_fi.visualise.server import figures
 
 
 class Receiver(NamedTuple):
@@ -69,10 +70,10 @@ class CSIApplication:
                         mostly used as a scheduler).
     """
 
-    def __init__(self, producer: CSIProducer, visualise_raw: bool = False) -> None:
+    def __init__(self, visualise_raw: bool = False) -> None:
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
 
-        self.producer = producer
+        self.producer = NoopCSIProducer()
 
         self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
             config.processing_sample_rate
@@ -89,16 +90,33 @@ class CSIApplication:
         self.processing_callback = None
 
         self.preprocessor = Preprocessor()
+        self.custom_figures: dict[
+            visualise.figures.FigureId, visualise.figures.SpecificFigure
+        ] = {}
+
+    def set_producer(self, producer: CSIProducer) -> None:
+        """
+        Set the producer for the application. This is used to change the data source
+        at runtime.
+        """
+        self.producer = producer
+
+    def register_figure(
+        self,
+        figure_id: visualise.figures.FigureId,
+        figure: visualise.figures.SpecificFigure,
+    ) -> None:
+        self.custom_figures[figure_id] = figure
+        figures.all_figures[figure_id] = figure
 
     def visualise_data(
-        self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
+        self, data: npt.NDArray[Any], dtype: visualise.figures.FigureId
     ) -> None:
         """
         Update a visualisation with the given data. The available visualisations are as
         per visualise.server.all_figures.
         """
-        if not self.webapp_queue.full():
-            self.webapp_queue.put(visualise.VisualiserData(data, dtype))
+        self.webapp_queue.put(visualise.VisualiserData(data, dtype))
 
     def on_raw(self, func: Callable[[CSIMatrix], None]) -> Callable[[CSIMatrix], None]:
         """
@@ -119,7 +137,6 @@ class CSIApplication:
         """
 
         def decorator(data: CSIMatrix) -> None:
-            self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
             func(data)
 
         self.preprocessed_csi_callback = decorator
@@ -155,9 +172,11 @@ class CSIApplication:
             self.raw_csi_callback(sample.matrix)
         if self.pre_merge_callback is not None:
             self.pre_merge_callback(sample.frames)
-        processed = self.preprocessor.preprocess(sample.matrix)
-        if self.visualise_raw:
-            self.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
+        processed = self.preprocessor.preprocess(
+            sample.matrix, sample.frames, visualiser=self.visualise_data
+        )
+        # if self.visualise_raw:
+        #     self.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
         if self.preprocessed_csi_callback is not None:
             self.preprocessed_csi_callback(processed)
 
@@ -186,17 +205,21 @@ class CSIApplication:
         self.buffer_thread.start()
         self.webapp_thread.start()
 
-        if self.processing_callback is not None:
+        def get_proc_sample() -> None:
+            sample = self.preprocessor.last_sample
+            if self.visualise_raw and sample is not None:
+                self.visualise_data(
+                    sample,
+                    visualise.figures.Figure.PROCESSED_CSI,
+                )
+            if sample is not None and self.processing_callback is not None:
+                self.processing_callback(sample)
+            else:
+                self.logger.warning("No samples to process")
 
-            def get_proc_sample() -> None:
-                sample = self.preprocessor.last_sample
-                if sample is not None and self.processing_callback is not None:
-                    self.processing_callback(sample)
-                else:
-                    self.logger.warning("No samples to process")
+        self.scheduler = Scheduler(config.processing_sample_rate, get_proc_sample)
+        self.scheduler.start()
 
-            self.scheduler = Scheduler(config.processing_sample_rate, get_proc_sample)
-            self.scheduler.start()
         try:
             self.webapp_thread.join()
         except KeyboardInterrupt:

where_fi/cli/__init__.py

@@ -1,18 +1,15 @@
+import importlib.util
 import logging
-from queue import Queue
-from typing import cast
+from pathlib import Path
 
 import numpy as np
 import numpy.typing as npt
 import torch
 import typer
 
-from where_fi.collection import CSIMatrix
+from where_fi.application import CSIApplication
+from where_fi.processing.aoa import AoA
 
-from ..application import CSIApplication
-from ..config import config
-from ..processing.aoa import AoA
-from ..visualise import server as visualise
 from . import file, globals
 
 cli = typer.Typer(callback=globals.main)
@@ -53,44 +50,43 @@ def heatmap() -> None:
 
 
 @cli.command()
-def phase_analysis(
-    subcarriers: list[int] = [0], rx_antenna: int = 0, tx_antenna: int = 0
-) -> None:
+def run(file: Path, app_name: str = "app") -> None:
     """
-    Visualise the phase information in the CSI data received from the antennas.
+    Run the application with the given file.
 
-    The data goes through the same preprocessing steps as the heatmap command, but
-    instead of going through the AoA estimation, we simply analyse the phase of the
-    selected subcarrier and antenna.
+    Can be used for running arbitrary CSI applications with non-default data streams
+    (e.g. from file or environment simulation).
     """
+    print(file.absolute())
+    if not file.exists():
+        print(f"Error: File '{file}' does not exist.")
+        raise typer.Exit(1)
+    module_name = file.stem
 
-    app = CSIApplication(globals.csi_producer)
+    spec = importlib.util.spec_from_file_location(module_name, str(file))
+    if spec is None:
+        print(f"Could not load spec from {file}")
+        raise typer.Exit(1)
 
-    if isinstance(subcarriers, int):
-        subcarriers = [subcarriers]
-    print("Starting phase analysis on subcarriers: ", subcarriers)
+    module = importlib.util.module_from_spec(spec)
+    if spec.loader is None:
+        print(f"Could not load module from {file}")
+        raise typer.Exit(1)
+    try:
+        spec.loader.exec_module(module)
+    except Exception as e:
+        print(f"Failed to execute {file}: {e}")
+        raise typer.Exit(1) from e
 
-    subcarrier_phase: dict[int, Queue[float]] = {
-        x: Queue(config.collection_sample_rate) for x in subcarriers
-    }
+    if not hasattr(module, app_name):
+        print(f"Error: '{app_name}' not defined in the module.")
+        raise typer.Exit(1)
+    if not isinstance(module.app, CSIApplication):
+        print(f"Error: '{app_name}' is not a CSIApplication.")
+        raise typer.Exit(1)
 
-    @app.on_sample
-    def _(sample: CSIMatrix) -> None:
-        for subcarrier in subcarriers:
-            phase = cast(float, np.angle(sample[subcarrier, rx_antenna, tx_antenna]))
-            if subcarrier_phase[subcarrier].full():
-                subcarrier_phase[subcarrier].get()
-            subcarrier_phase[subcarrier].put(phase)
-
-    @app.on_process
-    def _(_: CSIMatrix) -> None:
-        logger.info("Updating phase visualisation")
-        app.visualise_data(
-            np.array([x.queue for x in subcarrier_phase.values()]),
-            visualise.figures.Figure.PHASE_ANALYSIS,
-        )
-
-    app.start()
+    module.app.producer = globals.csi_producer
+    module.app.start()
 
 
 cli.add_typer(file.app, name="file", help="Commands for working with CSI files")

where_fi/cli/globals.py

@@ -1,15 +1,20 @@
 from pathlib import Path
 
 from .. import collection
-from ..collection import file, ingest
+from ..collection import file, ingest, raytracing
 
 csi_producer: collection.CSIProducer = collection.NoopCSIProducer()
 is_live = True
 
 
-def main(from_file: Path | None = None) -> None:
+def main(from_file: Path | None = None, from_environment: Path | None = None) -> None:
     global csi_producer, is_live
+    if from_file and from_environment:
+        raise ValueError("Cannot specify both a data file and an environment file")
     if from_file:
         csi_producer = file.FileCSIPRoducer(path=from_file)
+    elif from_environment:
+        environment = raytracing.Environment.from_config(from_environment)
+        csi_producer = raytracing.SimulatedCSIProducer(environment)
     else:
         csi_producer = ingest.RealtimeCSIProducer()

where_fi/collection/__init__.py

@@ -3,7 +3,7 @@ from typing import Iterator
 import numpy as np
 import numpy.typing as npt
 
-from . import file, ingest
+from . import file, ingest, raytracing
 from .protocols import CSIProducer, MergedCSI
 
 
@@ -21,4 +21,4 @@ class NoopCSIProducer:
 
 CSIMatrix = npt.NDArray[np.complex64]
 
-__all__ = ["file", "ingest", "NoopCSIProducer", "CSIProducer"]
+__all__ = ["file", "ingest", "raytracing", "NoopCSIProducer", "CSIProducer"]

where_fi/collection/csi_frame.py

@@ -43,8 +43,8 @@ class CSIHeader:
         self.num_rx = data[46]
         self.num_tx = data[47]
         self.num_subcarriers = struct.unpack("I", data[52:56])[0]
-        self.rssi1 = struct.unpack("I", data[60:64])[0]
-        self.rssi2 = struct.unpack("I", data[64:68])[0]
+        self.rssi1: int = struct.unpack("I", data[60:64])[0]
+        self.rssi2: int = struct.unpack("I", data[64:68])[0]
         self.source_mac = struct.unpack("BBBBBB", data[68:74])
         self.source_mac_string = "%02x:%02x:%02x:%02x:%02x:%02x" % struct.unpack(
             "BBBBBB", data[68:74]

where_fi/collection/raytracing/__init__.py

@@ -0,0 +1,259 @@
+import logging
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Iterator
+
+import numpy as np
+import numpy.typing as npt
+
+from where_fi.collection.protocols import CSIProducer, MergedCSI
+from where_fi.config import config
+
+C = 299_792_458
+MAX_DEPTH = 2
+LOSS_EXPONENT = 0.8
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PathComponent:
+    delay: float
+    phase_offset: float
+    attenuation: float
+
+
+class ChannelImpulseResponse:
+    def __init__(self, path_components: list[PathComponent]) -> None:
+        self.path_components = path_components
+
+    @staticmethod
+    def delayed(
+        other: "ChannelImpulseResponse", delay: float, reflect: bool = False
+    ) -> "ChannelImpulseResponse":
+        new_path_components: list[PathComponent] = []
+        distance = delay * C
+
+        for component in other.path_components:
+            new_path_components.append(
+                PathComponent(
+                    delay=component.delay + delay,
+                    phase_offset=component.phase_offset + (np.pi if reflect else 0),
+                    attenuation=component.attenuation
+                    * (1 / (1 + distance) ** LOSS_EXPONENT),
+                    # attenuation=component.attenuation
+                    # * (np.exp(-distance * LOSS_EXPONENT)),
+                )
+            )
+        return ChannelImpulseResponse(new_path_components)
+
+    def __add__(self, other: "ChannelImpulseResponse") -> "ChannelImpulseResponse":
+        return ChannelImpulseResponse(self.path_components + other.path_components)
+
+
+class PathObject:
+    def __init__(self, x: float, y: float, z: float) -> None:
+        self.x = x
+        self.y = y
+        self.z = z
+        self.cir = ChannelImpulseResponse([])
+
+    def distance(self, other: "PathObject") -> float:
+        return (
+            (self.x - other.x) ** 2 + (self.y - other.y) ** 2 + (self.z - other.z) ** 2
+        ) ** 0.5
+
+    def reflect(
+        self,
+        incoming_cir: ChannelImpulseResponse,
+        target: list["PathObject"],
+        depth: int = 0,
+    ) -> None:
+        self.cir = self.cir + incoming_cir
+
+        if depth > MAX_DEPTH:
+            return
+
+        for obj in target:
+            if id(obj) == id(self):
+                continue
+            distance = self.distance(obj)
+            assert distance > 0
+            delay = distance / C
+            obj.reflect(
+                ChannelImpulseResponse.delayed(incoming_cir, delay, reflect=True),
+                target,
+                depth + 1,
+            )
+
+
+class Transmitter(PathObject):
+    def __init__(self, x: float, y: float, z: float) -> None:
+        super().__init__(x, y, z)
+
+
+DELTA_T = 10
+GAMMA = np.pi / 4
+
+
+class Receiver(PathObject):
+    def __init__(self, x: float, y: float, z: float, ideal: bool = True) -> None:
+        super().__init__(x, y, z)
+        self.delta_t = 0 if ideal else DELTA_T
+        self.gamma = 0 if ideal else GAMMA
+
+    def get_cfr(self) -> npt.NDArray[np.complex64]:
+        """
+        Calculate the CSI matrix for this rx-tx pair. This is computed by the Fourier
+        Transform of the Channel Impulse Response (CIR). The CIR is calculated as in
+        [1], [2].
+
+        To get the FT of the CIR, we use the sifting property of the Dirac delta
+
+        [1] - https://tns.thss.tsinghua.edu.cn/wst/docs/pre
+        [2] - https://dl.acm.org/doi/10.1145/2543581.2543592, Equation 4
+        """
+        cfr = np.zeros(len(config.subcarrier_frequencies), dtype=np.complex64)
+        for i_sub, f_sub in enumerate(config.subcarrier_frequencies):
+            cfr[i_sub] = sum(
+                [
+                    component.attenuation
+                    * np.exp(1j * component.phase_offset)
+                    * np.exp(-1j * 2 * np.pi * f_sub * component.delay)
+                    for component in self.cir.path_components
+                ]
+            ) * np.exp(
+                1j
+                * (
+                    2
+                    * np.pi
+                    * (i_sub / len(config.subcarrier_frequencies))
+                    * self.delta_t
+                    + self.gamma
+                )
+            )
+        return cfr
+
+
+class Environment:
+    def __init__(self, objects: list[PathObject]) -> None:
+        self.transmitters = [obj for obj in objects if isinstance(obj, Transmitter)]
+        self.receivers = [obj for obj in objects if isinstance(obj, Receiver)]
+        self.objects = [
+            obj for obj in objects if obj not in self.transmitters + self.receivers
+        ]
+        self.count = 0
+
+    def move(self) -> None:
+        self.count += 1
+        if self.count > 300:
+            print("Moving objects")
+            for obj in self.objects:
+                obj.x += np.sin(self.count / 100 * 2 * np.pi) * 0.1
+
+    @staticmethod
+    def from_config(filename: Path) -> "Environment":
+        """
+        Read a config file that includes a scene description and create an envionment
+        based on that
+
+        The config file should be a text file where each line corresponds to an object.
+        The first word of each line should be the type of object (TX/RX/OBJ), followed
+        by the x and y coordinates of the object.
+        """
+        objects: list[PathObject] = []
+        with open(filename, "r") as f:
+            for line in f.readlines():
+                if line.startswith("#"):
+                    continue
+                parts = line.split(" ")
+                x, y, z = map(float, parts[1:])
+                if parts[0] == "TX":
+                    objects.append(Transmitter(x, y, z))
+                elif parts[0] == "RX":
+                    objects.append(Receiver(x, y, z))
+                else:
+                    objects.append(PathObject(x, y, z))
+
+        return Environment(objects)
+
+    def add_awgn(
+        self, signal: npt.NDArray[np.complex64], snr_dB: float
+    ) -> npt.NDArray[np.complex64]:
+        signal_power = np.mean(np.abs(signal) ** 2)
+        snr_linear = 10 ** (snr_dB / 10)
+        noise_power = signal_power / snr_linear
+        noise = np.sqrt(noise_power / 2) * (
+            np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape)
+        )
+        return signal + noise.astype(np.complex64)
+
+    def get_csi(self) -> npt.NDArray[np.complex64]:
+        """
+        Calculate the Channel State Information (CSI) matrix for the simulated
+        environment.
+
+        The returned matrix is of shape (num_subcarriers, num_receivers,
+                                         num_transmitters)
+
+        This is calculated by finding all paths leading to each receiver, and
+        calculating the CFR evaluated at each subcarrier.
+        """
+        csi = np.zeros(
+            (
+                len(config.subcarrier_frequencies),
+                len(self.receivers),
+                len(self.transmitters),
+            ),
+            dtype=np.complex64,
+        )
+        for i_tx, transmitter in enumerate(self.transmitters):
+            for obj in self.objects + self.receivers + self.transmitters:
+                obj.cir = ChannelImpulseResponse([])
+            transmitter.reflect(
+                ChannelImpulseResponse(
+                    [
+                        PathComponent(
+                            delay=0,
+                            phase_offset=0,
+                            # attenuation=100000000,
+                            attenuation=100,
+                        )
+                    ]
+                ),
+                self.objects + self.receivers,
+            )
+            for i_rx, receiver in enumerate(self.receivers):
+                logger.debug(f"RX {i_rx} paths: {len(receiver.cir.path_components)}")
+                csi[:, i_rx, i_tx] = self.add_awgn(receiver.get_cfr(), snr_dB=30)
+        return csi
+
+
+class SimulatedCSIProducer(CSIProducer):
+    pass
+
+    def __init__(self, environment: Environment) -> None:
+        self.environment = environment
+        self.active = True
+
+    def __call__(self) -> Iterator[MergedCSI]:
+        while self.active:
+            self.environment.move()
+            start = datetime.now()
+            csi = self.environment.get_csi()
+            yield MergedCSI(
+                frames={},
+                matrix=csi,
+            )
+            time.sleep(
+                max(
+                    0,
+                    1 / config.collection_sample_rate
+                    - (datetime.now() - start).total_seconds(),
+                )
+            )
+
+    def stop(self) -> None:
+        self.active = False

where_fi/config/models.py

@@ -1,3 +1,4 @@
+from functools import cached_property
 from typing import Literal, Self
 
 from pydantic import BaseModel, model_validator
@@ -6,8 +7,6 @@ Host = tuple[str, int]
 
 
 class Preprocessing(BaseModel):
-    moving_average_alpha: float
-
     class Bandpass(BaseModel):
         lowcut: int
         highcut: int
@@ -16,12 +15,48 @@ class Preprocessing(BaseModel):
         def bounds(self) -> tuple[int, int]:
             return (self.lowcut, self.highcut)
 
-    bandpass: Bandpass
-    subcarrier_step: int
+    bandpass: Bandpass | None = None
+    subcarrier_step: int = 1
+    denoising: Literal["none", "median", "mean"] = "median"
+    denoising_period: float = 1
+
+    steps: list[
+        Literal[
+            "fill_pilots",
+            "skip_subcarriers",
+            "remove_agc",
+            "remove_sfo",
+            "remove_sto",
+            "bandpass",
+        ]
+    ] = ["fill_pilots", "skip_subcarriers", "remove_agc", "remove_sfo"]
+
+    @model_validator(mode="after")
+    def skip_in_steps(self) -> Self:
+        if "skip_subcarriers" not in self.steps and self.subcarrier_step != 1:
+            raise ValueError(
+                "subcarrier_step must be 1 if skip_subcarriers is not in steps"
+            )
+        return self
+
+    @model_validator(mode="after")
+    def bandpass_in_steps(self) -> Self:
+        if "bandpass" not in self.steps and self.bandpass is not None:
+            raise ValueError(
+                "bandpass must be in preprocessing steps if bandpass configuration is "
+                "provided"
+            )
+
+        if "bandpass" in self.steps and self.bandpass is None:
+            raise ValueError(
+                "bandpass configuration must be provided if bandpass is in "
+                "preprocessing steps"
+            )
+        return self
 
 
 class MUSIC(BaseModel):
-    eigval_threshold: int
+    eigval_threshold: float
     window_size: int
 
     class Heatmap(BaseModel):
@@ -36,6 +71,10 @@ class Antennas(BaseModel):
     spacing: float
     order: list[tuple[Host, int]]
 
+    @property
+    def count(self) -> int:
+        return len(self.order)
+
 
 class Config(BaseModel):
     receive_hosts: list[Host]
@@ -52,11 +91,11 @@ class Config(BaseModel):
     preprocessing: Preprocessing
     music: MUSIC
 
-    @property
+    @cached_property
     def central_freq_hz(self) -> int:
         return self.central_freq * 1_000_000
 
-    @property
+    @cached_property
     def band(self) -> Literal["2.4", "5", "6"]:
         if self.central_freq in range(2412, 2484):
             return "2.4"
@@ -66,11 +105,49 @@ class Config(BaseModel):
             return "6"
         raise ValueError(f"{self.central_freq} is not a valid Wi-Fi channel")
 
-    @property
-    def delta_f(self) -> int:
+    @cached_property
+    def num_guards(self) -> int:
+        if self.channel_width == 20:
+            return 7
+        return 11
+
+    @cached_property
+    def subcarriers(self) -> int:
+        nulls = {
+            20: 1,
+            40: 3,
+        }
+        return len(self.subcarrier_frequencies) + nulls[self.channel_width]
+
+    @cached_property
+    def _delta_f_no_skipping(self) -> int:
         if self.frame_format == "HESU":
-            return 78_125 * self.preprocessing.subcarrier_step
-        return 312_500 * self.preprocessing.subcarrier_step
+            return 78_125
+        return 312_500
+
+    @cached_property
+    def delta_f(self) -> int:
+        return self._delta_f_no_skipping * self.preprocessing.subcarrier_step
+
+    @cached_property
+    def _subcarriers_no_skipping(self) -> list[int]:
+        used = {
+            20: (1, 29),
+            40: (2, 59),
+        }
+        subcarrier_indices = list(
+            range(-used[self.channel_width][1] + 1, -used[self.channel_width][0] + 1)
+        ) + list(range(used[self.channel_width][0], used[self.channel_width][1]))
+        print(subcarrier_indices)
+
+        return [
+            self.central_freq_hz + i * self._delta_f_no_skipping
+            for i in subcarrier_indices
+        ]
+
+    @cached_property
+    def subcarrier_frequencies(self) -> list[int]:
+        return self._subcarriers_no_skipping[:: self.preprocessing.subcarrier_step]
 
     @model_validator(mode="after")
     def channels(self) -> Self:

where_fi/processing/aoa.py

@@ -18,8 +18,8 @@ torch.set_default_device(device)
 class AoA:
     def __init__(self) -> None:
         self.historical_autocorr = torch.tensor([], dtype=torch.complex64)
-        self.N_subcarriers = -1
-        self.N_rx = -1
+        self.N_subcarriers = config.subcarriers
+        self.N_rx = config.antennas.count
         self.timestamp = datetime.now()
         pass
 
@@ -84,7 +84,7 @@ class AoA:
             * np.pi
             * config.central_freq_hz
             * config.antennas.spacing
-            * (1 - torch.cos(theta))
+            * (torch.sin(theta))
             / 299_792_458
         )
         assert omega_t.shape == phi_theta.shape == (N,)
@@ -133,7 +133,7 @@ class AoA:
         logger.debug(f"Eigenvalues: {eigvals}")
         E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]
 
-        logger.debug(f"Signal subspace: {E_n.shape}")
+        logger.info(f"Signal subspace: {E_n.shape}")
         steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1)
         steering_h = torch.conj(steering).permute(0, 2, 1)
 

where_fi/processing/preprocess.py

@@ -1,5 +1,4 @@
 import logging
-from queue import Queue
 from typing import Any, Callable
 
 import numpy as np
@@ -7,8 +6,10 @@ import numpy.typing as npt
 from scipy.signal import butter, correlate, sosfilt, sosfilt_zi
 
 from where_fi.collection import CSIMatrix
+from where_fi.collection.csi_frame import CSI
+from where_fi.collection.protocols import CSIHost
+from where_fi.config import config
 
-from ..config import config
 from ..visualise import server as visualise
 
 logger = logging.getLogger(__name__)
@@ -20,21 +21,42 @@ class Preprocessor:
     def __init__(self) -> None:
         self.short_term_avg = np.zeros((1,), dtype=np.complex64)
         self.long_term_avg = np.zeros((1,), dtype=np.complex64)
-        # self.filter = butter(
-        #     5,
-        #     config.preprocessing.bandpass.bounds,
-        #     fs=config.sample_rate,
-        #     btype="band",
-        #     output="sos",
-        # )
         self._last_sample = None
+        self.denoising_samples = int(
+            config.preprocessing.denoising_period * config.collection_sample_rate
+        )
+        self.circular_buffer = np.zeros(
+            (
+                self.denoising_samples,  # Number of samples
+                config.subcarriers,  # Number of subcarriers
+                config.antennas.count,  # Number of RX antennas
+                1,  # Number of TX antennas
+            ),
+            dtype=np.complex64,
+        )
+        self.sample_index = 0
 
     @property
     def last_sample(self) -> None | CSIMatrix:
         """
         The last sample of the preprocessor. This is used for low frequency processing
         """
-        return self._last_sample
+        match config.preprocessing.denoising:
+            case "none":
+                return self._last_sample
+            case "median":
+                # Return the median of the last samples
+                median_abs = np.median(np.abs(self.circular_buffer), axis=0)
+                median_angle = np.median(np.angle(self.circular_buffer), axis=0)
+                ans = median_abs * np.exp(1j * median_angle)
+                return ans
+            case "mean":
+                # Return the mean of the last samples
+                return np.mean(self.circular_buffer, axis=0)
+            case _:
+                raise ValueError(
+                    f"Invalid denoising method: {config.preprocessing.denoising}"
+                )
 
     def remove_sto(self, csi: CSIMatrix) -> CSIMatrix:
         """
@@ -93,57 +115,119 @@ class Preprocessor:
 
         return csi
 
+    def remove_agc(self, csi: CSIMatrix, frames: dict[CSIHost, CSI]) -> CSIMatrix:
+        """
+        Normalise the magnitude of the CSI data to compensate for the effect of the
+        Automatic Gain Control (AGC) of the receiver
+
+        References:
+        [1] -
+        """
+        rssi = [
+            frames[host].header.rssi1 if antenna == 0 else frames[host].header.rssi2
+            for host, antenna in config.antennas.order
+        ]
+        rssi_linear = np.reshape(10 ** (np.array(rssi) / 10), (1, -1, 1))
+        csi_power = np.sum(np.abs(csi) ** 2)
+        return csi * np.sqrt(rssi_linear / csi_power)
+
+    def skip_subcarriers(self, csi: CSIMatrix) -> CSIMatrix:
+        """
+        Sometimes beccause of the large amount of processing, we need to skip some
+        subcarriers for the processing to be able to run in real time.
+
+        The subcarriers to skip are defined in the config file, per subcarrier_step. If
+        subcarrier_step is set to 1 (default), no subcarriers are skipped.
+        """
+        return csi[:: config.preprocessing.subcarrier_step, :, :]
+
+    def fill_pilots(self, csi: CSIMatrix) -> CSIMatrix:
+        """
+        Fill the pilot subcarriers with the average of the surrounding subcarriers.
+
+        This is done by averaging the subcarriers before and after the pilot
+        subcarriers. Pilots are detected by checking that the value is exactly 0.
+
+        Also, it adds placeholders for the middle null subcarriers.
+        """
+
+        num_middle = 1 if config.channel_width == 20 else 3
+        assert csi.shape[0] + num_middle == config.subcarriers
+        with_middle = np.zeros(
+            (csi.shape[0] + num_middle, csi.shape[1], csi.shape[2]), dtype=np.complex64
+        )
+        with_middle[: csi.shape[0] // 2, :, :] = csi[: csi.shape[0] // 2, :, :]
+        with_middle[csi.shape[0] // 2 + num_middle :, :, :] = csi[
+            csi.shape[0] // 2 :, :, :
+        ]
+        if num_middle == 3:
+            with_middle[csi.shape[0] // 2 + 1, :, :] = (
+                csi[csi.shape[0] // 2 - 1, :, :] + csi[csi.shape[0] // 2, :, :]
+            ) / 2
+        return np.where(
+            np.expand_dims(with_middle[:, 0, 0] == 0, axis=(1, 2)),
+            correlate(with_middle, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"),
+            with_middle,
+        )
+
+    def bandpass(self, csi: CSIMatrix) -> CSIMatrix:
+        """
+        Apply a Butterworth bandpass filter to the CSI data.
+
+        Remove low-frequency noise (caused by static paths) and high-frequency noise
+        (measurement variance).
+        """
+        if not hasattr(self, "filter"):
+            assert config.preprocessing.bandpass is not None
+            self.filter = butter(
+                5,
+                config.preprocessing.bandpass.bounds,
+                fs=config.collection_sample_rate,
+                btype="band",
+                output="sos",
+            )
+
+        if not hasattr(self, "filter_zi"):
+            self.filter_zi = (
+                np.expand_dims(sosfilt_zi(self.filter), axis=(-1, -2, -3)) * csi
+            )
+
+        h_hat_filt, self.filter_zi = sosfilt(
+            self.filter, [csi], zi=self.filter_zi, axis=0
+        )
+
+        return h_hat_filt[0]
+
     def preprocess(
         self,
         h: CSIMatrix,
+        frames: dict[CSIHost, CSI],
         visualiser: None
         | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
     ) -> CSIMatrix:
         # CSI data is not available for pilot subcarriers.
-        h_hat: CSIMatrix = np.where(
-            np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)),
-            correlate(h, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"),
-            h,
-        )
+        h_hat = h
+        for step in config.preprocessing.steps:
+            match step:
+                case "skip_subcarriers":
+                    h_hat = self.skip_subcarriers(h_hat)
+                case "remove_agc":
+                    h_hat = self.remove_agc(h_hat, frames)
+                case "remove_sfo":
+                    h_hat = self.remove_sfo(h_hat, visualiser=visualiser)
+                case "remove_sto":
+                    h_hat = self.remove_sto(h_hat)
+                case "fill_pilots":
+                    h_hat = self.fill_pilots(h_hat)
+                case "bandpass":
+                    h_hat = self.bandpass(h_hat)
 
         logger.debug(f"CSI shape: {h_hat.shape}")
 
-        # h_hat = np.multiply(h_hat, h_hat.conj() / abs(h_hat.conj()))
-        h_hat = np.nan_to_num(h_hat)
-
-        # h_hat = correlate(h_hat, np.ones((3, 1, 1)) / 3)
-        # h_hat = correlate(h_hat, [[[1 / 4]], [[1 / 2]], [[1 / 4]]], mode="valid")
-
-        h_hat = self.remove_sfo(h_hat, visualiser=visualiser)
-
-        # Skip subcarrierss per config
-        h_hat = h_hat[:: config.preprocessing.subcarrier_step, :, :]
-
-        # h_hat = self.remove_sto(h_hat)
-
-        # Assume that all csi matrices will have the same shape
-        if self.long_term_avg.shape != h_hat.shape:
-            self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex64)
-
-        self.long_term_avg = (
-            self.long_term_avg * (1 - config.preprocessing.moving_average_alpha)
-            + h_hat * config.preprocessing.moving_average_alpha
-        )
-
         self._last_sample = h_hat
 
-        # Remove long term average, to remove static paths
+        if config.preprocessing.denoising != "none":
+            self.circular_buffer[self.sample_index] = h_hat
+            self.sample_index = (self.sample_index + 1) % self.denoising_samples
+
         return h_hat
-        h_hat -= self.long_term_avg
-
-        # Apply bandpass filter to remove low and high frequency noise
-        if not hasattr(self, "filter_zi"):
-            self.filter_zi = (
-                np.expand_dims(sosfilt_zi(self.filter), axis=(-1, -2, -3)) * h_hat
-            )
-
-        h_hat_filt, self.filter_zi = sosfilt(
-            self.filter, [h_hat], zi=self.filter_zi, axis=0
-        )
-
-        return h_hat_filt[0]

where_fi/visualise/server/__init__.py

@@ -18,7 +18,7 @@ from . import figures
 @dataclass
 class VisualiserData:
     data: npt.NDArray[Any]
-    dtype: figures.Figure
+    dtype: figures.FigureId
 
 
 class FigureServer(figure_pb2_grpc.FigureServiceServicer):
@@ -30,8 +30,8 @@ class FigureServer(figure_pb2_grpc.FigureServiceServicer):
     def GetFigure(
         self, request: figure_pb2.FigureRequest, context: grpc.ServicerContext
     ) -> Generator[figure_pb2.Figure, None, None]:
-        for figure_group in figures.Figure:
-            for figure in figure_group.figure_class().figures:
+        for figure_group in figures.all_figures.values():
+            for figure in figure_group.figures:
                 yield figure
 
     def GetFigureUpdate(
@@ -65,8 +65,12 @@ class Webapp:
         self.figure_server = FigureServer()
 
     def add_data(
-        self, dtype: figures.Figure, new_data: npt.NDArray[np.complex128]
+        self, dtype: figures.FigureId, new_data: npt.NDArray[np.complex128]
     ) -> None:
+        self.logger.debug(f"Adding data to figure server {dtype}")
+        if dtype not in figures.all_figures:
+            self.logger.error(f"Figure {dtype} not found")
+            return
         updates = figures.all_figures[dtype].update(new_data)
         for fig_id, update in updates.items():
             for client in self.figure_server.clients.get(fig_id, []):
@@ -76,7 +80,7 @@ class Webapp:
         while self.active:
             self.logger.debug("Listening for data")
             try:
-                data = data_queue.get(timeout=0.5)
+                data = data_queue.get(timeout=0.1)
                 self.add_data(data.dtype, data.data)
             except queue.Empty:
                 pass

where_fi/visualise/server/figures.py

@@ -103,7 +103,28 @@ class PerSubcarrierFigure(SpecificFigure):
         return reduce((lambda a, b: a | b), updates)
 
 
-class PerAntennaFigure(SpecificFigure):
+class LabelledMultiLineChart(SpecificFigure):
+    def __init__(
+        self,
+        figures: Sequence[SimpleLineChart],
+        funcs: list[Callable[[npt.NDArray[Any]], npt.NDArray[Any]]],
+        labels: Sequence[str],
+    ) -> None:
+        self.charts = figures
+        self.figures = [figure.figure for figure in figures]
+        self.funcs = funcs
+        self.labels = labels
+
+    def update(self, new_data: npt.NDArray[np.complex128]) -> FigureUpdate:
+        data_by_antenna = new_data[:, :, 0].T
+        updates = [
+            figure.update(func(data_by_antenna), labels=self.labels)
+            for func, figure in zip(self.funcs, self.charts, strict=True)
+        ]
+        return reduce((lambda a, b: a | b), updates)
+
+
+class PerAntennaFigure(LabelledMultiLineChart):
     """A figure that plots data for each antenna separately.
 
     This allows creating multiple figures, each having one line per antenna.
@@ -118,18 +139,9 @@ class PerAntennaFigure(SpecificFigure):
         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)
+        super().__init__(
+            figures, funcs, [f"Antenna {i + 1}" for i in range(config.antennas.count)]
+        )
 
 
 class MusicEigenvalueHistogram(SpecificFigure):
@@ -271,12 +283,15 @@ class Figure(Enum):
     MUSIC_EIGENVALUES = 3
     AOA_HEATMAP = 4
     PHASE_ANALYSIS = 5
+    MAGN_ANALYSIS = 6
 
     def figure_class(self) -> SpecificFigure:
         return all_figures[self]
 
 
-all_figures = {
+FigureId = str | Figure
+
+all_figures: dict[FigureId, SpecificFigure] = {
     Figure.RAW_CSI: PerAntennaFigure(
         [
             SimpleLineChart("Raw CSI Phase", "Subcarrier", "Phase"),
@@ -297,4 +312,5 @@ all_figures = {
     Figure.MUSIC_EIGENVALUES: MusicEigenvalueHistogram(),
     Figure.AOA_HEATMAP: HeatmapFigure(),
     Figure.PHASE_ANALYSIS: RandomVariable("Phase"),
+    Figure.MAGN_ANALYSIS: RandomVariable("Magnitude"),
 }