refactor csi producers

This commit is contained in:
Christos Falas 2025-05-01 14:22:00 +01:00
parent cc27972d81
commit 581aa64ee3
No known key found for this signature in database
7 changed files with 214 additions and 195 deletions

View File

@ -6,6 +6,8 @@ from typing import Any, Callable, NamedTuple
import numpy.typing as npt
from where_fi.collection.protocols import CSIProducer
from .collection import CSIMatrix, ingest
from .config import config
from .visualise import server as visualise
@ -66,32 +68,10 @@ class CSIApplication:
mostly used as a scheduler).
"""
def create_receivers(self) -> list[Receiver]:
"""
Create a thread for each receiver, so that data can be received in parallel.
The receivers are created with a queue that is used to store the received data,
to be processed later. The receivers are created based on the configuration
file, which contains the IP addresses of the receivers.
"""
feit_receivers = [
ingest.FeitReceiver(ip, mp.Queue(config.collection_sample_rate))
for ip in config.receive_hosts
]
return [
Receiver(
ip,
receiver,
threading.Thread(target=receiver.listen),
)
for receiver, ip in zip(feit_receivers, config.receive_hosts, strict=True)
]
def __init__(self) -> None:
def __init__(self, producer: CSIProducer) -> None:
self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
self.receivers = self.create_receivers()
self.transmitter = ingest.FeitTransmitter()
self.producer = producer
self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
config.processing_sample_rate
@ -105,10 +85,6 @@ class CSIApplication:
self.pre_merge_callback = None
self.processing_callback = None
self.buffer = ingest.CSIAntennaJoin(
{r.ip: r.receiver.queue for r in self.receivers}
)
def visualise_data(
self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
) -> None:
@ -148,12 +124,30 @@ class CSIApplication:
self.processing_callback = func
return func
def listen(self) -> None:
"""
Listen for incoming data from the selected ingestor, and call the appropriate
callback.
"""
data_gen = self.producer()
while True:
try:
sample = next(data_gen)
if self.csi_callback is not None:
self.csi_callback(sample.matrix)
if self.pre_merge_callback is not None:
self.pre_merge_callback(sample.frames)
except StopIteration:
break
except Exception as e:
self.logger.error(f"Error processing sample: {e}", exc_info=True)
break
def start(self) -> None:
for receiver in self.receivers:
receiver.thread.start()
# for receiver in self.receivers:
# receiver.thread.start()
self.buffer_thread = threading.Thread(
target=self.buffer.process_forever,
args=(self.csi_callback, self.pre_merge_callback),
target=self.listen,
)
self.buffer_thread.start()
self.webapp_thread.start()
@ -170,16 +164,12 @@ class CSIApplication:
def stop(self) -> None:
self.logger.info("Stopping application")
for receiver in self.receivers:
receiver.receiver.active = False
receiver.thread.join()
self.transmitter.active = False
self.producer.stop()
if hasattr(self, "scheduler"):
self.scheduler.stop()
self.scheduler.thread.join()
self.buffer.active = False
self.buffer_thread.join()
self.webapp.active = False

View File

@ -1,8 +1,8 @@
import logging
import multiprocessing as mp
from queue import Queue
import threading
import time
from queue import Queue
from typing import Any, cast
import numpy as np
@ -33,7 +33,7 @@ def antennas() -> None:
"""
from ..utils import antenna_order
if not globals.is_live:
if not globals.csi_producer.is_live:
raise NotImplementedError("This command only works with live data")
antenna_order.main()
@ -41,37 +41,28 @@ def antennas() -> None:
@cli.command()
def heatmap() -> None:
app = CSIApplication(globals.csi_producer)
preprocessor = Preprocessor()
aoa = AoA()
# Start webapp in background process
webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
config.processing_sample_rate
)
webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
webapp.start()
def visualise_data(data: npt.NDArray[Any], dtype: visualise.figures.Figure) -> None:
if not webapp_queue.full():
webapp_queue.put(visualise.VisualiserData(data, dtype))
def callback(antenna_data: npt.NDArray[np.complex64]) -> None:
@app.on_sample
def _(antenna_data: npt.NDArray[np.complex64]) -> None:
logger.info(f"Got final CSI data with shape {antenna_data.shape}")
visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI)
processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
app.visualise_data(antenna_data, visualise.figures.Figure.RAW_CSI)
processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
logger.info(f"Processed CSI data with shape {processed.shape}")
processed_tensor = torch.tensor(processed, device=device)
aoa.update(processed_tensor)
aoa.heatmap(visualiser=visualise_data)
aoa.heatmap(visualiser=app.visualise_data)
globals.csi_producer(csi_callback=callback)
app.start()
logger.info("Finished processing CSI data")
@cli.command()
def phase_analysis(
subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0
subcarriers: list[int] = [0], rx_antenna: int = 0, tx_antenna: int = 0
) -> None:
"""
Visualise the phase information in the CSI data received from the antennas.
@ -81,26 +72,34 @@ def phase_analysis(
selected subcarrier and antenna.
"""
app = CSIApplication()
app = CSIApplication(globals.csi_producer)
preprocessor = Preprocessor()
subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate)
if isinstance(subcarriers, int):
subcarriers = [subcarriers]
print("Starting phase analysis on subcarriers: ", subcarriers)
subcarrier_phase: dict[int, Queue[float]] = {
x: Queue(config.collection_sample_rate) for x in subcarriers
}
@app.on_sample
def _(antenna_data: npt.NDArray[np.complex64]) -> None:
processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
for subcarrier in subcarriers:
phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna]))
if subcarrier_phase.full():
subcarrier_phase.get()
subcarrier_phase.put(phase)
if subcarrier_phase[subcarrier].full():
subcarrier_phase[subcarrier].get()
subcarrier_phase[subcarrier].put(phase)
@app.on_process
def _() -> None:
logger.info(f"Updating phase visualisation")
logger.info("Updating phase visualisation")
app.visualise_data(
np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS
np.array([x.queue for x in subcarrier_phase.values()]),
visualise.figures.Figure.PHASE_ANALYSIS,
)
# globals.csi_producer(csi_callback=callback)

