342357bd0c
The original assumption that we can use eigh because the autocorrelation matrix is Hermitian is not true. Although the true autocorrelation matrix is Hermitian, this is not necessarily the case for the estimate we use.
186 lines
6.2 KiB
Python
186 lines
6.2 KiB
Python
import logging
|
|
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import torch
|
|
|
|
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 = 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) -> 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)
|
|
assert isinstance(eigvals, torch.Tensor)
|
|
assert isinstance(eigvecs, torch.Tensor)
|
|
logger.info(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) -> 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)),
|
|
)
|
|
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)
|
|
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
|