refactor into separate CSIApplication framework

where_fi/__main__.py

@@ -2,11 +2,10 @@ import logging
 
 from . import cli
 
-
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s %(name)-50s %(levelname)-8s %(message)s",
 )
 
 if __name__ == "__main__":
-    cli.app()
+    cli.cli()

where_fi/application.py

@@ -0,0 +1,187 @@
+import logging
+import multiprocessing as mp
+import threading
+import time
+from typing import Any, Callable, NamedTuple
+
+import numpy.typing as npt
+
+from .collection import CSIMatrix, ingest
+from .config import config
+from .visualise import server as visualise
+
+
+class Receiver(NamedTuple):
+    ip: ingest.Host
+    receiver: ingest.FeitReceiver
+    thread: threading.Thread
+
+
+class Scheduler:
+    """
+    A simple scheduler that runs a function at a given rate.
+    """
+
+    def __init__(self, rate: float, func: Callable[[], None]) -> None:
+        self.rate = rate
+        self.func = func
+        self.active = True
+
+    def run(self) -> None:
+        """
+        Run the scheduler in a loop, calling the function at the given rate.
+        """
+        while self.active:
+            self.func()
+            time.sleep(1 / self.rate)
+
+    def start(self) -> None:
+        """
+        Start the scheduler in a separate thread.
+        """
+        self.thread = threading.Thread(target=self.run)
+        self.thread.start()
+
+    def stop(self) -> None:
+        """
+        Stop the scheduler.
+        """
+        self.active = False
+
+
+class CSIApplication:
+    """
+    A high-level application that manages CSI data ingestion and processing.
+
+    This allows a simple interface to be used for the main CLI logic, without having to
+    deal with the threads and queues directly.
+
+    It can be used using decortors to register callbacks for different stages of the
+    processing pipeline:
+
+        - `on_pre_merge`: Called with the raw CSI data (including headers) from each
+                          receiver, before it is merged into a single matrix.
+        - `on_sample`: Called once per sample with the merged CSI data.
+        - `on_process`: Called at the processing sample rate, without any data (this is
+                        mostly used as a scheduler).
+    """
+
+    def create_receivers(self) -> list[Receiver]:
+        """
+        Create a thread for each receiver, so that data can be received in parallel.
+
+        The receivers are created with a queue that is used to store the received data,
+        to be processed later. The receivers are created based on the configuration
+        file, which contains the IP addresses of the receivers.
+        """
+        feit_receivers = [
+            ingest.FeitReceiver(ip, mp.Queue(config.collection_sample_rate))
+            for ip in config.receive_hosts
+        ]
+        return [
+            Receiver(
+                ip,
+                receiver,
+                threading.Thread(target=receiver.listen),
+            )
+            for receiver, ip in zip(feit_receivers, config.receive_hosts, strict=True)
+        ]
+
+    def __init__(self) -> None:
+        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
+
+        self.receivers = self.create_receivers()
+        self.transmitter = ingest.FeitTransmitter()
+
+        self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
+            config.processing_sample_rate
+        )
+        self.webapp = visualise.Webapp()
+        self.webapp_thread = threading.Thread(
+            target=self.webapp.start, args=(self.webapp_queue,)
+        )
+
+        self.csi_callback = None
+        self.pre_merge_callback = None
+        self.processing_callback = None
+
+        self.buffer = ingest.CSIAntennaJoin(
+            {r.ip: r.receiver.queue for r in self.receivers}
+        )
+
+    def visualise_data(
+        self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
+    ) -> 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))
+
+    def on_sample(
+        self, func: Callable[[CSIMatrix], None]
+    ) -> Callable[[CSIMatrix], None]:
+        """
+        Decorator to register a callback for the sample preprocessing.
+        """
+
+        def decorator(data: CSIMatrix) -> None:
+            self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
+            func(data)
+
+        self.csi_callback = decorator
+        return decorator
+
+    def on_pre_merge(
+        self,
+    ) -> Callable[[Callable[[dict[ingest.Host, ingest.CSI]], None]], None]:
+        def decorator(func: Callable[[dict[ingest.Host, ingest.CSI]], None]) -> None:
+            self.pre_merge_callback = func
+
+        return decorator
+
+    def on_process(self, func: Callable[[], None]) -> Callable[[], None]:
+        """
+        Decorator to register a callback for the sample processing.
+        """
+        self.processing_callback = func
+        return func
+
+    def start(self) -> None:
+        for receiver in self.receivers:
+            receiver.thread.start()
+        self.buffer_thread = threading.Thread(
+            target=self.buffer.process_forever,
+            args=(self.csi_callback, self.pre_merge_callback),
+        )
+        self.buffer_thread.start()
+        self.webapp_thread.start()
+
+        if self.processing_callback is not None:
+            self.scheduler = Scheduler(
+                config.processing_sample_rate, self.processing_callback
+            )
+            self.scheduler.start()
+        try:
+            self.webapp_thread.join()
+        except KeyboardInterrupt:
+            self.stop()
+
+    def stop(self) -> None:
+        self.logger.info("Stopping application")
+        for receiver in self.receivers:
+            receiver.receiver.active = False
+            receiver.thread.join()
+
+        self.transmitter.active = False
+
+        if hasattr(self, "scheduler"):
+            self.scheduler.stop()
+            self.scheduler.thread.join()
+        self.buffer.active = False
+        self.buffer_thread.join()
+
+        self.webapp.active = False
+        self.webapp_thread.join()
+        self.logger.info("Application stopped")

