cleanup types

pyproject.toml

@@ -0,0 +1,6 @@
+[tool.ruff]
+line-length = 88
+
+[tool.pyright]
+typeCheckingMode = "strict"
+reportMissingTypeStubs = "warning"

src/aoa.py

@@ -0,0 +1,17 @@
+import numpy as np
+import numpy.typing as npt
+import config
+
+
+class AoA:
+    def __init__(self):
+        self.historical_data = np.array([])
+        pass
+
+    def update(self, data: npt.NDArray[np.complex128]):
+        self.historical_data = np.concatenate([self.historical_data, data], axis=0)
+        if self.historical_data.shape[0] > config.AOA_SLIDING_WINDOW_SIZE:
+            self.historical_data = self.historical_data[
+                -config.AOA_SLIDING_WINDOW_SIZE :
+            ]
+        pass

src/config.py

@@ -0,0 +1,8 @@
+import os
+
+PREPROCESSING_SHORT_TERM_WINDOW_SIZE = 5
+PREPROCESSING_LONG_TERM_ALPHA = 0.1
+AOA_SLIDING_WINDOW_SIZE = 20
+
+FEITCSI_IP_ADDRESS = os.getenv("IP_ADDRESS", "10.0.12.62")
+FEITCSI_PORT = 8008

src/csi.py

@@ -1,5 +1,6 @@
 import struct
 import numpy as np
+import numpy.typing as npt
 
 
 RATE_MCS_MOD_TYPE_POS = 8
@@ -93,13 +94,13 @@ class CSIHeader:
 class CSI:
     @staticmethod
     def parseCsiData(data: bytes, header: CSIHeader):
-        csi_matrix = np.zeros(
+        csi_matrix: npt.NDArray[np.complex128] = np.zeros(
             (
                 header.num_subcarriers,
                 header.num_rx,
                 header.num_tx,
             ),
-            dtype=complex,
+            dtype=np.complex128,
         )
         pos = 0
         for j in range(header.num_rx):

src/main.py

@@ -1,10 +1,13 @@
 import logging
 import socket
+from typing import Callable
 from csi import CSI
+from aoa import AoA
 from preprocess import Preprocessor
 import visualise
 import threading
 import struct
+import config
 
 logging.basicConfig(
     level=logging.INFO,
@@ -21,7 +24,7 @@ class FeitServer:
         frame_format: str = "VHT",
     ):
         self.ip = ip
-        self.port = 8008
+        self.port = config.FEITCSI_PORT
         self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         self.server.connect((self.ip, self.port))
         self.start_string = (
@@ -32,8 +35,11 @@ class FeitServer:
         self.server.send(b"stop\n")
         self.server.send(self.start_string.encode())
 
-    def listen(self, callback):
+    def listen(self, callback: Callable[[CSI], None]):
         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)
@@ -42,15 +48,17 @@ class FeitServer:
                 logging.error("Failed to parse CSI data")
 
 
-def process_data(data):
+def process_data(data: CSI):
     logging.debug("Got CSI frame")
     processed = preprocess.preprocess(data)
     visualise.add_data(processed)
+    aoa.update(processed)
 
 
 if __name__ == "__main__":
-    server = FeitServer("10.0.12.62")
+    server = FeitServer(config.FEITCSI_IP_ADDRESS)
     preprocess = Preprocessor()
+    aoa = AoA()
 
     # Start webapp in background thread
     webapp = threading.Thread(target=visualise.start)

src/preprocess.py

@@ -2,19 +2,18 @@ import numpy as np
 from csi import CSI
 import logging
 from queue import Queue
+import config
+import numpy.typing as npt
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
 class Preprocessor:
-    NUM_SHORT_TERM_ENTRIES = 10
-    LONG_TERM_ALPHA = 0.1
-
     def __init__(self):
-        self.prev_entries = Queue(maxsize=100)
+        self.prev_entries: Queue[npt.NDArray[np.complex128]] = Queue(maxsize=100)
-        self.short_term_avg = np.zeros((1,), dtype=complex)
+        self.short_term_avg = np.zeros((1,), dtype=np.complex128)
-        self.long_term_avg = np.zeros((1,), dtype=complex)
+        self.long_term_avg = np.zeros((1,), dtype=np.complex128)
 
     def preprocess(self, csi: CSI):
         h = csi.matrix
@@ -22,18 +21,20 @@ class Preprocessor:
 
         # Assume that all csi matrices will have the same shape
         if self.short_term_avg.shape != h_hat.shape:
-            self.short_term_avg = np.zeros(h_hat.shape, dtype=complex)
+            self.short_term_avg = np.zeros(h_hat.shape, dtype=np.complex128)
-            self.long_term_avg = np.zeros(h_hat.shape, dtype=complex)
+            self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex128)
             self.prev_entries = Queue(maxsize=100)
 
         while self.prev_entries.full():
             old_value = self.prev_entries.get()
