refactor: switch to processes instead of threads

src/__main__.py

@@ -1,10 +1,11 @@
 import logging
-import threading
+import multiprocessing as mp
+from multiprocessing.connection import Connection
+from typing import NamedTuple
 
 from . import ingest
 from . import config
 from . import visualise
-from .preprocess import Preprocessor
 
 
 logging.basicConfig(
@@ -12,24 +13,36 @@ logging.basicConfig(
     format="%(asctime)s %(name)-40s %(levelname)-8s %(message)s",
 )
 
-receivers = [ingest.FeitReceiver(ip) for ip in config.RECEIVE_IP_ADDRESS_LIST]
-transmitter = ingest.FeitTransmitter()
-preprocess = Preprocessor()
 
-processor = ingest.CSIProcessor()
+class Receiver(NamedTuple):
+    ip: str
+    receiver: ingest.FeitReceiver
+    recv: Connection
+    send: Connection
 
-# Start webapp in background thread
-visualise.aoa = processor.aoa
-webapp = threading.Thread(target=visualise.start)
-webapp.start()
 
-receiver_threads = [
-    threading.Thread(target=server.listen, args=(processor.add_data,))
-    for server in receivers
+receivers = [
+    Receiver(ip, ingest.FeitReceiver(ip), *mp.Pipe(duplex=False))
+    for ip in config.RECEIVE_IP_ADDRESS_LIST
 ]
 
-for thread in receiver_threads:
-    thread.start()
+# Start injecting CSI frames
+transmitter = ingest.FeitTransmitter()
 
-processing_thread = threading.Thread(target=processor.process_forever)
+webapp_conns = mp.Pipe()
+
+processor = ingest.CSIProcessor({r.ip: r.recv for r in receivers}, webapp_conns[0])
+
+# Start webapp in background process
+webapp = mp.Process(target=visualise.start, args=(webapp_conns[1],))
+webapp.start()
+
+receiver_processes = [
+    mp.Process(target=r.receiver.listen, args=(r.send,)) for r in receivers
+]
+
+for proc in receiver_processes:
+    proc.start()
+
+processing_thread = mp.Process(target=processor.process_forever)
 processing_thread.start()

src/aoa.py

@@ -10,8 +10,8 @@ logger = logging.getLogger(__name__)
 class AoA:
     def __init__(self):
         self.historical_autocorr = np.array([])
-        self.N_subcarriers = -1
-        self.N_rx = -1
+        self.N_subcarriers = config.N_SUBCARRIERS - 2
+        self.N_rx = len(config.ANTENNA_ORDER)
         pass
 
     def smooth(self, data: npt.NDArray[np.complex128]):
@@ -21,8 +21,9 @@ class AoA:
         N = data.shape[1]  # Number of RX antennas
         T = data.shape[2]  # Number of TX antennas
 
-        self.N_subcarriers = M
-        self.N_rx = N
+        assert N == self.N_rx
+        assert M == self.N_subcarriers
+
         logger.debug(f"Smoothing: Subcarriers: {M}, RX antennas: {N}, TX antennas: {T}")
 
         # This only works with 1 TX antenna (i.e. no MIMO) - see #4 for more details

src/config.py

@@ -31,6 +31,8 @@ ANTENNA_SPACING = 0.0285  # 2.85 cm
 CHANNEL_WIDTH = 20
 FRAME_FORMAT = "HT"
 
+N_SUBCARRIERS = 56
+
 
 CENTRAL_FREQUENCY_HZ = CENTRAL_FREQUENCY_MHZ * 1_000_000
 C = 299_792_458  # m/s

src/ingest.py

@@ -1,8 +1,8 @@
 import logging
 import time
 import socket
-from typing import Callable
-import threading
+import multiprocessing as mp
+from multiprocessing.connection import Connection
 import struct
 from datetime import datetime
 import numpy as np
@@ -46,7 +46,7 @@ class FeitReceiver:
             f"{__name__}.{self.__class__.__name__}-{self.host}"
         )
 
-    def listen(self, callback: Callable[[str, CSI], None]):
+    def listen(self, conn: Connection):
         prev_time = datetime.now()
         while True:
             # This is the max size of a UDP packet. The size of the actual CSI
@@ -59,21 +59,26 @@ class FeitReceiver:
                     f"Received CSI data after {datetime.now() - prev_time}"
                 )
                 prev_time = datetime.now()
-                callback(self.host, csidata)
+                conn.send(csidata)
             except struct.error:
                 self.logger.error("Failed to parse CSI data")
 
 
 class CSIProcessor:
-    def __init__(self):
+    def __init__(
+        self, receiver_connections: dict[str, Connection], webserver: Connection
+    ):
         self.pending_data: dict[str, tuple[datetime, CSI]] = {}
-        self.pending_data_lock = threading.Lock()
+        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: str, data: CSI):
         if (
             host in self.pending_data
@@ -117,9 +122,13 @@ class CSIProcessor:
         processed = self.preprocess.preprocess(all_data)
         visualise.add_data(all_data, processed)
         self.aoa.update(processed)
+        self.webserver.send(self.aoa.E_n)
 
     def process_forever(self):
         while True:
+            for ip, conn in self.connections.items():
+                if conn.poll():
+                    self.add_data(ip, conn.recv())
             if self.is_ready():
                 self.process_data()
             else:

src/visualise/__init__.py

@@ -5,6 +5,7 @@ import numpy.typing as npt
 from simple_websocket import Server
 import time
 from datetime import datetime
+from multiprocessing.connection import Connection
 
 import matplotlib.pyplot as plt
 import io
@@ -19,6 +20,7 @@ matplotlib.use("agg")
 
 app = Flask(__name__)
 sock = Sock(app)
+aoa_conn = None
 
 logger = logging.getLogger(__name__)
 
@@ -110,11 +112,17 @@ def make_heatmap(max_tof: float):
 
 
 def gather_aoa(max_tof: float):
+    assert aoa_conn is not None
+
     prev_frame = datetime.now()
     while True:
         while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
             time.sleep(0.01)
-            pass
+        while aoa_conn.poll():
+            logger.debug("Receiving from aoa pipe")
+            E_n = aoa_conn.recv()
+        aoa.N_subcarriers = 54
+        aoa.E_n = E_n
         prev_frame = datetime.now()
         logger.debug("Generating heatmap")
         buf = make_heatmap(max_tof)
@@ -134,5 +142,8 @@ def aoa_tof():
     )
 
 
-def start():
+def start(conn: Connection):
+    global aoa_conn
+
+    aoa_conn = conn
     app.run(debug=True, use_reloader=False, host="0.0.0.0")