where_fi/cli/__init__.py

@@ -1,25 +1,28 @@
 import logging
 import multiprocessing as mp
-from typing import Any, cast
 from queue import Queue
+import threading
+import time
+from typing import Any, cast
 
 import numpy as np
 import numpy.typing as npt
 import torch
 import typer
 
+from ..application import CSIApplication
 from ..config import config
 from ..processing.aoa import AoA
 from ..processing.preprocess import Preprocessor
 from ..visualise import server as visualise
 from . import file, globals
 
-app = typer.Typer(callback=globals.main)
+cli = typer.Typer(callback=globals.main)
 logger = logging.getLogger(__name__)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-@app.command()
+@cli.command()
 def antennas() -> None:
     """Utility to help determine the order in which antennas are plugged in
 
@@ -31,18 +34,20 @@ def antennas() -> None:
     from ..utils import antenna_order
 
     if not globals.is_live:
-        raise ValueError("This command only works with live data")
+        raise NotImplementedError("This command only works with live data")
 
     antenna_order.main()
 
 
-@app.command()
+@cli.command()
 def heatmap() -> None:
     preprocessor = Preprocessor()
     aoa = AoA()
 
     # Start webapp in background process
-    webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(config.sample_rate)
+    webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
+        config.processing_sample_rate
+    )
     webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
     webapp.start()
 
@@ -64,7 +69,7 @@ def heatmap() -> None:
     logger.info("Finished processing CSI data")
 
 
-@app.command()
+@cli.command()
 def phase_analysis(
     subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0
 ) -> None:
@@ -76,38 +81,37 @@ def phase_analysis(
     selected subcarrier and antenna.
     """
 
+    app = CSIApplication()
     preprocessor = Preprocessor()
 
-    # Start webapp in background process
+    subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate)
-    webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(config.sample_rate)
-    webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
-    webapp.start()
 
-    subcarrier_phase: Queue[float] = Queue(config.sample_rate)
+    @app.on_sample
-
+    def _(antenna_data: npt.NDArray[np.complex64]) -> None:
-    def visualise_data(data: npt.NDArray[Any], dtype: visualise.figures.Figure) -> None:
+        processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
-        if not webapp_queue.full():
+        app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
-            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.figures.Figure.RAW_CSI)
-
-        processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
-        visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
 
         phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna]))
-
         if subcarrier_phase.full():
             subcarrier_phase.get()
         subcarrier_phase.put(phase)
 