-            self.short_term_avg -= old_value / self.NUM_SHORT_TERM_ENTRIES
+            self.short_term_avg -= (
+                old_value / config.PREPROCESSING_SHORT_TERM_WINDOW_SIZE
+            )
         self.prev_entries.put(h_hat)
-        self.short_term_avg += h_hat / self.NUM_SHORT_TERM_ENTRIES
+        self.short_term_avg += h_hat / config.PREPROCESSING_SHORT_TERM_WINDOW_SIZE
         self.long_term_avg = (
-            self.long_term_avg * (1 - self.LONG_TERM_ALPHA)
+            self.long_term_avg * (1 - config.PREPROCESSING_LONG_TERM_ALPHA)
-            + self.short_term_avg * self.LONG_TERM_ALPHA
+            + self.short_term_avg * config.PREPROCESSING_LONG_TERM_ALPHA
         )
 
         # Remove static components

src/visualise/__init__.py

@@ -1,8 +1,8 @@
-from flask import Flask, render_template, request
+from flask import Flask, render_template
 from flask_sock import Sock
 import numpy as np
 import numpy.typing as npt
-import json
+from simple_websocket import Server
 
 app = Flask(__name__)
 sock = Sock(app)
@@ -16,17 +16,21 @@ def preprocessed():
     return render_template("preprocessed.html")
 
 
-subscribers = {}
+type Subscriber = Server
+subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {}
 
 
 @sock.route("/data")
-def get_data(sock):
+def get_data(sock: Subscriber):
     while True:
         msg = sock.receive()
-        subscribers[sock] = list(map(int, msg.split()))
+        if len(msg.split()) != 3:
+            break
+        subcarrier, rx, tx = map(int, msg.split())
+        subscriber_settings[sock] = (subcarrier, rx, tx)
 
 
-def add_data(new_data):
+def add_data(new_data: npt.NDArray[np.complex128]):
     global data
     if data.size == 0:
         data = np.expand_dims(new_data, axis=0)
@@ -36,17 +40,17 @@ def add_data(new_data):
         # Only keep latest 100 entries
         if data.shape[0] > 100:
             data = data[-100:]
-    to_remove = []
+    to_remove: list[Subscriber] = []
-    for subscriber in subscribers:
+    for subscriber in subscriber_settings:
         try:
-            subcarrier, rx, tx = subscribers[subscriber]
+            subcarrier, rx, tx = subscriber_settings[subscriber]
             subscriber.send(new_data.real[subcarrier, rx, tx])
         except Exception as e:
             to_remove.append(subscriber)
             print(e)
 
     for subscriber in to_remove:
-        del subscribers[subscriber]
+        del subscriber_settings[subscriber]
 
 
 def start():

src/visualise/static/plot.js

@@ -12,7 +12,7 @@ Plotly.newPlot(
     {
       x: [],
       y: [],
-      type: "scatter", // Line chart
+      type: "scatter",
       mode: "lines+markers", // Line + markers
     },
   ],

typings/flask_sock/__init__.pyi

