Make heatmap on GPU

Still has some weird artifacts, some axis ordering might be incorrect

where_fi/cli/__init__.py

@@ -38,9 +38,12 @@ def antennas() -> None:
 def heatmap() -> None:
     preprocessor = Preprocessor()
     aoa = AoA()
-    manager = mp.Manager()
+
-    webapp_queue: "mp.Queue[AoA]" = manager.Queue(config.sample_rate)
     # Start webapp in background process
+    manager = mp.Manager()
+    webapp_queue: "mp.Queue[npt.NDArray[np.float32]]" = manager.Queue(
+        config.sample_rate
+    )
     webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
     webapp.start()
 
@@ -50,8 +53,10 @@ def heatmap() -> None:
         processed_tensor = torch.tensor(processed, device=device)
         # visualise.add_data(all_data, processed)
         aoa.update(processed_tensor)
+
         if not webapp_queue.full():
-            webapp_queue.put(aoa)
+            heatmap = aoa.heatmap()
+            webapp_queue.put(heatmap)
 
     globals.csi_producer(csi_callback=callback)
     logger.info("Finished processing CSI data")

where_fi/config/models.py

@@ -23,6 +23,13 @@ class MUSIC(BaseModel):
     eigval_threshold: int
     window_size: int
 
+    class Heatmap(BaseModel):
+        theta_resolution: int
+        tof_resolution: int
+        tof_max: float
+
+    heatmap: Heatmap
+
 
 class Antennas(BaseModel):
     spacing: float

where_fi/processing/aoa.py

@@ -68,53 +68,93 @@ class AoA:
         if self.historical_autocorr.shape[0] > WINDOW_SIZE:
             self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]
 
