diff --git a/src/python/analysis/waterfall.py b/src/python/analysis/waterfall.py new file mode 100644 index 0000000..d60f9ed --- /dev/null +++ b/src/python/analysis/waterfall.py @@ -0,0 +1,62 @@ +import sys + +import numpy as np +import matplotlib.pyplot as plt +import plotly.graph_objects as go +from scipy.signal import stft +from scipy.io import wavfile + +fs, rir = wavfile.read(sys.argv[1]) +rir = rir / np.max(np.abs(rir)) # Normalize +nfft = 1024 +frequencies, times, Zxx = stft(rir, fs=fs, nperseg=nfft) + +decay_time = np.zeros(frequencies.shape) +for i, freq in enumerate(frequencies): + magnitude = np.abs(Zxx[i, :]) + magnitude_db = 20 * np.log10(magnitude / np.max(magnitude)) + try: + decay_index = np.where(magnitude_db <= -20)[0][0] + decay_time[i] = times[decay_index] + except IndexError: + decay_time[i] = times[-1] # If it never decays by 60 dB + + +# Waterfall Plot +X, Y = np.meshgrid(times, frequencies) +Z = 20 * np.log10(np.abs(Zxx)/nfft) +Z -= np.max(Z) +Z = np.clip(Z, -60, 0) +# Z[Z<-60] = -100 + +fig = go.Figure(data=[go.Surface( + colorscale='viridis', + x=X, + y=Y, + z=Z +)]) +fig.update_layout( + title='Decay Times', + scene=dict( + xaxis_title='X: Time [s]', + yaxis_title='Y: Frequency [Hz]', + zaxis_title='Z: Amplitude [dB]', + + xaxis=dict(range=[0, 1],), + yaxis=dict(type='log',), + # zaxis=dict(range=[-60, 0],), + ), +) + +Z = 20 * np.log10(np.abs(Zxx/nfft)) +Z -= np.max(Z) + +plt.figure(figsize=(10, 6)) +plt.pcolormesh(times, frequencies, Z, shading='gouraud', vmin=-96, vmax=0) +plt.colorbar(label='Amplitude [dB]') +plt.title('Decay Times') +plt.xlabel('Time [s]') +plt.ylabel('Frequency [Hz]') +plt.yscale('log') +plt.ylim([frequencies[1], fs / 2]) +plt.show()