fix heatmap generation

This commit is contained in:
Christos Falas 2025-01-30 11:28:50 +00:00
parent 44080bf833
commit 1bc44ae43a
No known key found for this signature in database
2 changed files with 7 additions and 6 deletions

View File

@ -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

View File

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