Use PyTorch instead of NumPy #10

Merged
cfalas merged 5 commits from torch into main 2025-01-30 13:58:26 +02:00
2 changed files with 20 additions and 15 deletions
Showing only changes of commit 42ac105240 - Show all commits

View File

@ -50,8 +50,8 @@ def heatmap() -> 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)
logger.info(f"Processed CSI data with shape {processed.shape}")
processed_tensor = torch.tensor(processed, device=device)
# visualise.add_data(all_data, processed)
aoa.update(processed_tensor)
if not webapp_queue.full():

View File

@ -52,6 +52,7 @@ class AoA:
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.
@ -68,6 +69,8 @@ class AoA:
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
@ -83,7 +86,6 @@ class AoA:
/ 299_792_458
)
assert omega_t.shape == phi_theta.shape == (N,)
print(omega_t, phi_theta)
omega_t = torch.unsqueeze(omega_t, dim=-1)
phi_theta = torch.unsqueeze(phi_theta, dim=-1)
@ -114,20 +116,22 @@ class AoA:
# 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.info(f"Signal subspace: {E_n.shape}")
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.info(
f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, {E_n_H.shape}, {steering.shape}"
logger.debug(
f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, "
f"{E_n_H.shape}, {steering.shape}"
)
c: torch.Tensor = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
return torch.abs(c.real)
@ -143,17 +147,18 @@ class AoA:
dtype=np.float32,
)
thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs)
heatmap: npt.NDArray[np.float32] = (
self.evaluate(
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)),
)
.reshape(
logger.debug(f"Evaluated heatmap: {evaluated.shape}")
heatmap: npt.NDArray[np.float32] = evaluated.reshape(
config.music.heatmap.theta_resolution,
config.music.heatmap.tof_resolution,
)
.numpy(force=True)
)
).numpy(force=True)
return heatmap