From 08ffb6cf5bb9e8a8357c0beb197aa55d87b6e718 Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Tue, 31 Dec 2024 16:45:42 +0000 Subject: [PATCH] set up torch AoA estimation --- pyproject.toml | 1 + where_fi/cli/__init__.py | 7 +++- where_fi/collection/csi_frame.py | 6 +-- where_fi/collection/protocols.py | 2 +- where_fi/processing/aoa.py | 69 +++++++++++++++++-------------- where_fi/processing/preprocess.py | 10 ++--- 6 files changed, 54 insertions(+), 41 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2b192c6..d78a074 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "h5py>=3.12.1", "pyyaml>=6.0.2", "pydantic>=2.10.6", + "torch" ] [project.scripts] diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py index 3d46786..5caa2b7 100644 --- a/where_fi/cli/__init__.py +++ b/where_fi/cli/__init__.py @@ -3,6 +3,7 @@ import multiprocessing as mp import numpy as np import numpy.typing as npt +import torch import typer from .. import visualise @@ -13,6 +14,7 @@ from . import file, globals app = typer.Typer(callback=globals.main) logger = logging.getLogger(__name__) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @app.command() @@ -42,11 +44,12 @@ def heatmap() -> None: webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) webapp.start() - def callback(antenna_data: npt.NDArray[np.complex128]) -> None: + def callback(antenna_data: npt.NDArray[np.complex64]) -> None: logger.info(f"Got final CSI data with shape {antenna_data.shape}") processed = preprocessor.preprocess(antenna_data) + processed_tensor = torch.tensor(processed, device=device) # visualise.add_data(all_data, processed) - aoa.update(processed) + aoa.update(processed_tensor) if not webapp_queue.full(): webapp_queue.put(aoa) diff --git a/where_fi/collection/csi_frame.py b/where_fi/collection/csi_frame.py index 83291a9..4d84eef 100644 --- a/where_fi/collection/csi_frame.py +++ b/where_fi/collection/csi_frame.py @@ -93,14 +93,14 @@ class CSIHeader: class CSI: @staticmethod - def parseCsiData(data: bytes, header: CSIHeader) -> npt.NDArray[np.complex128]: - csi_matrix: npt.NDArray[np.complex128] = np.zeros( + def parseCsiData(data: bytes, header: CSIHeader) -> npt.NDArray[np.complex64]: + csi_matrix: npt.NDArray[np.complex64] = np.zeros( ( header.num_subcarriers, header.num_rx, header.num_tx, ), - dtype=np.complex128, + dtype=np.complex64, ) pos = 0 for j in range(header.num_rx): diff --git a/where_fi/collection/protocols.py b/where_fi/collection/protocols.py index d58505a..825f7a3 100644 --- a/where_fi/collection/protocols.py +++ b/where_fi/collection/protocols.py @@ -3,7 +3,7 @@ from typing import Callable, Protocol import numpy as np import numpy.typing as npt -CSICallback = Callable[[npt.NDArray[np.complex128]], None] +CSICallback = Callable[[npt.NDArray[np.complex64]], None] class CSIProducer(Protocol): diff --git a/where_fi/processing/aoa.py b/where_fi/processing/aoa.py index c6c5371..604e9d6 100644 --- a/where_fi/processing/aoa.py +++ b/where_fi/processing/aoa.py @@ -3,21 +3,26 @@ from datetime import datetime import numpy as np import numpy.typing as npt +import torch +import torch.linalg from ..config import config logger = logging.getLogger(__name__) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +torch.set_default_device(device) + class AoA: def __init__(self) -> None: - self.historical_autocorr = np.array([]) + self.historical_autocorr = torch.tensor([], dtype=torch.complex64) self.N_subcarriers = -1 self.N_rx = -1 self.timestamp = datetime.now() pass - def smooth(self, data: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]: + def smooth(self, data: torch.Tensor) -> torch.Tensor: assert len(data.shape) == 3 M = data.shape[0] # Number of subcarriers @@ -32,32 +37,32 @@ class AoA: # 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) + H_n = torch.zeros((N, M // 2, M // 2 + 1), dtype=torch.complex64) 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_rows = [np.hstack(H_n[i : i + N // 2 + 1]) for i in range(N // 2)] - H_sm = np.vstack(H_sm_rows) + H_sm_rows = [torch.hstack(list(H_n[i : i + N // 2 + 1])) for i in range(N // 2)] + H_sm = torch.vstack(H_sm_rows) logger.debug(f"Smoothed: {H_sm.shape}") return H_sm - def update(self, data: npt.NDArray[np.complex128]) -> None: + def update(self, data: torch.Tensor) -> None: self.timestamp = datetime.now() H_sm = self.smooth(data) - auto_corr = np.matmul(H_sm, np.conj(H_sm).T) + auto_corr = H_sm @ torch.conj(H_sm).T # This matrix is by definition Hermitian. # Therefore, all of its eigenvectors are orthogonal. - if self.historical_autocorr.size == 0: - self.historical_autocorr = np.expand_dims(auto_corr, 0) + if len(self.historical_autocorr.shape) <= 1: + self.historical_autocorr = torch.unsqueeze(auto_corr, 0) else: - self.historical_autocorr = np.append( - self.historical_autocorr, np.expand_dims(auto_corr, 0), axis=0 + self.historical_autocorr = torch.cat( + (self.historical_autocorr, torch.unsqueeze(auto_corr, 0)) ) WINDOW_SIZE = config.music.window_size @@ -65,7 +70,7 @@ class AoA: self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:] # Is the moving average also Hermitian? - R = np.mean(self.historical_autocorr, axis=0) + R = torch.mean(self.historical_autocorr, dim=0) # The smallest eigenvectors span the noise subspace, # and the largest span the signal subspace. @@ -75,8 +80,8 @@ class AoA: def steering_vector( self, theta: float, tof: float ) -> npt.NDArray[np.complexfloating]: - omega_t: npt.NDArray[np.complex128] = np.exp(-2j * np.pi * config.delta_f * tof) - phi_theta: npt.NDArray[np.complex128] = np.exp( + omega_t: npt.NDArray[np.complex64] = np.exp(-2j * np.pi * config.delta_f * tof) + phi_theta: npt.NDArray[np.complex64] = np.exp( 2j * np.pi * config.central_freq_hz @@ -85,39 +90,43 @@ class AoA: / 299_792_458 ) - omega_t = np.expand_dims(omega_t, axis=-1) - phi_theta = np.expand_dims(phi_theta, axis=-1) + omega_t = torch.unsqueeze(omega_t, dim=-1) + phi_theta = torch.unsqueeze(phi_theta, dim=-1) - antenna_v = omega_t ** np.arange(self.N_subcarriers // 2) - phis = phi_theta ** np.arange(self.N_rx // 2) - antenna_v = np.expand_dims(antenna_v, axis=-1) - steering = antenna_v * phis + antenna_v = omega_t ** torch.arange(self.N_subcarriers // 2) + phis = phi_theta ** torch.arange(self.N_rx // 2) + antenna_v = torch.unsqueeze(antenna_v, dim=-1) + print(antenna_v.shape, phis.shape) + steering = antenna_v[0] * phis + print(steering.shape) return steering.T.reshape(-1) def evaluate(self, theta: float, tof: float) -> float: try: steering = self.steering_vector(theta, tof) - steering_h = np.conj(steering).T + steering_h = torch.conj(steering).T except Exception as e: logger.exception(e) return 0 + + assert isinstance(self.E_n, torch.Tensor) E_n = self.E_n - E_n_H = np.conj(E_n).T + E_n_H = torch.conj(E_n).T c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering)) - return np.abs(c.real) + return torch.abs(c.real) def test_smoothing() -> None: - row, col = np.indices((4, 2)) + row, col = torch.indices((6, 4)) data = row + 1j * col + data = np.expand_dims(data, axis=2) + np.set_printoptions(linewidth=200) + print(data.shape) aoa = AoA() + aoa.N_subcarriers = 6 + aoa.N_rx = 4 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) + print(smoothed) def test_steering_vector() -> None: diff --git a/where_fi/processing/preprocess.py b/where_fi/processing/preprocess.py index e2c8508..01c4718 100644 --- a/where_fi/processing/preprocess.py +++ b/where_fi/processing/preprocess.py @@ -15,9 +15,9 @@ np.seterr(invalid="ignore") class Preprocessor: def __init__(self) -> None: - self.prev_entries: Queue[npt.NDArray[np.complex128]] = Queue(maxsize=100) - self.short_term_avg = np.zeros((1,), dtype=np.complex128) - self.long_term_avg = np.zeros((1,), dtype=np.complex128) + self.prev_entries: Queue[npt.NDArray[np.complex64]] = Queue(maxsize=100) + self.short_term_avg = np.zeros((1,), dtype=np.complex64) + self.long_term_avg = np.zeros((1,), dtype=np.complex64) self.filter = butter( 5, config.preprocessing.bandpass.bounds, @@ -26,7 +26,7 @@ class Preprocessor: output="sos", ) - def preprocess(self, h: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]: + def preprocess(self, h: npt.NDArray[np.complex64]) -> npt.NDArray[np.complex64]: # CSI data is not available for pilot subcarriers. h_hat = np.where( np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)), @@ -46,7 +46,7 @@ class Preprocessor: # Assume that all csi matrices will have the same shape if self.long_term_avg.shape != h_hat.shape: - self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex128) + self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex64) self.long_term_avg = ( self.long_term_avg * (1 - config.preprocessing.moving_average_alpha)