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 7 additions and 6 deletions
Showing only changes of commit 638440feff - Show all commits

View File

@ -134,8 +134,8 @@ class AoA:
f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, " f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, "
f"{E_n_H.shape}, {steering.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]: def heatmap(self) -> npt.NDArray[np.float32]:
thetas = np.linspace( thetas = np.linspace(
@ -157,8 +157,8 @@ class AoA:
) )
logger.debug(f"Evaluated heatmap: {evaluated.shape}") logger.debug(f"Evaluated heatmap: {evaluated.shape}")
heatmap: npt.NDArray[np.float32] = evaluated.reshape( heatmap: npt.NDArray[np.float32] = evaluated.reshape(
config.music.heatmap.theta_resolution,
config.music.heatmap.tof_resolution, config.music.heatmap.tof_resolution,
config.music.heatmap.theta_resolution,
).numpy(force=True) ).numpy(force=True)
return heatmap return heatmap

View File

@ -95,14 +95,15 @@ def add_data(
def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO: def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO:
logger.info(f"Making heatmap with aoa of {aoa.timestamp}") logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
fig = plt.figure() 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( r = np.linspace(
0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution 0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
) )
theta = np.linspace(0, np.pi, config.music.heatmap.theta_resolution) # Angle values 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() buf = io.BytesIO()
fig.savefig(buf, format="jpeg") fig.savefig(buf, format="jpeg")
plt.close(fig) plt.close(fig)