-        # Is the moving average also Hermitian?
+    def steering_vector(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
-        R = torch.mean(self.historical_autocorr, dim=0)
+        assert theta.shape == tof.shape
+        assert len(theta.shape) == 1
+        N = theta.shape[0]
 
-        # The smallest eigenvectors span the noise subspace,
+        omega_t: torch.Tensor = torch.exp(-2j * np.pi * config.delta_f * tof)
-        # and the largest span the signal subspace.
+        phi_theta: torch.Tensor = torch.exp(
-        eigvals, eigvecs = torch.linalg.eigh(R)
-        self.E_n = (
-            eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold].cpu().numpy()
-        )
-
-    def steering_vector(
-        self, theta: float, tof: float
-    ) -> npt.NDArray[np.complexfloating]:
-        omega_t: npt.NDArray[np.complex64] = np.exp(-2j * np.pi * config.delta_f * tof)
-        phi_theta: npt.NDArray[np.complex64] = np.exp(
             2j
             * np.pi
             * config.central_freq_hz
             * config.antennas.spacing
-            * (1 - np.cos(theta))
+            * (1 - torch.cos(theta))
             / 299_792_458
         )
+        assert omega_t.shape == phi_theta.shape == (N,)
+        print(omega_t, phi_theta)
 
-        omega_t = np.expand_dims(omega_t, axis=-1)
+        omega_t = torch.unsqueeze(omega_t, dim=-1)
-        phi_theta = np.expand_dims(phi_theta, axis=-1)
+        phi_theta = torch.unsqueeze(phi_theta, dim=-1)
 
-        antenna_v = omega_t ** torch.arange(self.N_subcarriers // 2)
+        assert omega_t.shape == phi_theta.shape == (N, 1)
-        phis = phi_theta ** torch.arange(self.N_rx // 2)
-        antenna_v = np.expand_dims(antenna_v, axis=-1)
-        print(antenna_v.shape, phis.shape)
-        steering = antenna_v[0] * phis
-        print(steering.shape)
-        return steering.T.reshape(-1)
 
-    def evaluate(self, theta: float, tof: float) -> float:
+        antenna_v = omega_t ** torch.arange(
-        try:
+            self.N_subcarriers // 2, dtype=torch.float32
-            steering = self.steering_vector(theta, tof)
+        )
-            steering_h = np.conj(steering).T
+        phis = phi_theta ** torch.arange(self.N_rx // 2, dtype=torch.float32)
-        except Exception as e:
-            logger.exception(e)
-            return 0
 
-        assert isinstance(self.E_n, torch.Tensor)
+        assert antenna_v.shape == (N, self.N_subcarriers // 2)
-        E_n = self.E_n
+        assert phis.shape == (N, self.N_rx // 2)
-        E_n_H = np.conj(E_n).T
+
-        c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
+        antenna_v = torch.unsqueeze(antenna_v, dim=1)
-        return np.abs(c.real).item()
+        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.
+        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}")
+        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}"
+        )
+        c: torch.Tensor = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
+        return torch.abs(c.real)
+
+    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)
+        heatmap: npt.NDArray[np.float32] = (
+            self.evaluate(
+                torch.tensor(thetas_mesh.reshape(-1)),
+                torch.tensor(tofs_mesh.reshape(-1)),
+            )
+            .reshape(
+                config.music.heatmap.theta_resolution,
+                config.music.heatmap.tof_resolution,
+            )
+            .numpy(force=True)
+        )
+        return heatmap
 
 
 def test_smoothing() -> None:
@@ -133,8 +173,8 @@ def test_smoothing() -> None:
 def test_steering_vector() -> None:
     aoa = AoA()
     aoa.N_subcarriers = 10
-    aoa.N_rx = 2
+    aoa.N_rx = 4
-    tau = 1
+    tau = torch.Tensor([1, 0])
-    theta = 0
+    theta = torch.Tensor([0, 1])
     print(aoa.steering_vector(theta, tau))
     assert False

where_fi/visualise/__init__.py

@@ -13,7 +13,7 @@ from flask import Flask, Response, render_template, request
 from flask_sock import Sock
 from simple_websocket import Server
 
-from .. import config
+from ..config import config
 from ..processing.aoa import AoA
 
 matplotlib.use("agg")
@@ -92,18 +92,17 @@ def add_data(
         del subscriber_settings[subscriber]
 
 
-def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
+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)
-    r = np.linspace(0, max_tof, 100)  # Radius values
+    r = np.linspace(
-    theta = np.linspace(0, np.pi, 50)  # Angle values
+        0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
-    R, Theta = np.meshgrid(r, theta)  # Create a 2D grid of r and theta
+    )
+    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
 
-    # Compute the function values
+    ax.pcolormesh(Y, X, heatmap, edgecolors="face")
-    Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))
-
-    ax.pcolormesh(Theta, R, Z, edgecolors="face")
     buf = io.BytesIO()
     fig.savefig(buf, format="jpeg")
     plt.close(fig)
@@ -112,36 +111,23 @@ def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
     return buf
 
 
-def gather_aoa(max_tof: float) -> Generator[bytes, None, None]:
+def gather_aoa() -> Generator[bytes, None, None]:
     assert aoa_queue is not None
 
-    prev_frame = datetime.now()
     while True:
-        while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
-            time.sleep(0.01)
-        while not aoa_queue.empty():
         logger.debug("Receiving from aoa pipe")
-            aoa = aoa_queue.get()
+        heatmap = aoa_queue.get()
-        prev_frame = datetime.now()
+        buf = plot_heatmap(heatmap)
-        logger.debug(f"Generating heatmap of time {aoa.timestamp}")
-        buf = make_heatmap(aoa, max_tof)
         yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
         buf.close()
 
 
 @app.route("/aoa_tof")
 def aoa_tof() -> Response:
-    max_tof_str = request.args.get("max_tof")
+    return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame")
-    try:
-        max_tof = float(max_tof_str)
-    except Exception:
-        max_tof = 5e-8
-    return Response(
-        gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"
-    )
 
 
-def start(conn: "mp.Queue[AoA]") -> None:
+def start(conn: "mp.Queue[npt.NDArray[np.float32]]") -> None:
     global app, aoa_queue
 
     aoa_queue = conn