feat: add cli util for finding correct antenna order

Pipfile

@@ -11,6 +11,7 @@ pytest = "*"
 matplotlib = "*"
 scipy = "*"
 scipy-stubs = "*"
+typer = "*"
 
 [dev-packages]
 

Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "c6ec3616be9da08134fd0a1eb5319ad8a9020abd1a929dd1196bbf5bf077f2e5"
+            "sha256": "dfb89f15ac6b380498f7a51a26ef1ba4b6e3a0cc01c8770dec4f2ef4b852371f"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -292,6 +292,14 @@
             "markers": "python_version >= '3.10'",
             "version": "==1.4.8"
         },
+        "markdown-it-py": {
+            "hashes": [
+                "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1",
+                "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb"
+            ],
+            "markers": "python_version >= '3.8'",
+            "version": "==3.0.0"
+        },
         "markupsafe": {
             "hashes": [
                 "sha256:0bff5e0ae4ef2e1ae4fdf2dfd5b76c75e5c2fa4132d05fc1b0dabcd20c7e28c4",
@@ -400,6 +408,14 @@
             "markers": "python_version >= '3.10'",
             "version": "==3.10.0"
         },
+        "mdurl": {
+            "hashes": [
+                "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8",
+                "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"
+            ],
+            "markers": "python_version >= '3.7'",
+            "version": "==0.1.2"
+        },
         "numpy": {
             "hashes": [
                 "sha256:059e6a747ae84fce488c3ee397cee7e5f905fd1bda5fb18c66bc41807ff119b2",
@@ -567,6 +583,14 @@
             "markers": "python_version >= '3.8'",
             "version": "==1.5.0"
         },
+        "pygments": {
+            "hashes": [
+                "sha256:61c16d2a8576dc0649d9f39e089b5f02bcd27fba10d8fb4dcc28173f7a45151f",
+                "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c"
+            ],
+            "markers": "python_version >= '3.8'",
+            "version": "==2.19.1"
+        },
         "pyparsing": {
             "hashes": [
                 "sha256:93d9577b88da0bbea8cc8334ee8b918ed014968fd2ec383e868fb8afb1ccef84",
@@ -592,6 +616,14 @@
             "markers": "python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2'",
             "version": "==2.9.0.post0"
         },
+        "rich": {
+            "hashes": [
+                "sha256:439594978a49a09530cff7ebc4b5c7103ef57baf48d5ea3184f21d9a2befa098",
+                "sha256:6049d5e6ec054bf2779ab3358186963bac2ea89175919d699e378b99738c2a90"
+            ],
+            "markers": "python_full_version >= '3.8.0'",
+            "version": "==13.9.4"
+        },
         "scipy": {
             "hashes": [
                 "sha256:0c2f95de3b04e26f5f3ad5bb05e74ba7f68b837133a4492414b3afd79dfe540e",
@@ -641,6 +673,14 @@
             "markers": "python_version >= '3.10'",
             "version": "==1.14.1.6"
         },
+        "shellingham": {
+            "hashes": [
+                "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686",
+                "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de"
+            ],
+            "markers": "python_version >= '3.7'",
+            "version": "==1.5.4"
+        },
         "simple-websocket": {
             "hashes": [
                 "sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c",
@@ -657,6 +697,23 @@
             "markers": "python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2'",
             "version": "==1.17.0"
         },
+        "typer": {
+            "hashes": [
+                "sha256:7994fb7b8155b64d3402518560648446072864beefd44aa2dc36972a5972e847",
+                "sha256:a0588c0a7fa68a1978a069818657778f86abe6ff5ea6abf472f940a08bfe4f0a"
+            ],
+            "index": "pypi",
+            "markers": "python_version >= '3.7'",
+            "version": "==0.15.1"
+        },
+        "typing-extensions": {
+            "hashes": [
+                "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d",
+                "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"
+            ],
+            "markers": "python_version >= '3.8'",
+            "version": "==4.12.2"
+        },
         "werkzeug": {
             "hashes": [
                 "sha256:54b78bf3716d19a65be4fceccc0d1d7b89e608834989dfae50ea87564639213e",

pyproject.toml

@@ -1,6 +1,12 @@
 [tool.ruff]
 line-length = 88
 
+[tool.ruff.lint]
+extend-select = ["I", "E", "W", "F", "B", "Q", "ANN"]
+
+[tool.ruff.lint.pycodestyle]
+max-doc-length = 88
+
 [tool.pyright]
 typeCheckingMode = "strict"
 reportMissingTypeStubs = "warning"

src/__main__.py

@@ -1,11 +1,6 @@
 import logging
-import multiprocessing as mp
-from typing import NamedTuple
 
-from . import ingest
-from . import config
-from . import visualise
-from . import aoa
+from . import cli
 
 
 logging.basicConfig(
@@ -13,35 +8,5 @@ logging.basicConfig(
     format="%(asctime)s %(name)-40s %(levelname)-8s %(message)s",
 )
 
-
-class Receiver(NamedTuple):
-    ip: ingest.Host
-    receiver: ingest.FeitReceiver
-    queue: "mp.Queue[ingest.CSI]"
-
-
-receivers = [
-    Receiver(ip, ingest.FeitReceiver(ip), mp.Queue(config.SAMPLE_RATE))
-    for ip in config.RECEIVE_HOSTS
-]
-
-# Start injecting CSI frames
-transmitter = ingest.FeitTransmitter()
-
-webapp_queue: "mp.Queue[aoa.AoA]" = mp.Queue(config.SAMPLE_RATE)
-
-processor = ingest.CSIProcessor({r.ip: r.queue for r in receivers}, webapp_queue)
-
-# Start webapp in background process
-webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
-webapp.start()
-
-receiver_processes = [
-    mp.Process(target=r.receiver.listen, args=(r.queue,)) for r in receivers
-]
-
-for proc in receiver_processes:
-    proc.start()
-
-processing_thread = mp.Process(target=processor.process_forever)
-processing_thread.start()
+if __name__ == "__main__":
+    cli.app()

src/cli.py

@@ -0,0 +1,48 @@
+import logging
+import multiprocessing as mp
+
+import numpy as np
+import numpy.typing as npt
+import typer
+
+from . import config, visualise
+from .collection import ingest
+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()
+def heatmap() -> None:
+    preprocessor = Preprocessor()
+    aoa = AoA()
+    webapp_queue: "mp.Queue[AoA]" = mp.Queue(config.SAMPLE_RATE)
+    # Start webapp in background process
+    webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
+    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)
+        if not webapp_queue.full():
+            webapp_queue.put(aoa)
+
+    ingest.start_processing(callback)

src/collection/csi_frame.py

@@ -1,8 +1,8 @@
 import struct
+
 import numpy as np
 import numpy.typing as npt
 
-
 RATE_MCS_MOD_TYPE_POS = 8
 RATE_MCS_MOD_TYPE_MSK = 0x7 << RATE_MCS_MOD_TYPE_POS
 RATE_MCS_CCK_MSK = 0 << RATE_MCS_MOD_TYPE_POS
@@ -37,7 +37,7 @@ RATE_MCS_BEAMF_MSK = 1 << RATE_MCS_BEAMF_POS
 
 
 class CSIHeader:
-    def __init__(self, data: bytes):
+    def __init__(self, data: bytes) -> None:
         self.csi_size = struct.unpack("I", data[0:4])[0]
         self.ftm_clock = struct.unpack("I", data[8:12])[0]
         self.num_rx = data[46]
@@ -93,7 +93,7 @@ class CSIHeader:
 
 class CSI:
     @staticmethod
-    def parseCsiData(data: bytes, header: CSIHeader):
+    def parseCsiData(data: bytes, header: CSIHeader) -> npt.NDArray[np.complex128]:
         csi_matrix: npt.NDArray[np.complex128] = np.zeros(
             (
                 header.num_subcarriers,
@@ -113,7 +113,7 @@ class CSI:
 
         return csi_matrix
 
-    def __init__(self, data: bytes):
+    def __init__(self, data: bytes) -> None:
         self.header = CSIHeader(data[:272])
         self.matrix = self.parseCsiData(
             data[272 : 272 + self.header.csi_size], self.header

src/collection/ingest.py

@@ -0,0 +1,213 @@
+import logging
+import multiprocessing as mp
+import socket
+import struct
+import subprocess
+import threading
+import time
+from datetime import datetime
+from typing import Callable, NamedTuple, NoReturn
+
+import numpy as np
+import numpy.typing as npt
+
+from .. import config
+from .csi_frame import CSI
+
+Host = tuple[str, int]
+CSICallback = Callable[[npt.NDArray[np.complex128]], None]
+
+
+class FeitHost:
+    def __init__(self, host: Host, command: str) -> None:
+        self.command = command
+        self.host = host
+        self.logger = logging.getLogger(
+            f"{__name__}.{self.__class__.__name__}-{self.host[0]}"
+        )
+        self.checker = threading.Thread(target=self.check_continuous)
+        self.checker.start()
+
+    def check_connection(self) -> bool:
+        feitcsi_status = subprocess.run(
+            f"ssh root@{self.host[0]} pgrep feitcsi",
+            check=False,
+            stdout=subprocess.DEVNULL,
+            shell=True,
+        )
+        return feitcsi_status.returncode == 0
+
+    def connect(self) -> None:
+        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.server.connect(self.host)
+        self.server.send(b"stop\n")
+        self.server.send(self.command.encode())
+        self.logger.info(f"Connected to {self.host}")
+
+    def check_continuous(self) -> NoReturn:
+        """
+        Repeatedly check if the FeitCSI service is running
+
+        Because FeitCSI is using TCP, we will get no information if the service stops,
+        or the computer is not reachable. In order to make debugging easier, this checks
+        and logs continuously if the service is running.
+        """
+        last_status = False
+        while True:
+            if not self.check_connection():
+                self.logger.error(f"FeitCSI is not running on {self.host[0]}")
+                last_status = False
+            else:
+                if not last_status:
+                    self.connect()
+                last_status = True
+            time.sleep(1)
+
+
+class FeitTransmitter(FeitHost):
+    def __init__(self) -> None:
+        command = (
+            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
+            f"--channel-width {config.CHANNEL_WIDTH} "
+            f"--format {config.FRAME_FORMAT} "
+            f"--mode inject -s 1 --verbose "
+            f"--inject-delay {1_000_000 // config.SAMPLE_RATE}"
+        )
+        super().__init__(config.INJECT_HOST, command)
+
+
+class FeitReceiver(FeitHost):
+    def __init__(self, host: Host) -> None:
+        command = (
+            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
+            f"--channel-width {config.CHANNEL_WIDTH} "
+            f"--format {config.FRAME_FORMAT} "
+            f"--mode measure"
+        )
+        super().__init__(host, command)
+
+    def listen(self, queue: "mp.Queue[CSI]") -> NoReturn:
+        prev_time = datetime.now()
+        self.logger.info("Listening for CSI data")
+        while True:
+            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()
+                queue.put(csidata)
+            except struct.error:
+                self.logger.error("Failed to parse CSI data")
+
+
+class CSIProcessor:
+    def __init__(
+        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
+
+    def add_data(self, host: Host, data: CSI) -> None:
+        if (
+            host in self.pending_data
+            and self.last_processed < self.pending_data[host][0]
+        ):
+            self.logger.warning(
+                f"Skipping data from {host} at {self.pending_data[host][0]}"
+            )
+
+        with self.pending_data_lock:
+            self.pending_data[host] = (datetime.now(), data)
+
+    def is_ready(self) -> bool:
+        for host in config.RECEIVE_HOSTS:
+            if (
+                host not in self.pending_data
+                or self.pending_data[host][0] <= self.last_processed
+            ):
+                return False
+        return True
+
+    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
+        ]
+
+        # We have data from all servers
+        all_data = np.concat(antenna_data, axis=1)
+
+        if callback:
+            callback(all_data)
+
+    def process_forever(
+        self,
+        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, pre_merge_callback=pre_merge_callback)
+            else:
+                self.logger.debug("Not all data is ready")
+            time.sleep(0.0005)
+
+
+class Receiver(NamedTuple):
+    ip: Host
+    receiver: FeitReceiver
+    queue: "mp.Queue[CSI]"
+
+
+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
+    ]
+
+    FeitTransmitter()
+
+    # Start injecting CSI frames
+    processor = CSIProcessor({r.ip: r.queue for r in receivers})
+
+    receiver_processes = [
+        threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
+    ]
+
+    for proc in receiver_processes:
+        proc.start()
+
+    processing_thread = mp.Process(
+        target=processor.process_forever, args=(csi_callback, pre_merge_callback)
+    )
+    processing_thread.start()
+    try:
+        processing_thread.join()
+    except KeyboardInterrupt:
+        return

src/ingest.py

@@ -1,140 +0,0 @@
-import logging
-import time
-import socket
-import multiprocessing as mp
-import struct
-from datetime import datetime
-import numpy as np
-
-from .csi import CSI
-from .aoa import AoA
-from .preprocess import Preprocessor
-from . import config
-
-
-Host = tuple[str, int]
-
-
-class FeitTransmitter:
-    def __init__(self):
-        self.inject_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        self.inject_server.connect(config.INJECT_HOST)
-        inject_start_string = (
-            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
-            f"--channel-width {config.CHANNEL_WIDTH} "
-            f"--format {config.FRAME_FORMAT} "
-            f"--mode inject -s 1 --verbose "
-            f"--inject-delay {1_000_000 // config.SAMPLE_RATE}"
-        )
-        self.inject_server.send(b"stop\n")
-        self.inject_server.send(inject_start_string.encode())
-
-
-class FeitReceiver:
-    def __init__(self, host: Host):
-        self.host = host
-        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        self.server.connect(host)
-        self.start_string = (
-            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
-            f"--channel-width {config.CHANNEL_WIDTH} "
-            f"--format {config.FRAME_FORMAT} "
-            f"--mode measure"
-        )
-        self.server.send(b"stop\n")
-        self.server.send(self.start_string.encode())
-        self.logger = logging.getLogger(
-            f"{__name__}.{self.__class__.__name__}-{self.host}"
-        )
-
-    def listen(self, queue: "mp.Queue[CSI]"):
-        prev_time = datetime.now()
-        while True:
-            # 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()
-                queue.put(csidata)
-            except struct.error:
-                self.logger.error("Failed to parse CSI data")
-
-
-class CSIProcessor:
-    def __init__(
-        self,
-        receiver_connections: dict[Host, "mp.Queue[CSI]"],
-        webserver: "mp.Queue[AoA]",
-    ):
-        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.preprocess = Preprocessor()
-        self.aoa = AoA()
-
-        self.connections = receiver_connections
-        self.webserver = webserver
-
-    def add_data(self, host: Host, data: CSI):
-        if (
-            host in self.pending_data
-            and self.last_processed < self.pending_data[host][0]
-        ):
-            self.logger.warning(
-                f"Skipping data from {host} at {self.pending_data[host][0]}"
-            )
-
-        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):
-        for host in config.RECEIVE_HOSTS:
-            if (
-                host not in self.pending_data
-                or self.pending_data[host][0] <= self.last_processed
-            ):
-                return False
-        return True
-
-    def process_data(self):
-        self.last_processed = datetime.now()
-        antenna_data = [
-            np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
-            for ip, antenna in config.ANTENNA_ORDER
-        ]
-
-        # 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}")
-
-        processed = self.preprocess.preprocess(all_data)
-        visualise.add_data(all_data, processed)
-        self.aoa.update(processed)
-        if not self.webserver.full():
-            self.webserver.put(self.aoa)
-
-    def process_forever(self):
-        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()
-            else:
-                self.logger.debug("Not all data is ready")
-            time.sleep(0.001)

src/processing/aoa.py

@@ -1,22 +1,23 @@
-import numpy as np
-import numpy.typing as npt
-from . import config
+import logging
 from datetime import datetime
 
-import logging
+import numpy as np
+import numpy.typing as npt
+
+from .. import config
 
 logger = logging.getLogger(__name__)
 
 
 class AoA:
-    def __init__(self):
+    def __init__(self) -> None:
         self.historical_autocorr = np.array([])
         self.N_subcarriers = -1
         self.N_rx = -1
         self.timestamp = datetime.now()
         pass
 
-    def smooth(self, data: npt.NDArray[np.complex128]):
+    def smooth(self, data: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]:
         assert len(data.shape) == 3
 
         M = data.shape[0]  # Number of subcarriers
@@ -44,7 +45,7 @@ class AoA:
 
         return H_sm
 
-    def update(self, data: npt.NDArray[np.complex128]):
+    def update(self, data: npt.NDArray[np.complex128]) -> None:
         self.timestamp = datetime.now()
         H_sm = self.smooth(data)
 
@@ -76,9 +77,11 @@ class AoA:
         2j * np.pi * config.CENTRAL_FREQUENCY_HZ * config.ANTENNA_SPACING / config.C
     )
 
-    def steering_vector(self, theta: float, tof: float):
-        omega_t = np.exp(-2j * np.pi * config.DELTA_F * tof)
-        phi_theta = np.exp(
+    def steering_vector(
+        self, theta: float, tof: float
+    ) -> npt.NDArray[np.complexfloating]:
+        omega_t: npt.NDArray[np.complex128] = np.exp(-2j * np.pi * config.DELTA_F * tof)
+        phi_theta: npt.NDArray[np.complex128] = np.exp(
             2j
             * np.pi
             * config.CENTRAL_FREQUENCY_HZ
@@ -96,7 +99,7 @@ class AoA:
         steering = antenna_v * phis
         return steering.T.reshape(-1)
 
-    def evaluate(self, theta: float, tof: float):
+    def evaluate(self, theta: float, tof: float) -> float:
         try:
             steering = self.steering_vector(theta, tof)
             steering_h = np.conj(steering).T
@@ -109,7 +112,7 @@ class AoA:
         return np.abs(c.real)
 
 
-def test_smoothing():
+def test_smoothing() -> None:
     row, col = np.indices((4, 2))
     data = row + 1j * col
     aoa = AoA()
@@ -120,10 +123,9 @@ def test_smoothing():
     expected = np.hstack([H_01, H_12])
     print(expected)
     assert np.allclose(smoothed, expected)
-    pass
 
 
-def test_steering_vector():
+def test_steering_vector() -> None:
     aoa = AoA()
     aoa.N_subcarriers = 10
     aoa.N_rx = 2

src/processing/preprocess.py

@@ -1,11 +1,11 @@
-import numpy as np
 import logging
 from queue import Queue
-from . import config
-import numpy.typing as npt
-from scipy.signal import correlate
-from scipy.signal import butter, sosfilt_zi, sosfilt
 
+import numpy as np
+import numpy.typing as npt
+from scipy.signal import butter, correlate, sosfilt, sosfilt_zi
+
+from .. import config
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -14,7 +14,7 @@ np.seterr(invalid="ignore")
 
 
 class Preprocessor:
-    def __init__(self):
+    def __init__(self) -> None:
         self.prev_entries: Queue[npt.NDArray[np.complex128]] = Queue(maxsize=100)
         self.short_term_avg = np.zeros((1,), dtype=np.complex128)
         self.long_term_avg = np.zeros((1,), dtype=np.complex128)

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)

src/visualise/__init__.py

@@ -11,7 +11,7 @@ import matplotlib.pyplot as plt
 import io
 import logging
 
-from ..aoa import AoA
+from ..processing.aoa import AoA
 from .. import config
 
 import matplotlib