add preprocessing step selection

where_fi/config/models.py

@@ -6,8 +6,6 @@ Host = tuple[str, int]
 
 
 class Preprocessing(BaseModel):
-    moving_average_alpha: float
-
     class Bandpass(BaseModel):
         lowcut: int
         highcut: int
@@ -16,11 +14,45 @@ class Preprocessing(BaseModel):
         def bounds(self) -> tuple[int, int]:
             return (self.lowcut, self.highcut)
 
-    bandpass: Bandpass
+    bandpass: Bandpass | None = None
-    subcarrier_step: int
+    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: float
@@ -72,9 +104,17 @@ class Config(BaseModel):
             return "6"
         raise ValueError(f"{self.central_freq} is not a valid Wi-Fi channel")
 
+    @property
+    def num_pilots(self) -> int:
+        if self.channel_width == 20:
+            return 8
+        if self.channel_width == 40:
+            return 14
+        return 24
+
     @property
     def subcarriers(self) -> int:
-        return int((self.channel_width * 1e6) // self.delta_f) - 8
+        return int((self.channel_width * 1e6) // self.delta_f) - self.num_pilots
 
     @property
     def _delta_f_no_skipping(self) -> int:

where_fi/processing/preprocess.py

@@ -1,5 +1,4 @@
 import logging
-from queue import Queue
 from typing import Any, Callable
 
 import numpy as np
@@ -22,13 +21,6 @@ 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
@@ -123,9 +115,7 @@ class Preprocessor:
 
         return csi
 
-    def normalise_magnitude(
+    def remove_agc(self, csi: CSIMatrix, frames: dict[CSIHost, CSI]) -> CSIMatrix:
-        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
@@ -141,6 +131,57 @@ class Preprocessor:
         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.
+        """
+        return np.where(
+            np.expand_dims(csi[:, 0, 0] == 0, axis=(1, 2)),
+            correlate(csi, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"),
+            csi,
+        )
+
+    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,
@@ -149,55 +190,28 @@ class Preprocessor:
         | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
     ) -> CSIMatrix:
         # CSI data is not available for pilot subcarriers.
-        h_hat: CSIMatrix = np.where(
+        h_hat = h
-            np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)),
+        for step in config.preprocessing.steps:
-            correlate(h, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"),
+            match step:
-            h,
+                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.normalise_magnitude(h_hat, frames)
-        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
-        )
-
-        # Remove long term average, to remove static paths
-        # h_hat -= self.long_term_avg
-
         self._last_sample = h_hat
 
         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
-
-        # 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]