From a7a34a613879f642314284797be5db636b861096 Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Thu, 26 Dec 2024 14:01:24 +0000 Subject: [PATCH] Add basic MUSIC implementation --- src/aoa.py | 119 +++++++++++++++++++++++++++++++++++--- src/visualise/__init__.py | 50 +++++++++++++++- 2 files changed, 160 insertions(+), 9 deletions(-) diff --git a/src/aoa.py b/src/aoa.py index 91340ee..68442c3 100644 --- a/src/aoa.py +++ b/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 diff --git a/src/visualise/__init__.py b/src/visualise/__init__.py index 05732b6..c58ab85 100644 --- a/src/visualise/__init__.py +++ b/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)