refactor csi producers

where_fi/application.py

@@ -6,6 +6,8 @@ from typing import Any, Callable, NamedTuple
 
 import numpy.typing as npt
 
+from where_fi.collection.protocols import CSIProducer
+
 from .collection import CSIMatrix, ingest
 from .config import config
 from .visualise import server as visualise
@@ -66,32 +68,10 @@ class CSIApplication:
                         mostly used as a scheduler).
     """
 
-    def create_receivers(self) -> list[Receiver]:
+    def __init__(self, producer: CSIProducer) -> None:
-        """
-        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.producer = producer
-        self.transmitter = ingest.FeitTransmitter()
 
         self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
             config.processing_sample_rate
@@ -105,10 +85,6 @@ class CSIApplication:
         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:
@@ -148,12 +124,30 @@ class CSIApplication:
         self.processing_callback = func
         return func
 
+    def listen(self) -> None:
+        """
+        Listen for incoming data from the selected ingestor, and call the appropriate
+        callback.
+        """
+        data_gen = self.producer()
+        while True:
+            try:
+                sample = next(data_gen)
+                if self.csi_callback is not None:
+                    self.csi_callback(sample.matrix)
+                if self.pre_merge_callback is not None:
+                    self.pre_merge_callback(sample.frames)
+            except StopIteration:
+                break
+            except Exception as e:
+                self.logger.error(f"Error processing sample: {e}", exc_info=True)
+                break
+
     def start(self) -> None:
-        for receiver in self.receivers:
+        # for receiver in self.receivers:
-            receiver.thread.start()
+        #     receiver.thread.start()
         self.buffer_thread = threading.Thread(
-            target=self.buffer.process_forever,
+            target=self.listen,
-            args=(self.csi_callback, self.pre_merge_callback),
         )
         self.buffer_thread.start()
         self.webapp_thread.start()
@@ -170,16 +164,12 @@ class CSIApplication:
 
     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
+        self.producer.stop()
 
         if hasattr(self, "scheduler"):
             self.scheduler.stop()
             self.scheduler.thread.join()
-        self.buffer.active = False
         self.buffer_thread.join()
 
         self.webapp.active = False

where_fi/cli/__init__.py

@@ -1,8 +1,8 @@
 import logging
 import multiprocessing as mp
-from queue import Queue
 import threading
 import time
+from queue import Queue
 from typing import Any, cast
 
 import numpy as np
@@ -33,7 +33,7 @@ def antennas() -> None:
     """
     from ..utils import antenna_order
 
-    if not globals.is_live:
+    if not globals.csi_producer.is_live:
         raise NotImplementedError("This command only works with live data")
 
     antenna_order.main()
@@ -41,37 +41,28 @@ def antennas() -> None:
 
 @cli.command()
 def heatmap() -> None:
+    app = CSIApplication(globals.csi_producer)
     preprocessor = Preprocessor()
     aoa = AoA()
 
-    # Start webapp in background process
+    @app.on_sample
-    webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
+    def _(antenna_data: npt.NDArray[np.complex64]) -> None:
-        config.processing_sample_rate
-    )
-    webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
-    webapp.start()
-
-    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.figures.Figure.RAW_CSI)
+        app.visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI)
-        processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
+        processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
-        visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
+        app.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)
-        aoa.heatmap(visualiser=visualise_data)
+        aoa.heatmap(visualiser=app.visualise_data)
 
-    globals.csi_producer(csi_callback=callback)
+    app.start()
     logger.info("Finished processing CSI data")
 
 
 @cli.command()
 def phase_analysis(
-    subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0
+    subcarriers: list[int] = [0], rx_antenna: int = 0, tx_antenna: int = 0
 ) -> None:
     """
     Visualise the phase information in the CSI data received from the antennas.