-        visualise_data(
+    @app.on_process
-            np.array(subcarrier_phase), visualise.figures.Figure.PHASE_ANALYSIS
+    def _() -> None:
+        logger.info(f"Updating phase visualisation")
+        app.visualise_data(
+            np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS
         )
 
-    globals.csi_producer(csi_callback=callback)
+    # globals.csi_producer(csi_callback=callback)
+    app.start()
     logger.info("Finished processing CSI data")
+    while threading.active_count() > 1:
+        names = [
+            t.name for t in threading.enumerate() if t != threading.current_thread()
+        ]
+        logger.info("Waiting for threads to close: " + ",".join(names))
+        time.sleep(2)
 
 
-app.add_typer(file.app, name="file", help="Commands for working with CSI files")
+cli.add_typer(file.app, name="file", help="Commands for working with CSI files")

where_fi/collection/__init__.py

@@ -1,3 +1,6 @@
+import numpy as np
+import numpy.typing as npt
+
 from . import file, ingest
 from .protocols import CSICallback, CSIProducer
 
@@ -6,4 +9,6 @@ def noop(csi_callback: CSICallback | None = None) -> None:
     del csi_callback
 
 
+CSIMatrix = npt.NDArray[np.complex64]
+
 __all__ = ["file", "ingest", "noop", "CSIProducer", "CSICallback"]

where_fi/collection/ingest.py

@@ -29,13 +29,17 @@ class FeitHost:
         self.checker.start()
 
     def check_connection(self) -> bool:
-        feitcsi_status = subprocess.run(
+        try:
-            f"ssh root@{self.host[0]} pgrep feitcsi",
+            feitcsi_status = subprocess.run(
-            check=False,
+                f"ssh root@{self.host[0]} pgrep feitcsi",
-            stdout=subprocess.DEVNULL,
+                check=False,
-            shell=True,
+                stdout=subprocess.DEVNULL,
-        )
+                shell=True,
-        return feitcsi_status.returncode == 0
+                timeout=1,
+            )
+            return feitcsi_status.returncode == 0
+        except subprocess.TimeoutExpired:
+            return False
 
     def connect(self) -> None:
         self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -54,6 +58,7 @@ class FeitHost:
         """
         last_status = False
         while self.active:
+            self.logger.debug(f"Checking connection to {self.host[0]}")
             if not self.check_connection():
                 self.logger.error(f"FeitCSI is not running on {self.host[0]}")
                 last_status = False
@@ -72,22 +77,23 @@ class FeitTransmitter(FeitHost):
             f"--channel-width {config.channel_width} "
             f"--format {config.frame_format} "
             f"--mode inject -s 1 --verbose "
-            f"--inject-delay {1_000_000 // config.sample_rate}"
+            f"--inject-delay {1_000_000 // config.collection_sample_rate}"
         )
         super().__init__(config.transmit_host, command)
 
 
 class FeitReceiver(FeitHost):
-    def __init__(self, host: Host) -> None:
+    def __init__(self, host: Host, queue: "mp.Queue[CSI]") -> None:
         command = (
             f"feitcsi --frequency {config.central_freq} "
             f"--channel-width {config.channel_width} "
             f"--format {config.frame_format} "
             f"--mode measure"
         )
+        self.queue = queue
         super().__init__(host, command)
 
-    def listen(self, queue: "mp.Queue[CSI]") -> None:
+    def listen(self) -> None:
         prev_time = datetime.now()
         self.logger.info("Listening for CSI data")
         while self.active:
@@ -103,13 +109,24 @@ class FeitReceiver(FeitHost):
                     f"Received CSI data after {datetime.now() - prev_time}"
                 )
                 prev_time = datetime.now()
-                queue.put(csidata)
+                self.queue.put(csidata)
             except struct.error:
                 self.logger.error("Failed to parse CSI data")
         self.logger.info("Stopping CSI receiver")
 
 
-class CSIProcessor:
+class CSIAntennaJoin:
+    """
+    This class is used to buffer a set of CSI data from multiple receivers, and once a
+    set of readings is ready for each receiver, it will be merged into a single sample
+    using the antenna order in the configuration.
+
+    When the data is ready, the following callbacks are called:
+        - `pre_merge_callback`: Called with the data before merging. This is useful for
+          per-antenna analysis, e.g. finding the correct antenna order
+        - `sample_callback`: Called with the merged data
+    """
+
     def __init__(
         self,
         receiver_connections: dict[Host, "mp.Queue[CSI]"],
@@ -144,7 +161,7 @@ class CSIProcessor:
 
     def process_data(
         self,
-        callback: CSICallback | None = None,
+        sample_callback: CSICallback | None = None,
         pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
     ) -> None:
         self.last_processed = datetime.now()
@@ -160,26 +177,23 @@ class CSIProcessor:
         # We have data from all servers
         all_data = np.concat(antenna_data, axis=1)
 
-        if callback:
+        if sample_callback:
-            callback(all_data)
+            sample_callback(all_data)
 
     def process_forever(
         self,
         callback: CSICallback | None = None,
         pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
     ) -> None:
-        try:
+        while self.active:
-            while self.active:
+            for ip, queue in self.connections.items():
-                for ip, queue in self.connections.items():
+                while not queue.empty():
-                    while not queue.empty():
+                    self.add_data(ip, queue.get())
-                        self.add_data(ip, queue.get())
+            if self.is_ready():
-                if self.is_ready():
+                self.process_data(callback, pre_merge_callback=pre_merge_callback)
-                    self.process_data(callback, pre_merge_callback=pre_merge_callback)
+            else:
-                else:
+                self.logger.debug("Not all data is ready")
-                    self.logger.debug("Not all data is ready")
+            time.sleep(0.0005)
-                time.sleep(0.0005)
-        except KeyboardInterrupt:
-            self.logger.info("Exiting CSI processing")
 
 
 class Receiver(NamedTuple):
@@ -193,26 +207,26 @@ def start_processing(
     pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
 ) -> None:
     receivers = [
-        Receiver(ip, FeitReceiver(ip), mp.Queue(config.sample_rate))
+        Receiver(ip, FeitReceiver(ip), mp.Queue(config.collection_sample_rate))
         for ip in config.receive_hosts
     ]
 
     transmitter = FeitTransmitter()
 
     # Start injecting CSI frames
-    processor = CSIProcessor({r.ip: r.queue for r in receivers})
+    buffer = CSIAntennaJoin({r.ip: r.queue for r in receivers})
 
-    receiver_processes = [
+    receiver_threads = [
         threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
     ]
 
-    for proc in receiver_processes:
+    for proc in receiver_threads:
         proc.start()
 
-    processing_thread = mp.Process(
+    buffering_thread = threading.Thread(
-        target=processor.process_forever, args=(csi_callback, pre_merge_callback)
+        target=buffer.process_forever, args=(csi_callback, pre_merge_callback)
     )
-    processing_thread.start()
+    buffering_thread.start()
     try:
         while True:
             time.sleep(100)
@@ -220,5 +234,5 @@ def start_processing(
         for r in receivers:
             r.receiver.active = False
         transmitter.active = False
-        processor.active = False
+        buffer.active = False
         return

where_fi/config/models.py

@@ -43,7 +43,8 @@ class Config(BaseModel):
 
     antennas: Antennas
 
-    sample_rate: int
+    collection_sample_rate: int
+    processing_sample_rate: int
     central_freq: int
     channel_width: Literal[20, 40, 80, 160]
     frame_format: Literal["NOHT", "HT", "VHT", "HESU"]

where_fi/processing/preprocess.py

@@ -6,6 +6,8 @@ import numpy as np
 import numpy.typing as npt
 from scipy.signal import butter, correlate, sosfilt, sosfilt_zi
 
+from where_fi.collection import CSIMatrix
+
 from ..config import config
 from ..visualise import server as visualise
 
@@ -17,33 +19,56 @@ np.seterr(invalid="ignore")
 
 class Preprocessor:
     def __init__(self) -> None:
-        self.prev_entries: Queue[npt.NDArray[np.complex64]] = Queue(maxsize=100)
+        self.prev_entries: Queue[CSIMatrix] = Queue(maxsize=100)
         self.short_term_avg = np.zeros((1,), dtype=np.complex64)
         self.long_term_avg = np.zeros((1,), dtype=np.complex64)
-        self.filter = butter(
+        # self.filter = butter(
-            5,
+        #     5,
-            config.preprocessing.bandpass.bounds,
+        #     config.preprocessing.bandpass.bounds,
-            fs=config.sample_rate,
+        #     fs=config.sample_rate,
-            btype="band",
+        #     btype="band",
-            output="sos",
+        #     output="sos",
-        )
+        # )
+
+    def remove_sto(self, csi: CSIMatrix) -> CSIMatrix:
+        """
+        Remove sampling time offsets caused by:
+            - Sampling frequency offset
+            - Packet detection delay
+
+        According to [1]:
+            > Conjugate multiplication and division are the only two methods to
+            > eliminate the SFO and PDD.
+
+        No citation or explanation is provided, so not sure why/whether it works.
+        Something similar is also done in [2] without explanation.
+
+        [1] - https://tns.thss.tsinghua.edu.cn/wst/docs/sanitization
+        [2] - https://doi.org/10.1109/ICC51166.2024.10623053
+        """
+        # Conjugate multiplication
+        csi_remove_sto = np.zeros_like(csi)
+        for antenna in range(csi.shape[1]):
+            antenna_nxt = (antenna + 1) % csi.shape[1]
+            csi_remove_sto[:, antenna, :] = np.multiply(
+                csi[:, antenna, :], csi[:, antenna_nxt, :].conj()
+            )
+        return csi_remove_sto
 
     def preprocess(
         self,
-        h: npt.NDArray[np.complex64],
+        h: CSIMatrix,
         visualiser: None
-        | Callable[[npt.NDArray[Any], visualise.DataType], None] = None,
+        | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
-    ) -> npt.NDArray[np.complex64]:
+    ) -> CSIMatrix:
         # CSI data is not available for pilot subcarriers.
-        h_hat: npt.NDArray[np.complex64] = np.where(
+        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,
         )
 
-        # Skip subcarrierss per config
+        logger.debug(f"CSI shape: {h_hat.shape}")
-        h_hat = h_hat[:: config.preprocessing.subcarrier_step, :, :]
-        # logger.info(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)
@@ -52,12 +77,11 @@ class Preprocessor:
         # h_hat = correlate(h_hat, [[[1 / 4]], [[1 / 2]], [[1 / 4]]], mode="valid")
 
         # Unwrap phase and remove linear fit
-        print(h_hat.shape)
         unwrapped = np.unwrap(np.angle(h_hat[:, :, 0]), axis=0).reshape(
             h_hat.shape[0], h_hat.shape[1], 1
         )
         if visualiser:
-            visualiser(unwrapped, visualise.DataType.UNWRAPPED_PHASE)
+            visualiser(unwrapped, visualise.figures.Figure.UNWRAPPED_PHASE)
         for antenna in range(h_hat.shape[1]):
             tau, rho = np.linalg.lstsq(
                 np.vstack([np.arange(h_hat.shape[0]), np.ones(h_hat.shape[0])]).T,
@@ -71,7 +95,11 @@ class Preprocessor:
                 )
             )
 
-        return h_hat
+        # 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)
@@ -82,6 +110,7 @@ class Preprocessor:
         )
 
         # Remove long term average, to remove static paths
+        return h_hat
         h_hat -= self.long_term_avg
 
         # Apply bandpass filter to remove low and high frequency noise

where_fi/utils/antenna_order.py

@@ -11,11 +11,19 @@ AntennaIdentifier = tuple[ingest.Host, int]
 order: list[AntennaIdentifier] = []
 prev_unplugged: set[AntennaIdentifier] = set()
 long_antenna_average: dict[AntennaIdentifier, deque[int]] = {
-    (host, 0): deque(maxlen=15 * config.sample_rate) for host in config.receive_hosts
+    (host, 0): deque(maxlen=15 * config.collection_sample_rate)
-} | {(host, 1): deque(maxlen=15 * config.sample_rate) for host in config.receive_hosts}
+    for host in config.receive_hosts
+} | {
+    (host, 1): deque(maxlen=15 * config.collection_sample_rate)
+    for host in config.receive_hosts
+}
 short_antenna_average: dict[AntennaIdentifier, deque[int]] = {
-    (host, 0): deque(maxlen=2 * config.sample_rate) for host in config.receive_hosts
+    (host, 0): deque(maxlen=2 * config.collection_sample_rate)
-} | {(host, 1): deque(maxlen=2 * config.sample_rate) for host in config.receive_hosts}
+    for host in config.receive_hosts
+} | {
+    (host, 1): deque(maxlen=2 * config.collection_sample_rate)
+    for host in config.receive_hosts
+}
 
 RSSI_THRESHOLD = 10
 

where_fi/visualise/server/__init__.py

@@ -2,6 +2,7 @@ import logging
 import multiprocessing as mp
 import queue
 import threading
+import time
 from concurrent import futures
 from dataclasses import dataclass
 from typing import Any, Generator
@@ -13,11 +14,6 @@ import numpy.typing as npt
 from . import figures
 from .generated import figure_pb2, figure_pb2_grpc
 
-logger = logging.getLogger(__name__)
-
-clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {}
-clients_lock = threading.Lock()
-
 
 @dataclass
 class VisualiserData:
@@ -57,29 +53,48 @@ class FigureServer(figure_pb2_grpc.FigureServiceServicer):
                 clients[request.uuid].remove(q)
 
 
-def add_data(dtype: figures.Figure, new_data: npt.NDArray[np.complex128]) -> None:
+class Webapp:
-    updates = figures.all_figures[dtype].update(new_data)
+    def __init__(self) -> None:
-    for fig_id, update in updates.items():
+        self.logger = logging.getLogger(__name__)
-        for client in clients.get(fig_id, []):
+        self.active = True
-            client.put(update)
 
+        self.clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {}
+        self.clients_lock = threading.Lock()
 
-def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None:
+    def add_data(
-    while True:
+        self, dtype: figures.Figure, new_data: npt.NDArray[np.complex128]
-        data = data_queue.get()
+    ) -> None:
-        add_data(data.dtype, data.data)
+        updates = figures.all_figures[dtype].update(new_data)
+        for fig_id, update in updates.items():
+            for client in self.clients.get(fig_id, []):
+                client.put(update)
 
+    def listen_for_data(self, data_queue: "mp.Queue[VisualiserData]") -> None:
+        while self.active:
+            self.logger.debug("Listening for data")
+            try:
+                data = data_queue.get(timeout=0.5)
+                self.add_data(data.dtype, data.data)
+            except queue.Empty:
+                pass
 
-def start(data_queue: "mp.Queue[VisualiserData]") -> None:
+        while not data_queue.empty():
-    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+            data = data_queue.get()
-    figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server)
+
-    server.add_insecure_port("[::]:50051")
+    def start(self, data_queue: "mp.Queue[VisualiserData]") -> None:
-    server.add_insecure_port("0.0.0.0:50051")
+        server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
-    logger.info("Starting server on port 50051")
+        figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server)
-    server.start()
+        server.add_insecure_port("[::]:50051")
-    logger.info("Server started")
+        server.add_insecure_port("0.0.0.0:50051")
-    threading.Thread(target=listen_for_data, args=(data_queue,), daemon=True).start()
+        self.logger.info("Starting server on port 50051")
-    try:
+        server.start()
-        server.wait_for_termination()
+        self.logger.info("Server started")
-    except KeyboardInterrupt:
+
-        logger.info("Exiting visualisation server")
+        data_thread = threading.Thread(target=self.listen_for_data, args=(data_queue,))
+        data_thread.start()
+
+        while self.active:
+            time.sleep(1)
+            self.logger.debug("Server is running")
+        server.stop(0.5)
+        data_thread.join()