fix heatmap generation

The mean of Hermitian matrices is also Hermitian

where_fi/cli/__init__.py

@@ -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():

where_fi/processing/aoa.py

@@ -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)
@@ -111,26 +113,29 @@ class AoA:
 
     def evaluate(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
         R = torch.mean(self.historical_autocorr, dim=0)
+        assert (R == torch.conj(R).T).all()
 
         # The smallest eigenvectors span the noise subspace,
         # and the largest span the signal subspace.
-        eigvals, eigvecs = torch.linalg.eig(R)
+        logger.debug(f"Calculating eigenvectors of R: {R.shape}")
+        eigvals, eigvecs = torch.linalg.eigh(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(
+        logger.debug(
-            f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, {E_n_H.shape}, {steering.shape}"
+            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))
+        c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
-        return torch.abs(c.real)
+        return torch.abs(c)[:, 0, 0]
 
     def heatmap(self) -> npt.NDArray[np.float32]:
         thetas = np.linspace(
@@ -143,17 +148,18 @@ class AoA:
             dtype=np.float32,
         )
         thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs)
-        heatmap: npt.NDArray[np.float32] = (
+        logger.debug(
-            self.evaluate(
+            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}")
-                config.music.heatmap.theta_resolution,
+        heatmap: npt.NDArray[np.float32] = evaluated.reshape(
             config.music.heatmap.tof_resolution,
-            )
+            config.music.heatmap.theta_resolution,
-            .numpy(force=True)
+        ).numpy(force=True)
-        )
         return heatmap
 
 

where_fi/visualise/__init__.py

@@ -95,14 +95,15 @@ def add_data(
 def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO:
     logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
     fig = plt.figure()
-    ax = fig.add_axes([0, 0, 1, 1], polar=True)
+    ax = fig.add_axes([0.1, 0.1, 0.9, 0.9])  # , polar=True)
     r = np.linspace(
         0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
     )
     theta = np.linspace(0, np.pi, config.music.heatmap.theta_resolution)  # Angle values
-    X, Y = np.meshgrid(r, theta)  # Create a 2D grid of r and theta
 
-    ax.pcolormesh(Y, X, heatmap, edgecolors="face")
+    mesh = ax.pcolormesh(theta, r, heatmap, edgecolors="face", vmin=0, vmax=50)
+
+    fig.colorbar(mesh, ax=ax)
     buf = io.BytesIO()
     fig.savefig(buf, format="jpeg")
     plt.close(fig)