View File

@ -1,18 +1,15 @@
from functools import partial
from pathlib import Path
from .. import collection
from ..collection import file, ingest
csi_producer: collection.CSIProducer = collection.noop
csi_producer: collection.CSIProducer = collection.NoopCSIProducer()
is_live = True
def main(from_file: Path | None = None) -> None:
global csi_producer, is_live
if from_file:
csi_producer = partial(file.start_processing, file_path=from_file)
is_live = False
csi_producer = file.FileCSIPRoducer(path=from_file)
else:
csi_producer = ingest.start_processing
is_live = True
csi_producer = ingest.RealtimeCSIProducer()

View File

@ -1,14 +1,24 @@
from typing import Iterator
import numpy as np
import numpy.typing as npt
from . import file, ingest
from .protocols import CSICallback, CSIProducer
from .protocols import CSIProducer, MergedCSI
def noop(csi_callback: CSICallback | None = None) -> None:
del csi_callback
class NoopCSIProducer:
"""A no-op CSI producer that does nothing"""
is_live = False
def __call__(self) -> Iterator[MergedCSI]:
return iter([])
def stop(self) -> None:
pass
CSIMatrix = npt.NDArray[np.complex64]
__all__ = ["file", "ingest", "noop", "CSIProducer", "CSICallback"]
__all__ = ["file", "ingest", "NoopCSIProducer", "CSIProducer"]

View File