@@ -0,0 +1,56 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from functools import wraps
+from flask import Blueprint, Response, current_app, request
+from simple_websocket import ConnectionClosed, Server
+from typing import Callable
+
+class Sock:
+    """Instantiate the Flask-Sock extension.
+
+    :param app: The Flask application instance. If not provided, it must be
+                initialized later by calling the :func:`Sock.init_app` method.
+    """
+    def __init__(self, app=...) -> None: ...
+    def init_app(self, app):  # -> None:
+        """Initialize the Flask-Socket extension.
+
+
+        :param app: The Flask application instance. This method only needs to
+                    be called if the application instance was not passed as
+                    an argument to the constructor.
+        """
+        ...
+
+    def route(
+        self, path: str, bp: Blueprint | None = None
+    ) -> Callable[
+        [Callable[[Server], None]], Callable[[Server], None]
+    ]:  # -> Callable[..., None]:
+        """Decorator to create a WebSocket route.
+
+        The decorated function will be invoked when a WebSocket client
+        establishes a connection, with a WebSocket connection object passed
+        as an argument. Example::
+
+            @sock.route('/ws')
+            def websocket_route(ws):
+                # The ws object has the following methods:
+                # - ws.send(data)
+                # - ws.receive(timeout=None)
+                # - ws.close(reason=None, message=None)
+
+        If the route has variable components, the ``ws`` argument needs to be
+        included before them.
+
+        :param path: the URL associated with the route.
+        :param bp: the blueprint on which to register the route. If not given,
+                   the route is attached directly to the Flask application
+                   instance. When a blueprint is used, the application is
+                   responsible for the blueprint's registration.
+        :param kwargs: additional route options. See the Flask documentation
+                       for the ``app.route`` decorator for details.
+        """
+        ...

typings/simple_websocket/__init__.pyi

@@ -0,0 +1,16 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from .ws import Client, Server
+from .aiows import AioClient, AioServer
+from .errors import ConnectionClosed, ConnectionError
+
+__all__ = [
+    "Client",
+    "Server",
+    "AioClient",
+    "AioServer",
+    "ConnectionClosed",
+    "ConnectionError",
+]

typings/simple_websocket/aiows.pyi