@@ -81,26 +72,34 @@ def phase_analysis(
     selected subcarrier and antenna.
     """
 
-    app = CSIApplication()
+    app = CSIApplication(globals.csi_producer)
     preprocessor = Preprocessor()
 
-    subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate)
+    if isinstance(subcarriers, int):
+        subcarriers = [subcarriers]
+    print("Starting phase analysis on subcarriers: ", subcarriers)
+
+    subcarrier_phase: dict[int, Queue[float]] = {
+        x: Queue(config.collection_sample_rate) for x in subcarriers
+    }
 
     @app.on_sample
     def _(antenna_data: npt.NDArray[np.complex64]) -> None:
         processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
         app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
 
+        for subcarrier in subcarriers:
             phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna]))
-        if subcarrier_phase.full():
+            if subcarrier_phase[subcarrier].full():
-            subcarrier_phase.get()
+                subcarrier_phase[subcarrier].get()
-        subcarrier_phase.put(phase)
+            subcarrier_phase[subcarrier].put(phase)
 
     @app.on_process
     def _() -> None:
-        logger.info(f"Updating phase visualisation")
+        logger.info("Updating phase visualisation")
         app.visualise_data(
-            np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS
+            np.array([x.queue for x in subcarrier_phase.values()]),
+            visualise.figures.Figure.PHASE_ANALYSIS,
         )
 
     # globals.csi_producer(csi_callback=callback)

where_fi/cli/globals.py

@@ -1,18 +1,15 @@
-from functools import partial
 from pathlib import Path
 
 from .. import collection
 from ..collection import file, ingest
 
-csi_producer: collection.CSIProducer = collection.noop
+csi_producer: collection.CSIProducer = collection.NoopCSIProducer()
 is_live = True
 
 
 def main(from_file: Path | None = None) -> None:
     global csi_producer, is_live
     if from_file:
-        csi_producer = partial(file.start_processing, file_path=from_file)
+        csi_producer = file.FileCSIPRoducer(path=from_file)
-        is_live = False
     else:
-        csi_producer = ingest.start_processing
+        csi_producer = ingest.RealtimeCSIProducer()
-        is_live = True

where_fi/collection/__init__.py

@@ -1,14 +1,24 @@
+from typing import Iterator
+
 import numpy as np
 import numpy.typing as npt
 
 from . import file, ingest
-from .protocols import CSICallback, CSIProducer
+from .protocols import CSIProducer, MergedCSI
 
 
-def noop(csi_callback: CSICallback | None = None) -> None:
+class NoopCSIProducer:
-    del csi_callback
+    """A no-op CSI producer that does nothing"""
+
+    is_live = False
+
+    def __call__(self) -> Iterator[MergedCSI]:
+        return iter([])
+
+    def stop(self) -> None:
+        pass
 
 
 CSIMatrix = npt.NDArray[np.complex64]
 
-__all__ = ["file", "ingest", "noop", "CSIProducer", "CSICallback"]
+__all__ = ["file", "ingest", "NoopCSIProducer", "CSIProducer"]

where_fi/collection/file.py

@@ -2,6 +2,7 @@ import logging
 import time
 from datetime import datetime, timedelta
 from pathlib import Path
+from typing import Iterator
 
 import h5py
 
@@ -10,12 +11,20 @@ from . import protocols
 logger = logging.getLogger(__name__)
 
 
-def start_processing(
+class FileCSIPRoducer:
-    file_path: Path,
+    is_live = False
-    csi_callback: protocols.CSICallback | None = None,
+
-) -> None:
+    def __init__(self, path: Path) -> None:
-    logger.info(f"Replaying CSI data from {file_path}")
+        self.path = path
-    with h5py.File(file_path, "r") as file:
+        if not self.path.exists():
+            raise FileNotFoundError(f"File {self.path} does not exist")
+
+    def __call__(self) -> Iterator[protocols.MergedCSI]:
+        """
+        Replay the CSI data from the file.
+        """
+        logger.info(f"Replaying CSI data from {self.path}")
+        with h5py.File(self.path, "r") as file:
             try:
                 for key in file:
                     datetime.fromisoformat(key)
@@ -28,7 +37,13 @@ def start_processing(
             target_offset = datetime.now() - start_time
             for key in file:
                 logger.debug(f"Sending data from {key} at {datetime.now().isoformat()}")
-            csi_callback(file[key][:])
+
+                # TODO: add raw frames
+                yield protocols.MergedCSI(
+                    frames={},
+                    matrix=file[key][:],
+                )
+
                 curr_time_virtual = datetime.fromisoformat(key)
                 new_offset = datetime.now() - curr_time_virtual
                 logger.debug(f"New offset {new_offset}, target is {target_offset}")
@@ -37,3 +52,9 @@ def start_processing(
                         f"Data is {new_offset - target_offset} behind, lagging behind..."
                     )
                 time.sleep(max(0, (target_offset - new_offset).total_seconds()))
+
+    def stop(self) -> None:
+        """
+        Stop the producer.
+        """
+        logger.info("Stopping file CSI producer")

where_fi/collection/ingest.py

@@ -1,18 +1,18 @@
 import logging
 import multiprocessing as mp
+import selectors
 import socket
-import struct
 import subprocess
 import threading
 import time
 from datetime import datetime
-from typing import Callable, NamedTuple
+from typing import Iterator
 
 import numpy as np
 
 from ..config import config
 from .csi_frame import CSI
-from .protocols import CSICallback
+from .protocols import MergedCSI
 
 Host = tuple[str, int]
 
@@ -27,6 +27,7 @@ class FeitHost:
         self.active = True
         self.checker = threading.Thread(target=self.check_continuous)
         self.checker.start()
+        self.selectors: list[selectors.BaseSelector] = []
 
     def check_connection(self) -> bool:
         try:
@@ -46,6 +47,10 @@ class FeitHost:
         self.server.connect(self.host)
         self.server.send(b"stop\n")
         self.server.send(self.command.encode())
+
+        for selector in self.selectors:
+            selector.register(self.server, selectors.EVENT_READ, data=self)
+
         self.logger.info(f"Connected to {self.host}")
 
     def check_continuous(self) -> None:
@@ -83,40 +88,49 @@ class FeitTransmitter(FeitHost):
 
 
 class FeitReceiver(FeitHost):
-    def __init__(self, host: Host, queue: "mp.Queue[CSI]") -> None:
+    """
+    A class used to connect to the host running FeitCSI and receive CSI data over a UDP
+    socket.
+    """
+
+    def __init__(self, host: Host) -> 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) -> None:
+    def recv(self) -> CSI:
-        prev_time = datetime.now()
-        self.logger.info("Listening for CSI data")
-        while self.active:
-            while not hasattr(self, "server"):
-                time.sleep(0.1)
-            # This is the max size of a UDP packet. The size of the actual CSI
-            # packet will depend on the frame format and channel width, which
-            # changes the number of subcarriers
-            data = self.server.recv(65535)
-            try:
-                csidata = CSI(data)
-                self.logger.debug(
-                    f"Received CSI data after {datetime.now() - prev_time}"
-                )
-                prev_time = datetime.now()
-                self.queue.put(csidata)
-            except struct.error:
-                self.logger.error("Failed to parse CSI data")
-        self.logger.info("Stopping CSI receiver")
-
-
-class CSIAntennaJoin:
         """
+        Block until a CSI frame is received and return it.
+        """
+
+        # Receive the data from the socket
+        data = self.server.recv(65535)
+        # Decode the data using the CSI frame format
+        csi = CSI(data)
+        return csi
+
+    def register(self, selector: selectors.BaseSelector) -> None:
+        """
+        Register the receiver as a selector. This will allow us to multiplex the I/O
+        operations on a single thread.
+
+        This allows us to block until any of the receivers receive data
+        """
+        self.selectors.append(selector)
+
+        if hasattr(self, "server"):
+            selector.register(self.server, selectors.EVENT_READ, data=self)
+
+
+class CSIAntennaArray:
+    """
+    Represents a set of receiver hosts that are used to receive CSI data, assumed to be
+    part of a linear antenna array.
+
     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.
@@ -127,15 +141,12 @@ class CSIAntennaJoin:
         - `sample_callback`: Called with the merged data
     """
 
-    def __init__(
+    def __init__(self, receivers: list[FeitReceiver]) -> None:
-        self,
-        receiver_connections: dict[Host, "mp.Queue[CSI]"],
-    ) -> None:
         self.pending_data: dict[Host, tuple[datetime, CSI]] = {}
         self.pending_data_lock = mp.Lock()
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
         self.last_processed = datetime.now()
-        self.connections = receiver_connections
+        self.receivers = receivers
         self.active = True
 
     def add_data(self, host: Host, data: CSI) -> None:
@@ -159,16 +170,13 @@ class CSIAntennaJoin:
                 return False
         return True
 
-    def process_data(
+    def process_data(self) -> None | MergedCSI:
-        self,
+        if not self.is_ready():
-        sample_callback: CSICallback | None = None,
+            self.logger.debug("Not all data is ready")
-        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
+            return None
-    ) -> None:
         self.last_processed = datetime.now()
-        if pre_merge_callback is not None:
+
-            pre_merge_callback(
+        frames = {host: data[1] for host, data in self.pending_data.items()}
-                {host: data[1] for host, data in self.pending_data.items()}
-            )
         antenna_data = [
             np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
             for ip, antenna in config.antennas.order
@@ -176,63 +184,45 @@ class CSIAntennaJoin:
 
         # We have data from all servers
         all_data = np.concat(antenna_data, axis=1)
+        return MergedCSI(frames=frames, matrix=all_data)
 
-        if sample_callback:
+    def process_forever(self) -> Iterator[MergedCSI]:
-            sample_callback(all_data)
+        # Register the receivers with the selector
+        self.logger.info("Starting to process data from receivers")
+        selector = selectors.DefaultSelector()
+
+        for receiver in self.receivers:
+            receiver.register(selector)
 
-    def process_forever(
-        self,
-        callback: CSICallback | None = None,
-        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
-    ) -> None:
         while self.active:
-            for ip, queue in self.connections.items():
+            events = selector.select(timeout=0.1)
-                while not queue.empty():
+            if not events:
-                    self.add_data(ip, queue.get())
+                self.logger.warning("No events received in the last 0.1 seconds")
-            if self.is_ready():
+                continue
-                self.process_data(callback, pre_merge_callback=pre_merge_callback)
+
-            else:
+            for key, _ in events:
-                self.logger.debug("Not all data is ready")
+                receiver = key.data
-            time.sleep(0.0005)
+                self.add_data(receiver.host, receiver.recv())
+
+            sample = self.process_data()
+            if sample is not None:
+                yield sample
 
 
-class Receiver(NamedTuple):
+class RealtimeCSIProducer:
-    ip: Host
+    def __init__(self) -> None:
-    receiver: FeitReceiver
+        self.receivers = [FeitReceiver(ip) for ip in config.receive_hosts]
-    queue: "mp.Queue[CSI]"
+        self.transmitter = FeitTransmitter()
-
-
-def start_processing(
-    csi_callback: CSICallback | None = None,
-    pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
-) -> None:
-    receivers = [
-        Receiver(ip, FeitReceiver(ip), mp.Queue(config.collection_sample_rate))
-        for ip in config.receive_hosts
-    ]
-
-    transmitter = FeitTransmitter()
 
         # Start injecting CSI frames
-    buffer = CSIAntennaJoin({r.ip: r.queue for r in receivers})
+        self.buffer = CSIAntennaArray(self.receivers)
+        self.is_live = True
 
-    receiver_threads = [
+    def __call__(self) -> Iterator[MergedCSI]:
-        threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
+        return self.buffer.process_forever()
-    ]
 
-    for proc in receiver_threads:
+    def stop(self) -> None:
-        proc.start()
+        for r in self.receivers:
-
+            r.active = False
-    buffering_thread = threading.Thread(
+        self.transmitter.active = False
-        target=buffer.process_forever, args=(csi_callback, pre_merge_callback)
+        self.buffer.active = False
-    )
-    buffering_thread.start()
-    try:
-        while True:
-            time.sleep(100)
-    except KeyboardInterrupt:
-        for r in receivers:
-            r.receiver.active = False
-        transmitter.active = False
-        buffer.active = False
-        return

where_fi/collection/protocols.py

@@ -1,10 +1,22 @@
-from typing import Callable, Protocol
+from typing import Iterator, NamedTuple, Protocol
 
 import numpy as np
 import numpy.typing as npt
 
-CSICallback = Callable[[npt.NDArray[np.complex64]], None]
+from .csi_frame import CSI
+
+CSIHost = tuple[str, int]
+
+
+class MergedCSI(NamedTuple):
+    """Merged CSI data from all antennas"""
+
+    frames: dict[CSIHost, CSI]
+    matrix: npt.NDArray[np.complex64]
 
 
 class CSIProducer(Protocol):
-    def __call__(self, csi_callback: CSICallback | None = None) -> None: ...
+    is_live: bool
+
+    def __call__(self) -> Iterator[MergedCSI]: ...
+    def stop(self) -> None: ...