@ -2,6 +2,7 @@ import logging
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Iterator
import h5py
@ -10,12 +11,20 @@ from . import protocols
logger = logging.getLogger(__name__)
def start_processing(
file_path: Path,
csi_callback: protocols.CSICallback | None = None,
) -> None:
logger.info(f"Replaying CSI data from {file_path}")
with h5py.File(file_path, "r") as file:
class FileCSIPRoducer:
is_live = False
def __init__(self, path: Path) -> None:
self.path = path
if not self.path.exists():
raise FileNotFoundError(f"File {self.path} does not exist")
def __call__(self) -> Iterator[protocols.MergedCSI]:
"""
Replay the CSI data from the file.
"""
logger.info(f"Replaying CSI data from {self.path}")
with h5py.File(self.path, "r") as file:
try:
for key in file:
datetime.fromisoformat(key)
@ -28,7 +37,13 @@ def start_processing(
target_offset = datetime.now() - start_time
for key in file:
logger.debug(f"Sending data from {key} at {datetime.now().isoformat()}")
csi_callback(file[key][:])
# TODO: add raw frames
yield protocols.MergedCSI(
frames={},
matrix=file[key][:],
)
curr_time_virtual = datetime.fromisoformat(key)
new_offset = datetime.now() - curr_time_virtual
logger.debug(f"New offset {new_offset}, target is {target_offset}")
@ -37,3 +52,9 @@ def start_processing(
f"Data is {new_offset - target_offset} behind, lagging behind..."
)
time.sleep(max(0, (target_offset - new_offset).total_seconds()))
def stop(self) -> None:
"""
Stop the producer.
"""
logger.info("Stopping file CSI producer")

View File

@ -1,18 +1,18 @@
import logging
import multiprocessing as mp
import selectors
import socket
import struct
import subprocess
import threading
import time
from datetime import datetime
from typing import Callable, NamedTuple
from typing import Iterator
import numpy as np
from ..config import config
from .csi_frame import CSI
from .protocols import CSICallback
from .protocols import MergedCSI
Host = tuple[str, int]
@ -27,6 +27,7 @@ class FeitHost:
self.active = True
self.checker = threading.Thread(target=self.check_continuous)
self.checker.start()
self.selectors: list[selectors.BaseSelector] = []
def check_connection(self) -> bool:
try:
@ -46,6 +47,10 @@ class FeitHost:
self.server.connect(self.host)
self.server.send(b"stop\n")
self.server.send(self.command.encode())
for selector in self.selectors:
selector.register(self.server, selectors.EVENT_READ, data=self)
self.logger.info(f"Connected to {self.host}")
def check_continuous(self) -> None:
@ -83,40 +88,49 @@ class FeitTransmitter(FeitHost):
class FeitReceiver(FeitHost):
def __init__(self, host: Host, queue: "mp.Queue[CSI]") -> None:
"""
A class used to connect to the host running FeitCSI and receive CSI data over a UDP
socket.
"""
def __init__(self, host: Host) -> None:
command = (
f"feitcsi --frequency {config.central_freq} "
f"--channel-width {config.channel_width} "
f"--format {config.frame_format} "
f"--mode measure"
)
self.queue = queue
super().__init__(host, command)
def listen(self) -> None:
prev_time = datetime.now()
self.logger.info("Listening for CSI data")
while self.active:
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()
self.queue.put(csidata)
except struct.error:
self.logger.error("Failed to parse CSI data")
self.logger.info("Stopping CSI receiver")
class CSIAntennaJoin:
def recv(self) -> CSI:
"""
Block until a CSI frame is received and return it.
"""
# Receive the data from the socket
data = self.server.recv(65535)
# Decode the data using the CSI frame format
csi = CSI(data)
return csi
def register(self, selector: selectors.BaseSelector) -> None:
"""
Register the receiver as a selector. This will allow us to multiplex the I/O
operations on a single thread.
This allows us to block until any of the receivers receive data
"""
self.selectors.append(selector)
if hasattr(self, "server"):
selector.register(self.server, selectors.EVENT_READ, data=self)
class CSIAntennaArray:
"""
Represents a set of receiver hosts that are used to receive CSI data, assumed to be
part of a linear antenna array.
This class is used to buffer a set of CSI data from multiple receivers, and once a
set of readings is ready for each receiver, it will be merged into a single sample
using the antenna order in the configuration.
@ -127,15 +141,12 @@ class CSIAntennaJoin:
- `sample_callback`: Called with the merged data
"""
def __init__(
self,
receiver_connections: dict[Host, "mp.Queue[CSI]"],
) -> None:
def __init__(self, receivers: list[FeitReceiver]) -> 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
self.receivers = receivers
self.active = True
def add_data(self, host: Host, data: CSI) -> None:
@ -159,16 +170,13 @@ class CSIAntennaJoin:
return False
return True
def process_data(
self,
sample_callback: CSICallback | None = None,
pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
) -> None:
def process_data(self) -> None | MergedCSI:
if not self.is_ready():
self.logger.debug("Not all data is ready")
return 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()}
)
frames = {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.antennas.order
@ -176,63 +184,45 @@ class CSIAntennaJoin:
# We have data from all servers
all_data = np.concat(antenna_data, axis=1)
return MergedCSI(frames=frames, matrix=all_data)
if sample_callback:
sample_callback(all_data)
def process_forever(self) -> Iterator[MergedCSI]:
# Register the receivers with the selector
self.logger.info("Starting to process data from receivers")
selector = selectors.DefaultSelector()
for receiver in self.receivers:
receiver.register(selector)
def process_forever(
self,
callback: CSICallback | None = None,
pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
) -> None:
while self.active:
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)
events = selector.select(timeout=0.1)
if not events:
self.logger.warning("No events received in the last 0.1 seconds")
continue
for key, _ in events:
receiver = key.data
self.add_data(receiver.host, receiver.recv())
sample = self.process_data()
if sample is not None:
yield sample
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.collection_sample_rate))
for ip in config.receive_hosts
]
transmitter = FeitTransmitter()
class RealtimeCSIProducer:
def __init__(self) -> None:
self.receivers = [FeitReceiver(ip) for ip in config.receive_hosts]
self.transmitter = FeitTransmitter()
# Start injecting CSI frames
buffer = CSIAntennaJoin({r.ip: r.queue for r in receivers})
self.buffer = CSIAntennaArray(self.receivers)
self.is_live = True
receiver_threads = [
threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
]
def __call__(self) -> Iterator[MergedCSI]:
return self.buffer.process_forever()
for proc in receiver_threads:
proc.start()
buffering_thread = threading.Thread(
target=buffer.process_forever, args=(csi_callback, pre_merge_callback)
)
buffering_thread.start()
try:
while True:
time.sleep(100)
except KeyboardInterrupt:
for r in receivers:
r.receiver.active = False
transmitter.active = False
buffer.active = False
return
def stop(self) -> None:
for r in self.receivers:
r.active = False
self.transmitter.active = False
self.buffer.active = False

View File

@ -1,10 +1,22 @@
from typing import Callable, Protocol
from typing import Iterator, NamedTuple, Protocol
import numpy as np
import numpy.typing as npt
CSICallback = Callable[[npt.NDArray[np.complex64]], None]
from .csi_frame import CSI
CSIHost = tuple[str, int]
class MergedCSI(NamedTuple):
"""Merged CSI data from all antennas"""
frames: dict[CSIHost, CSI]
matrix: npt.NDArray[np.complex64]
class CSIProducer(Protocol):
def __call__(self, csi_callback: CSICallback | None = None) -> None: ...
is_live: bool
def __call__(self) -> Iterator[MergedCSI]: ...
def stop(self) -> None: ...