bins, sensors, _ = target_psd_matrix.shape
beamforming_vector = np.empty((bins, sensors), dtype=np.complex)
for f in range(bins):
eigenvals, eigenvecs = eigh(target_psd_matrix[f, :, :])
beamforming_vector[f, :] = eigenvecs[:, np.argmax(eigenvals)]
return beamforming_vector
After Change
// Find max eigenvals
vals = np.argmax(eigenvals, axis=-1)
// Select eigenvec for max eigenval
beamforming_vector = np.array([eigenvecs[i, :, vals[i]] for i in range(eigenvals.shape[0])])
// Reconstruct original shape
beamforming_vector = np.reshape(beamforming_vector, shape[:-1])