unwrap CSI phase

where_fi/processing/preprocess.py

@@ -1,11 +1,13 @@
 import logging
 from queue import Queue
+from typing import Any, Callable
 
 import numpy as np
 import numpy.typing as npt
 from scipy.signal import butter, correlate, sosfilt, sosfilt_zi
 
 from ..config import config
+from ..visualise import server as visualise
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -26,9 +28,14 @@ class Preprocessor:
             output="sos",
         )
 
-    def preprocess(self, h: npt.NDArray[np.complex64]) -> npt.NDArray[np.complex64]:
+    def preprocess(
+        self,
+        h: npt.NDArray[np.complex64],
+        visualiser: None
+        | Callable[[npt.NDArray[Any], visualise.DataType], None] = None,
+    ) -> npt.NDArray[np.complex64]:
         # CSI data is not available for pilot subcarriers.
-        h_hat = np.where(
+        h_hat: npt.NDArray[np.complex64] = np.where(
             np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)),
             correlate(h, [[[1 / 2]], [[0]], [[1 / 2]]], mode="same"),
             h,
@@ -38,12 +45,33 @@ class Preprocessor:
         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.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 = 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)
+        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)
+                )
+            )
+
+        return 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)