From 1bc44ae43a34b52b099599e169b83b673e879df0 Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Thu, 30 Jan 2025 11:28:50 +0000 Subject: [PATCH] fix heatmap generation --- where_fi/processing/aoa.py | 6 +++--- where_fi/visualise/__init__.py | 7 ++++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/where_fi/processing/aoa.py b/where_fi/processing/aoa.py index d8303c1..a03bb0f 100644 --- a/where_fi/processing/aoa.py +++ b/where_fi/processing/aoa.py @@ -134,8 +134,8 @@ class AoA: 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) + 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( @@ -157,8 +157,8 @@ class AoA: ) logger.debug(f"Evaluated heatmap: {evaluated.shape}") heatmap: npt.NDArray[np.float32] = evaluated.reshape( - config.music.heatmap.theta_resolution, config.music.heatmap.tof_resolution, + config.music.heatmap.theta_resolution, ).numpy(force=True) return heatmap diff --git a/where_fi/visualise/__init__.py b/where_fi/visualise/__init__.py index 611e945..a95791c 100644 --- a/where_fi/visualise/__init__.py +++ b/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)