move preprocessing to CSIApplication

where_fi/application.py

@@ -6,11 +6,11 @@ from typing import Any, Callable, NamedTuple
 
 import numpy.typing as npt
 
-from where_fi.collection.protocols import CSIProducer
+from where_fi.collection import CSIMatrix, ingest
-
+from where_fi.collection.protocols import CSIProducer, MergedCSI
-from .collection import CSIMatrix, ingest
+from where_fi.config import config
-from .config import config
+from where_fi.processing.preprocess import Preprocessor
-from .visualise import server as visualise
+from where_fi.visualise import server as visualise
 
 
 class Receiver(NamedTuple):
@@ -68,7 +68,7 @@ class CSIApplication:
                         mostly used as a scheduler).
     """
 
-    def __init__(self, producer: CSIProducer) -> None:
+    def __init__(self, producer: CSIProducer, visualise_raw: bool = False) -> None:
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
 
         self.producer = producer
@@ -80,11 +80,15 @@ class CSIApplication:
         self.webapp_thread = threading.Thread(
             target=self.webapp.start, args=(self.webapp_queue,)
         )
+        self.visualise_raw = visualise_raw
 
-        self.csi_callback = None
         self.pre_merge_callback = None
+        self.raw_csi_callback = None
+        self.preprocessed_csi_callback = None
         self.processing_callback = None
 
+        self.preprocessor = Preprocessor()
+
     def visualise_data(
         self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
     ) -> None:
@@ -95,6 +99,17 @@ class CSIApplication:
         if not self.webapp_queue.full():
             self.webapp_queue.put(visualise.VisualiserData(data, dtype))
 
+    def on_raw(self, func: Callable[[CSIMatrix], None]) -> Callable[[CSIMatrix], None]:
+        """
+        Decorator to register a callback for the sample preprocessing.
+        """
+
+        def decorator(data: CSIMatrix) -> None:
+            func(data)
+
+        self.raw_csi_callback = decorator
+        return decorator
+
     def on_sample(
         self, func: Callable[[CSIMatrix], None]
     ) -> Callable[[CSIMatrix], None]:
@@ -106,7 +121,7 @@ class CSIApplication:
             self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
             func(data)
 
-        self.csi_callback = decorator
+        self.preprocessed_csi_callback = decorator
         return decorator
 
     def on_pre_merge(
@@ -117,13 +132,33 @@ class CSIApplication:
 
         return decorator
 
-    def on_process(self, func: Callable[[], None]) -> Callable[[], None]:
+    def on_process(
+        self, func: Callable[[CSIMatrix], None]
+    ) -> Callable[[CSIMatrix], None]:
         """
         Decorator to register a callback for the sample processing.
         """
         self.processing_callback = func
         return func
 
+    def process_sample(self, sample: MergedCSI) -> None:
+        """Data pipeline for processing a sample.
+        This is the entry point for a sample being received from the producer. Depending
+        on the callbacks which are registered, this will call the appropriate functions
+        to pre-process and visualise the sample.
+        """
+        if self.visualise_raw:
+            self.visualise_data(sample.matrix, visualise.figures.Figure.RAW_CSI)
+        if self.raw_csi_callback is not None:
+            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)
+        if self.preprocessed_csi_callback is not None:
+            self.preprocessed_csi_callback(processed)
+
     def listen(self) -> None:
         """
         Listen for incoming data from the selected ingestor, and call the appropriate
@@ -133,10 +168,7 @@ class CSIApplication:
         while True:
             try:
                 sample = next(data_gen)
-                if self.csi_callback is not None:
+                self.process_sample(sample)
-                    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:
@@ -153,9 +185,15 @@ class CSIApplication:
         self.webapp_thread.start()
 
         if self.processing_callback is not None:
-            self.scheduler = Scheduler(
+
-                config.processing_sample_rate, self.processing_callback
+            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()
         try:
             self.webapp_thread.join()

where_fi/cli/__init__.py

@@ -1,19 +1,17 @@
 import logging
-import multiprocessing as mp
-import threading
-import time
 from queue import Queue
-from typing import Any, cast
+from typing import cast
 
 import numpy as np
 import numpy.typing as npt
 import torch
 import typer
 
+from where_fi.collection import CSIMatrix
+
 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
 
@@ -41,18 +39,12 @@ def antennas() -> None:
 
 @cli.command()
 def heatmap() -> None:
-    app = CSIApplication(globals.csi_producer)
+    app = CSIApplication(globals.csi_producer, visualise_raw=True)
-    preprocessor = Preprocessor()
     aoa = AoA()
 
     @app.on_sample
-    def _(antenna_data: npt.NDArray[np.complex64]) -> None:
+    def _(sample: npt.NDArray[np.complex64]) -> None:
-        logger.info(f"Got final CSI data with shape {antenna_data.shape}")
+        processed_tensor = torch.tensor(sample, device=device)
-        app.visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI)
-        processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
-        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=app.visualise_data)
 