@@ -0,0 +1,163 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+class AioBase:
+    def __init__(self, connection_type=..., receive_bytes=..., ping_interval=..., max_message_size=...) -> None:
+        ...
+    
+    async def connect(self): # -> None:
+        ...
+    
+    async def handshake(self): # -> None:
+        ...
+    
+    async def send(self, data): # -> None:
+        """Send data over the WebSocket connection.
+
+        :param data: The data to send. If ``data`` is of type ``bytes``, then
+                     a binary message is sent. Else, the message is sent in
+                     text format.
+        """
+        ...
+    
+    async def receive(self, timeout=...): # -> None:
+        """Receive data over the WebSocket connection.
+
+        :param timeout: Amount of time to wait for the data, in seconds. Set
+                        to ``None`` (the default) to wait indefinitely. Set
+                        to 0 to read without blocking.
+
+        The data received is returned, as ``bytes`` or ``str``, depending on
+        the type of the incoming message.
+        """
+        ...
+    
+    async def close(self, reason=..., message=...): # -> None:
+        """Close the WebSocket connection.
+
+        :param reason: A numeric status code indicating the reason of the
+                       closure, as defined by the WebSocket specification. The
+                       default is 1000 (normal closure).
+        :param message: A text message to be sent to the other side.
+        """
+        ...
+    
+    def choose_subprotocol(self, request): # -> None:
+        ...
+    
+
+
+class AioServer(AioBase):
+    """This class implements a WebSocket server.
+
+    Instead of creating an instance of this class directly, use the
+    ``accept()`` class method to create individual instances of the server,
+    each bound to a client request.
+    """
+    def __init__(self, request, subprotocols=..., receive_bytes=..., ping_interval=..., max_message_size=...) -> None:
+        ...
+    
+    @classmethod
+    async def accept(cls, aiohttp=..., asgi=..., sock=..., headers=..., subprotocols=..., receive_bytes=..., ping_interval=..., max_message_size=...): # -> WebSocketASGI | Self:
+        """Accept a WebSocket connection from a client.
+
+        :param aiohttp: The request object from aiohttp. If this argument is
+                        provided, ``asgi``, ``sock`` and ``headers`` must not
+                        be set.
+        :param asgi: A (scope, receive, send) tuple from an ASGI request. If
+                     this argument is provided, ``aiohttp``, ``sock`` and
+                     ``headers`` must not be set.
+        :param sock: A connected socket to use. If this argument is provided,
+                     ``aiohttp`` and ``asgi`` must not be set. The ``headers``
+                     argument must be set with the incoming request headers.
+        :param headers: A dictionary with the incoming request headers, when
+                        ``sock`` is used.
+        :param subprotocols: A list of supported subprotocols, or ``None`` (the
+                             default) to disable subprotocol negotiation.
+        :param receive_bytes: The size of the receive buffer, in bytes. The
+                              default is 4096.
+        :param ping_interval: Send ping packets to clients at the requested
+                              interval in seconds. Set to ``None`` (the
+                              default) to disable ping/pong logic. Enable to
+                              prevent disconnections when the line is idle for
+                              a certain amount of time, or to detect
+                              unresponsive clients and disconnect them. A
+                              recommended interval is 25 seconds.
+        :param max_message_size: The maximum size allowed for a message, in
+                                 bytes, or ``None`` for no limit. The default
+                                 is ``None``.
+        """
+        ...
+    
+    async def handshake(self): # -> None:
+        ...
+    
+    def choose_subprotocol(self, request): # -> None:
+        """Choose a subprotocol to use for the WebSocket connection.
+
+        The default implementation selects the first protocol requested by the
+        client that is accepted by the server. Subclasses can override this
+        method to implement a different subprotocol negotiation algorithm.
+
+        :param request: A ``Request`` object.
+
+        The method should return the subprotocol to use, or ``None`` if no
+        subprotocol is chosen.
+        """
+        ...
+    
+
+
+class AioClient(AioBase):
+    """This class implements a WebSocket client.
+
+    Instead of creating an instance of this class directly, use the
+    ``connect()`` class method to create an instance that is connected to a
+    server.
+    """
+    def __init__(self, url, subprotocols=..., headers=..., receive_bytes=..., ping_interval=..., max_message_size=..., ssl_context=...) -> None:
+        ...
+    
+    @classmethod
+    async def connect(cls, url, subprotocols=..., headers=..., receive_bytes=..., ping_interval=..., max_message_size=..., ssl_context=..., thread_class=..., event_class=...): # -> Self:
+        """Returns a WebSocket client connection.
+
+        :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are
+                    accepted.
+        :param subprotocols: The name of the subprotocol to use, or a list of
+                             subprotocol names in order of preference. Set to
+                             ``None`` (the default) to not use a subprotocol.
+        :param headers: A dictionary or list of tuples with additional HTTP
+                        headers to send with the connection request. Note that
+                        custom headers are not supported by the WebSocket
+                        protocol, so the use of this parameter is not
+                        recommended.
+        :param receive_bytes: The size of the receive buffer, in bytes. The
+                              default is 4096.
+        :param ping_interval: Send ping packets to the server at the requested
+                              interval in seconds. Set to ``None`` (the
+                              default) to disable ping/pong logic. Enable to
+                              prevent disconnections when the line is idle for
+                              a certain amount of time, or to detect an
+                              unresponsive server and disconnect. A recommended
+                              interval is 25 seconds. In general it is
+                              preferred to enable ping/pong on the server, and
+                              let the client respond with pong (which it does
+                              regardless of this setting).
+        :param max_message_size: The maximum size allowed for a message, in
+                                 bytes, or ``None`` for no limit. The default
+                                 is ``None``.
+        :param ssl_context: An ``SSLContext`` instance, if a default SSL
+                            context isn't sufficient.
+        """
+        ...
+    
+    async def handshake(self): # -> None:
+        ...
+    
+    async def close(self, reason=..., message=...): # -> None:
+        ...
+    
+
+

