Add basic MUSIC implementation

src/aoa.py

@@ -1,17 +1,120 @@
 import numpy as np
 import numpy.typing as npt
-import config
+from . import config
+
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class AoA:
     def __init__(self):
-        self.historical_data = np.array([])
+        self.historical_autocorr = np.array([])
+        self.N_subcarriers = -1
+        self.N_rx = -1
         pass
 
+    def smooth(self, data: npt.NDArray[np.complex128]):
+        assert len(data.shape) == 3
+
+        M = data.shape[0]  # Number of subcarriers
+        N = data.shape[1]  # Number of RX antennas
+        T = data.shape[2]  # Number of TX antennas
+
+        self.N_subcarriers = M
+        self.N_rx = N
+        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
+        assert T == 1, "The current implementation only supports 1 TX antenna"
+
+        H_n = np.zeros((N, M // 2, M // 2 + 1), dtype=np.complex128)
+
+        for i in range(N):
+            for j in range(M // 2):
+                H_n[i, j] = data[j : j + M // 2 + 1, i, 0]
+
+        H_sm = np.vstack([H_n[0], H_n[1]])
+
+        # This would normally have to be arranged as follows:
+        # H_0 H_1 ... H_{N//2 - 1}
+        # H_1 H_2 ... H_{N//2}
+        # ...
+        # H_{N//2} ... H{N-1}
+        # But this doesn't work when we only have 2 receiving antennas
+        assert N == 2, "The current implementation only supports 2 RX antennas"
+
+        return H_sm
+
     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
+        H_sm = self.smooth(data)
+
+        auto_corr = np.matmul(H_sm, np.conj(H_sm).T)
+
+        if self.historical_autocorr.size == 0:
+            self.historical_autocorr = np.expand_dims(auto_corr, 0)
+        else:
+            self.historical_autocorr = np.append(
+                self.historical_autocorr, np.expand_dims(auto_corr, 0), axis=0
+            )
+
+        WINDOW_SIZE = config.AOA_SLIDING_WINDOW_SIZE
+        if self.historical_autocorr.shape[0] > WINDOW_SIZE:
+            self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]
+
+        R = np.mean(self.historical_autocorr, axis=0)
+        RR_h = np.matmul(R, np.conj(R).T)
+        # This matrix is by definition Hermitian.
+        # Therefore, all of its eigenvectors are orthogonal.
+
+        # The smallest eigenvectors span the noise subspace,
+        # and the largest span the signal subspace.
+        eigvals, eigvecs = np.linalg.eig(RR_h)
+        self.E_n = eigvecs[:, eigvals < config.EIGVAL_THRESHOLD]
+
+    omega_base = np.exp(-2j * np.pi * config.DELTA_F)
+    phi_base = np.exp(
+        2j * np.pi * config.CENTRAL_FREQUENCY_HZ * config.ANTENNA_SPACING / config.C
+    )
+
+    def steering_vector(self, theta: float, tof: float):
+        omega_t = self.omega_base**tof
+        phi_theta = self.phi_base ** (1 - np.cos(theta))
+
+        assert phi_theta.shape == omega_t.shape
+
+        omega_t = np.expand_dims(omega_t, axis=-1)
+        phi_theta = np.expand_dims(phi_theta, axis=-1)
+
+        antenna_v = omega_t ** np.arange(self.N_subcarriers // 2)
+        phis = phi_theta ** np.arange(self.N_rx)
+        antenna_v = np.expand_dims(antenna_v, axis=-1)
+        phis = np.expand_dims(phis, axis=-2)
+        steering = antenna_v * phis
+        return steering.reshape(-1)
+
+    def evaluate(self, theta: float, tof: float):
+        try:
+            steering = self.steering_vector(theta, tof)
+            steering_h = np.conj(steering).T
+        except Exception as e:
+            logger.exception(e)
+            return 0
+        E_n = self.E_n
+        E_n_H = np.conj(self.E_n).T
+        c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
+        return c.real
+
+
+def test_smoothing():
+    row, col = np.indices((4, 2))
+    data = row + 1j * col
+    aoa = AoA()
+    smoothed = aoa.smooth(data)
+    H_0 = np.array([[0 + 0j, 0 + 1j, 0 + 2j], [0 + 1j, 0 + 2j, 0 + 3j]])
+    H_01 = np.vstack([H_0, H_0 + 1])
+    H_12 = np.vstack([H_0 + 1, H_0 + 2])
+    expected = np.hstack([H_01, H_12])
+    print(expected)
+    assert np.allclose(smoothed, expected)
+    pass

src/visualise/__init__.py

@@ -1,14 +1,29 @@
-from flask import Flask, render_template
+from flask import Flask, render_template, Response
 from flask_sock import Sock
 import numpy as np
 import numpy.typing as npt
 from simple_websocket import Server
+import time
+
+import matplotlib.pyplot as plt
+import io
+from PIL import Image
+import logging
+
+from ..aoa import AoA
+
+import matplotlib
+
+matplotlib.use("agg")
 
 app = Flask(__name__)
 sock = Sock(app)
 
+logger = logging.getLogger(__name__)
+
 
 data: npt.NDArray[np.complex128] = np.array([], dtype=complex)
+aoa: AoA = AoA()
 
 
 @app.route("/preprocessed")
@@ -53,5 +68,38 @@ def add_data(new_data: npt.NDArray[np.complex128]):
         del subscriber_settings[subscriber]
 
 
+def make_heatmap():
+    fig = plt.figure()
+    ax = fig.add_axes([0, 0, 1, 1], polar=True)
+    r = np.linspace(0, 3e-8, 100)  # Radius values
+    theta = np.linspace(0, np.pi, 50)  # Angle values
+    R, Theta = np.meshgrid(r, theta)  # Create a 2D grid of r and theta
+
+    # Compute the function values
+    Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))
+
+    ax.pcolormesh(Theta, R, Z, edgecolors="face")
+    buf = io.BytesIO()
+    fig.savefig(buf, format="jpeg")
+    plt.close(fig)
+
+    buf.seek(0)
+    return buf
+
+
+def gather_aoa():
+    while True:
+        # time.sleep(0.05)
+        logger.info("Got AoA heatmap")
+        buf = make_heatmap()
+        yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
+        buf.close()
+
+
+@app.route("/aoa_tof")
+def aoa_tof():
+    return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame")
+
+
 def start():
     app.run(debug=True, use_reloader=False)