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 = config.subcarriers self.N_rx = config.antennas.count 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 * (torch.sin(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.info(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