feat: add cli util for finding correct antenna order

src/cli.py

@@ -1,3 +1,4 @@
+import logging
 import multiprocessing as mp
 
 import numpy as np
@@ -10,6 +11,21 @@ from .processing.aoa import AoA
 from .processing.preprocess import Preprocessor
 
 app = typer.Typer()
+logger = logging.getLogger(__name__)
+
+
+@app.command()
+def antennas() -> None:
+    """Utility to help determine the order in which antennas are plugged in
+
+    Once the script is running, unplug and replug antennas from left to right, to get
+    the correct order. Every time an antenna is unplugged and replugged, the script will
+    print the antenna identifier. When you are done, press Ctrl+C to stop the script and
+    get the final order.
+    """
+    from .utils import antenna_order
+
+    antenna_order.main()
 
 
 @app.command()
@@ -22,6 +38,7 @@ def heatmap() -> None:
     webapp.start()
 
     def callback(antenna_data: npt.NDArray[np.complex128]) -> None:
+        logger.info(f"Got final CSI data with shape {antenna_data.shape}")
         processed = preprocessor.preprocess(antenna_data)
         # visualise.add_data(all_data, processed)
         aoa.update(processed)
@@ -29,7 +46,3 @@ def heatmap() -> None:
             webapp_queue.put(aoa)
 
     ingest.start_processing(callback)
-
-
-if __name__ == "__main__":
-    app()

src/collection/ingest.py

@@ -130,14 +130,6 @@ class CSIProcessor:
         with self.pending_data_lock:
             self.pending_data[host] = (datetime.now(), data)
 
-            # Useful for figuring out the correct antenna order - RSSI values will
-            # decrease when the specific antenna is disconnected
-            rssis = [
-                (ip, csi.header.rssi1, csi.header.rssi2)
-                for ip, (_, csi) in sorted(self.pending_data.items())
-            ]
-            self.logger.debug("Antenna RSSI values: {}".format(rssis))
-
     def is_ready(self) -> bool:
         for host in config.RECEIVE_HOSTS:
             if (
@@ -147,8 +139,16 @@ class CSIProcessor:
                 return False
         return True
 
-    def process_data(self, callback: CSICallback) -> None:
+    def process_data(
+        self,
+        callback: CSICallback | None = None,
+        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
+    ) -> None:
         self.last_processed = datetime.now()
+        if pre_merge_callback is not None:
+            pre_merge_callback(
+                {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.ANTENNA_ORDER
@@ -156,20 +156,21 @@ class CSIProcessor:
 
         # We have data from all servers
         all_data = np.concat(antenna_data, axis=1)
-        self.logger.info(f"Got final CSI data with shape {all_data.shape}")
 
-        callback(all_data)
+        if callback:
+            callback(all_data)
 
     def process_forever(
         self,
-        callback: CSICallback,
+        callback: CSICallback | None = None,
+        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
     ) -> NoReturn:
         while True:
             for ip, queue in self.connections.items():
                 while not queue.empty():
                     self.add_data(ip, queue.get())
             if self.is_ready():
-                self.process_data(callback)
+                self.process_data(callback, pre_merge_callback=pre_merge_callback)
             else:
                 self.logger.debug("Not all data is ready")
             time.sleep(0.0005)
@@ -181,7 +182,10 @@ class Receiver(NamedTuple):
     queue: "mp.Queue[CSI]"
 
 
-def start_processing(csi_callback: CSICallback) -> None:
+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.SAMPLE_RATE))
         for ip in config.RECEIVE_HOSTS
@@ -200,7 +204,7 @@ def start_processing(csi_callback: CSICallback) -> None:
         proc.start()
 
     processing_thread = mp.Process(
-        target=processor.process_forever, args=(csi_callback,)
+        target=processor.process_forever, args=(csi_callback, pre_merge_callback)
     )
     processing_thread.start()
     try:

src/utils/antenna_order.py

@@ -0,0 +1,56 @@
+import logging
+
+from ..collection import ingest
+
+logger = logging.getLogger(__name__)
+
+AntennaIdentifier = tuple[ingest.Host, int]
+order: list[AntennaIdentifier] = []
+prev_unplugged: set[AntennaIdentifier] = set()
+antenna_average: dict[AntennaIdentifier, float] = {}
+
+RSSI_THRESHOLD = 8
+
+
+def callback(antenna_data: dict[ingest.Host, ingest.CSI]) -> None:
+    global prev_unplugged
+    global antenna_average
+    global order
+    unplugged: set[tuple[ingest.Host, int]] = set()
+    logger.debug(
+        "Antenna RSSI: "
+        + str(
+            {
+                host: (csi.header.rssi1, csi.header.rssi2)
+                for host, csi in antenna_data.items()
+            }
+        )
+    )
+    for host, csi in antenna_data.items():
+        antenna_average[(host, 0)] = (
+            antenna_average.get((host, 0), csi.header.rssi1) * 0.9
+            + csi.header.rssi1 * 0.1
+        )
+        antenna_average[(host, 1)] = (
+            antenna_average.get((host, 1), csi.header.rssi2) * 0.9
+            + csi.header.rssi2 * 0.1
+        )
+        if csi.header.rssi1 > antenna_average[(host, 0)] + RSSI_THRESHOLD:
+            unplugged.add((host, 0))
+        if csi.header.rssi2 > antenna_average[(host, 1)] + RSSI_THRESHOLD:
+            unplugged.add((host, 1))
+    if prev_unplugged != unplugged:
+        if len(prev_unplugged) > len(unplugged):
+            logger.info(f"Antenna plugged in: {prev_unplugged - unplugged}")
+        else:
+            logger.info(f"Antenna unplugged: {unplugged - prev_unplugged}")
+            diff = list(unplugged - prev_unplugged)
+            if len(diff) == 1 and diff[0] not in order:
+                host, antenna = diff.pop()
+                order.append((host, antenna))
+                logger.info(f"Current order: {order}")
+    prev_unplugged = unplugged
+
+
+def main() -> None:
+    ingest.start_processing(pre_merge_callback=callback)