8fe5995c5f
Split off figure-serving from the data transformations needed to make the figures themselves (e.g. binning)
203 lines
6.7 KiB
Python
203 lines
6.7 KiB
Python
import logging
|
|
from datetime import datetime
|
|
from typing import Any, Callable
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import torch
|
|
|
|
from ..config import config
|
|
from ..visualise import server as visualise
|
|
|
|
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 = torch.tensor([], dtype=torch.complex64)
|
|
self.N_subcarriers = -1
|
|
self.N_rx = -1
|
|
self.timestamp = datetime.now()
|
|
pass
|
|
|
|
def smooth(self, data: torch.Tensor) -> torch.Tensor:
|
|
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 = 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 = [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: torch.Tensor) -> None:
|
|
self.timestamp = datetime.now()
|
|
H_sm = self.smooth(data)
|
|
logger.debug(f"Calculated smoothed CSI matrix: {H_sm.shape}")
|
|
|
|
auto_corr = H_sm @ torch.conj(H_sm).T
|
|
# This matrix is by definition Hermitian.
|
|
# Therefore, all of its eigenvectors are orthogonal.
|
|
|
|
if len(self.historical_autocorr.shape) <= 1:
|
|
self.historical_autocorr = torch.unsqueeze(auto_corr, 0)
|
|
else:
|
|
self.historical_autocorr = torch.cat(
|
|
(self.historical_autocorr, torch.unsqueeze(auto_corr, 0))
|
|
)
|
|
|
|
WINDOW_SIZE = config.music.window_size
|
|
if self.historical_autocorr.shape[0] > WINDOW_SIZE:
|
|
self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]
|
|
|
|
logger.debug("Finished updating autocorrelation matrix")
|
|
|
|
def steering_vector(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
|
|
assert theta.shape == tof.shape
|
|
assert len(theta.shape) == 1
|
|
N = theta.shape[0]
|
|
|
|
omega_t: torch.Tensor = torch.exp(-2j * np.pi * config.delta_f * tof)
|
|
phi_theta: torch.Tensor = torch.exp(
|
|
2j
|
|
* np.pi
|
|
* config.central_freq_hz
|
|
* config.antennas.spacing
|
|
* (1 - torch.cos(theta))
|
|
/ 299_792_458
|
|
)
|
|
assert omega_t.shape == phi_theta.shape == (N,)
|
|
|
|
omega_t = torch.unsqueeze(omega_t, dim=-1)
|
|
phi_theta = torch.unsqueeze(phi_theta, dim=-1)
|
|
|
|
assert omega_t.shape == phi_theta.shape == (N, 1)
|
|
|
|
antenna_v = omega_t ** torch.arange(
|
|
self.N_subcarriers // 2, dtype=torch.float32
|
|
)
|
|
phis = phi_theta ** torch.arange(self.N_rx // 2, dtype=torch.float32)
|
|
|
|
assert antenna_v.shape == (N, self.N_subcarriers // 2)
|
|
assert phis.shape == (N, self.N_rx // 2)
|
|
|
|
antenna_v = torch.unsqueeze(antenna_v, dim=1)
|
|
phis = torch.unsqueeze(phis, dim=-1)
|
|
|
|
assert antenna_v.shape == (N, 1, self.N_subcarriers // 2)
|
|
assert phis.shape == (N, self.N_rx // 2, 1)
|
|
|
|
steering = torch.bmm(phis, antenna_v)
|
|
|
|
assert steering.shape == (N, self.N_rx // 2, self.N_subcarriers // 2)
|
|
return steering.reshape(N, -1)
|
|
|
|
def evaluate(
|
|
self,
|
|
theta: torch.Tensor,
|
|
tof: torch.Tensor,
|
|
visualiser: None
|
|
| Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
|
|
) -> torch.Tensor:
|
|
R = torch.mean(self.historical_autocorr, dim=0)
|
|
|
|
# The smallest eigenvectors span the noise subspace,
|
|
# and the largest span the signal subspace.
|
|
logger.debug(f"Calculating eigenvectors of R: {R.shape}")
|
|
eigvals, eigvecs = torch.linalg.eig(R)
|
|
if visualiser:
|
|
visualiser(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
|
|
assert isinstance(eigvals, torch.Tensor)
|
|
assert isinstance(eigvecs, torch.Tensor)
|
|
logger.debug(f"Eigenvalues: {eigvals}")
|
|
E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]
|
|
|
|
logger.debug(f"Signal subspace: {E_n.shape}")
|
|
steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1)
|
|
steering_h = torch.conj(steering).permute(0, 2, 1)
|
|
|
|
E_n = E_n.unsqueeze(0)
|
|
E_n_H = torch.conj(E_n).permute(0, 2, 1)
|
|
logger.debug(
|
|
f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, "
|
|
f"{E_n_H.shape}, {steering.shape}"
|
|
)
|
|
c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
|
|
return torch.abs(c)[:, 0, 0]
|
|
|
|
def heatmap(
|
|
self,
|
|
visualiser: None | Callable[[npt.NDArray[Any], visualise.figures.Figure], None],
|
|
) -> npt.NDArray[np.float32]:
|
|
thetas = np.linspace(
|
|
0, np.pi, config.music.heatmap.theta_resolution, dtype=np.float32
|
|
)
|
|
tofs = np.linspace(
|
|
0,
|
|
config.music.heatmap.tof_max,
|
|
config.music.heatmap.tof_resolution,
|
|
dtype=np.float32,
|
|
)
|
|
thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs)
|
|
logger.debug(
|
|
f"Calculating heatmap with {thetas_mesh.shape} and {tofs_mesh.shape}"
|
|
)
|
|
evaluated = self.evaluate(
|
|
torch.tensor(thetas_mesh.reshape(-1)),
|
|
torch.tensor(tofs_mesh.reshape(-1)),
|
|
visualiser=visualiser,
|
|
)
|
|
logger.debug(f"Evaluated heatmap: {evaluated.shape}")
|
|
heatmap: npt.NDArray[np.float32] = evaluated.reshape(
|
|
config.music.heatmap.tof_resolution,
|
|
config.music.heatmap.theta_resolution,
|
|
).numpy(force=True)
|
|
|
|
if visualiser:
|
|
visualiser(heatmap, visualise.figures.Figure.AOA_HEATMAP)
|
|
return heatmap
|
|
|
|
|
|
def test_smoothing() -> None:
|
|
row, col = np.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)
|
|
print(smoothed)
|
|
|
|
|
|
def test_steering_vector() -> None:
|
|
aoa = AoA()
|
|
aoa.N_subcarriers = 10
|
|
aoa.N_rx = 4
|
|
tau = torch.Tensor([1, 0])
|
|
theta = torch.Tensor([0, 1])
|
|
print(aoa.steering_vector(theta, tau))
|
|
assert False
|