typings/simple_websocket/asgi.pyi

@@ -0,0 +1,23 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+class WebSocketASGI:
+    def __init__(self, scope, receive, send, subprotocols=...) -> None:
+        ...
+    
+    @classmethod
+    async def accept(cls, scope, receive, send, subprotocols=...): # -> WebSocketASGI:
+        ...
+    
+    async def receive(self):
+        ...
+    
+    async def send(self, data): # -> None:
+        ...
+    
+    async def close(self): # -> None:
+        ...
+    
+
+

typings/simple_websocket/errors.pyi

@@ -0,0 +1,22 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+class SimpleWebsocketError(RuntimeError):
+    ...
+
+
+class ConnectionError(SimpleWebsocketError):
+    """Connection error exception class."""
+    def __init__(self, status_code=...) -> None:
+        ...
+    
+
+
+class ConnectionClosed(SimpleWebsocketError):
+    """Connection closed exception class."""
+    def __init__(self, reason=..., message=...) -> None:
+        ...
+    
+
+

typings/simple_websocket/ws.pyi

@@ -0,0 +1,211 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+class Base:
+    def __init__(
+        self,
+        sock=...,
+        connection_type=...,
+        receive_bytes=...,
+        ping_interval=...,
+        max_message_size=...,
+        thread_class=...,
+        event_class=...,
+        selector_class=...,
+    ) -> None: ...
+    def handshake(self):  # -> None:
+        ...
+    def send(self, data: bytes | str) -> None:  # -> None:
+        """Send data over the WebSocket connection.
+
+        :param data: The data to send. If ``data`` is of type ``bytes``, then
+                     a binary message is sent. Else, the message is sent in
+                     text format.
+        """
+        ...
+
+    def receive(self, timeout: float | None = ...) -> bytes | str:  # -> None:
+        """Receive data over the WebSocket connection.
+
+        :param timeout: Amount of time to wait for the data, in seconds. Set
+                        to ``None`` (the default) to wait indefinitely. Set
+                        to 0 to read without blocking.
+
+        The data received is returned, as ``bytes`` or ``str``, depending on
+        the type of the incoming message.
+        """
+        ...
+
+    def close(self, reason=..., message=...):  # -> None:
+        """Close the WebSocket connection.
+
+        :param reason: A numeric status code indicating the reason of the
+                       closure, as defined by the WebSocket specification. The
+                       default is 1000 (normal closure).
+        :param message: A text message to be sent to the other side.
+        """
+        ...
+
+    def choose_subprotocol(self, request):  # -> None:
+        ...
+
+class Server(Base):
+    """This class implements a WebSocket server.
+
+    Instead of creating an instance of this class directly, use the
+    ``accept()`` class method to create individual instances of the server,
+    each bound to a client request.
+    """
+    def __init__(
+        self,
+        environ,
+        subprotocols=...,
+        receive_bytes=...,
+        ping_interval=...,
+        max_message_size=...,
+        thread_class=...,
+        event_class=...,
+        selector_class=...,
+    ) -> None: ...
+    @classmethod
+    def accept(
+        cls,
+        environ,
+        subprotocols=...,
+        receive_bytes=...,
+        ping_interval=...,
+        max_message_size=...,
+        thread_class=...,
+        event_class=...,
+        selector_class=...,
+    ):  # -> Self:
+        """Accept a WebSocket connection from a client.
+
+        :param environ: A WSGI ``environ`` dictionary with the request details.
+                        Among other things, this class expects to find the
+                        low-level network socket for the connection somewhere
+                        in this dictionary. Since the WSGI specification does
+                        not cover where or how to store this socket, each web
+                        server does this in its own different way. Werkzeug,
+                        Gunicorn, Eventlet and Gevent are the only web servers
+                        that are currently supported.
+        :param subprotocols: A list of supported subprotocols, or ``None`` (the
+                             default) to disable subprotocol negotiation.
+        :param receive_bytes: The size of the receive buffer, in bytes. The
+                              default is 4096.
+        :param ping_interval: Send ping packets to clients at the requested
+                              interval in seconds. Set to ``None`` (the
+                              default) to disable ping/pong logic. Enable to
+                              prevent disconnections when the line is idle for
+                              a certain amount of time, or to detect
+                              unresponsive clients and disconnect them. A
+                              recommended interval is 25 seconds.
+        :param max_message_size: The maximum size allowed for a message, in
+                                 bytes, or ``None`` for no limit. The default
+                                 is ``None``.
+        :param thread_class: The ``Thread`` class to use when creating
+                             background threads. The default is the
+                             ``threading.Thread`` class from the Python
+                             standard library.
+        :param event_class: The ``Event`` class to use when creating event
+                            objects. The default is the `threading.Event``
+                            class from the Python standard library.
+        :param selector_class: The ``Selector`` class to use when creating
+                               selectors. The default is the
+                               ``selectors.DefaultSelector`` class from the
+                               Python standard library.
+        """
+        ...
+
+    def handshake(self):  # -> None:
+        ...
+    def choose_subprotocol(self, request):  # -> None:
+        """Choose a subprotocol to use for the WebSocket connection.
+
+        The default implementation selects the first protocol requested by the
+        client that is accepted by the server. Subclasses can override this
+        method to implement a different subprotocol negotiation algorithm.
+
+        :param request: A ``Request`` object.
+
+        The method should return the subprotocol to use, or ``None`` if no
+        subprotocol is chosen.
+        """
+        ...
+
+class Client(Base):
+    """This class implements a WebSocket client.
+
+    Instead of creating an instance of this class directly, use the
+    ``connect()`` class method to create an instance that is connected to a
+    server.
+    """
+    def __init__(
+        self,
+        url,
+        subprotocols=...,
+        headers=...,
+        receive_bytes=...,
+        ping_interval=...,
+        max_message_size=...,
+        ssl_context=...,
+        thread_class=...,
+        event_class=...,
+    ) -> None: ...
+    @classmethod
+    def connect(
+        cls,
+        url: str,
+        subprotocols=...,
+        headers=...,
+        receive_bytes=...,
+        ping_interval=...,
+        max_message_size=...,
+        ssl_context=...,
+        thread_class=...,
+        event_class=...,
+    ):  # -> Self:
+        """Returns a WebSocket client connection.
+
+        :param url: The connection URL. Both ``ws://`` and ``wss://`` URLs are
+                    accepted.
+        :param subprotocols: The name of the subprotocol to use, or a list of
+                             subprotocol names in order of preference. Set to
+                             ``None`` (the default) to not use a subprotocol.
+        :param headers: A dictionary or list of tuples with additional HTTP
+                        headers to send with the connection request. Note that
+                        custom headers are not supported by the WebSocket
+                        protocol, so the use of this parameter is not
+                        recommended.
+        :param receive_bytes: The size of the receive buffer, in bytes. The
+                              default is 4096.
+        :param ping_interval: Send ping packets to the server at the requested
+                              interval in seconds. Set to ``None`` (the
+                              default) to disable ping/pong logic. Enable to
+                              prevent disconnections when the line is idle for
+                              a certain amount of time, or to detect an
+                              unresponsive server and disconnect. A recommended
+                              interval is 25 seconds. In general it is
+                              preferred to enable ping/pong on the server, and
+                              let the client respond with pong (which it does
+                              regardless of this setting).
+        :param max_message_size: The maximum size allowed for a message, in
+                                 bytes, or ``None`` for no limit. The default
+                                 is ``None``.
+        :param ssl_context: An ``SSLContext`` instance, if a default SSL
+                            context isn't sufficient.
+        :param thread_class: The ``Thread`` class to use when creating
+                             background threads. The default is the
+                             ``threading.Thread`` class from the Python
+                             standard library.
+        :param event_class: The ``Event`` class to use when creating event
+                            objects. The default is the `threading.Event``
+                            class from the Python standard library.
+        """
+        ...
+
+    def handshake(self):  # -> None:
+        ...
+    def close(self, reason=..., message=...):  # -> None:
+        ...