set up torch AoA estimation
This commit is contained in:
parent
814321f524
commit
08ffb6cf5b
@ -15,6 +15,7 @@ dependencies = [
|
||||
"h5py>=3.12.1",
|
||||
"pyyaml>=6.0.2",
|
||||
"pydantic>=2.10.6",
|
||||
"torch"
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user