@@ -73,7 +65,6 @@ def phase_analysis(
     """
 
     app = CSIApplication(globals.csi_producer)
-    preprocessor = Preprocessor()
 
     if isinstance(subcarriers, int):
         subcarriers = [subcarriers]
@@ -84,33 +75,22 @@ def phase_analysis(
     }
 
     @app.on_sample
-    def _(antenna_data: npt.NDArray[np.complex64]) -> None:
+    def _(sample: CSIMatrix) -> 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]))
+            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 _() -> None:
+    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,
         )
 
-    # 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)
 
 
 cli.add_typer(file.app, name="file", help="Commands for working with CSI files")

where_fi/processing/preprocess.py

@@ -12,14 +12,12 @@ from ..config import config
 from ..visualise import server as visualise
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
 
 np.seterr(invalid="ignore")
 
 
 class Preprocessor:
     def __init__(self) -> None:
-        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(
@@ -29,6 +27,14 @@ class Preprocessor:
         #     btype="band",
         #     output="sos",
         # )
+        self._last_sample = None
+
+    @property
+    def last_sample(self) -> None | CSIMatrix:
+        """
+        The last sample of the preprocessor. This is used for low frequency processing
+        """
+        return self._last_sample
 
     def remove_sto(self, csi: CSIMatrix) -> CSIMatrix:
         """
@@ -36,6 +42,9 @@ class Preprocessor:
             - Sampling frequency offset
             - Packet detection delay
 
+        This is done by multiplying the CSI matrices of consecutive antennas in the
+        array.
+
         According to [1]:
             > Conjugate multiplication and division are the only two methods to
             > eliminate the SFO and PDD.
@@ -46,7 +55,7 @@ class Preprocessor:
         [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]
@@ -55,6 +64,35 @@ class Preprocessor:
             )
         return csi_remove_sto
 
+    def remove_sfo(
+        self,
+        csi: CSIMatrix,
+        visualiser: None
+        | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
+    ) -> CSIMatrix:
+        """
+        Remove sampling frequency offsets by linear regression.
+
+        This is caused by the difference in sampling frequency between the transmitter
+        and receiver, and this is linear in frequency. We can estimate it using linear
+        regression and compensate for it.
+        """
+        N_st, N_rx, _ = csi.shape
+        unwrapped = np.unwrap(np.angle(csi[:, :, 0]), axis=0).reshape(N_st, N_rx, 1)
+
+        if visualiser:
+            visualiser(unwrapped, visualise.figures.Figure.UNWRAPPED_PHASE)
+
+        X = np.vstack([np.arange(N_st), np.ones(N_st)]).T
+        for antenna in range(csi.shape[1]):
+            tau, rho = np.linalg.lstsq(X, unwrapped[:, antenna, 0])[0]
+            csi[:, antenna, 0] = np.abs(csi[:, antenna, 0]) * np.exp(
+                1j
+                * (np.angle(csi[:, antenna, 0]) - (tau * np.arange(csi.shape[0]) + rho))
+            )
+
+        return csi
+
     def preprocess(
         self,
         h: CSIMatrix,
@@ -76,24 +114,7 @@ class Preprocessor:
         # h_hat = correlate(h_hat, np.ones((3, 1, 1)) / 3)
         # h_hat = correlate(h_hat, [[[1 / 4]], [[1 / 2]], [[1 / 4]]], mode="valid")
 
-        # Unwrap phase and remove linear fit
+        h_hat = self.remove_sfo(h_hat, visualiser=visualiser)
-        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.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,
-                unwrapped[:, antenna, 0],
-            )[0]
-            h_hat[:, antenna, 0] = np.abs(h_hat[:, antenna, 0]) * np.exp(
-                1j
-                * (
-                    np.angle(h_hat[:, antenna, 0])
-                    - (tau * np.arange(h_hat.shape[0]) + rho)
-                )
-            )
 
         # Skip subcarrierss per config
         h_hat = h_hat[:: config.preprocessing.subcarrier_step, :, :]
@@ -109,6 +130,8 @@ class Preprocessor:
             + h_hat * config.preprocessing.moving_average_alpha
         )
 
+        self._last_sample = h_hat
+
         # Remove long term average, to remove static paths
         return h_hat
         h_hat -= self.long_term_avg