diff --git a/src/ria_toolkit_oss/viz/recording.py b/src/ria_toolkit_oss/viz/recording.py index 62422cd..9c166a6 100644 --- a/src/ria_toolkit_oss/viz/recording.py +++ b/src/ria_toolkit_oss/viz/recording.py @@ -189,3 +189,155 @@ def constellation(rec: Recording) -> Figure: ) return fig + + +def power_spectral_density(rec: Recording) -> Figure: + """Create a Power Spectral Density (PSD) plot from the recording. + + :param rec: Input signal to plot. + :type rec: ria_toolkit_oss.datatypes.Recording + + :return: PSD plot, as a Plotly Figure. + """ + complex_signal = rec.data[0] + center_frequency = int(rec.metadata.get("center_frequency", 0)) + sample_rate = int(rec.metadata.get("sample_rate", 1)) + + # Calculate PSD using Welch's method + frequencies, psd = signal.welch( + complex_signal, + fs=sample_rate, + nperseg=min(1024, len(complex_signal)), + return_onesided=False, + scaling="density", + ) + + # Shift frequencies and PSD for proper visualization + frequencies_shifted = fftshift(frequencies) + center_frequency + psd_shifted = fftshift(psd) + + # Convert to dB scale + psd_db = 10 * np.log10(psd_shifted + 1e-10) + + fig = go.Figure() + fig.add_trace( + go.Scatter(x=frequencies_shifted, y=psd_db, mode="lines", name="PSD", line=dict(width=0.8, color="#00D9FF")) + ) + + fig.update_layout( + title="Power Spectral Density", + xaxis_title="Frequency [Hz]", + yaxis_title="Power/Frequency [dB/Hz]", + template="plotly_dark", + height=300, + width=800, + showlegend=False, + ) + + return fig + + +def fft_plot(rec: Recording) -> Figure: + """Create an FFT magnitude plot from the recording. + + :param rec: Input signal to plot. + :type rec: ria_toolkit_oss.datatypes.Recording + + :return: FFT plot, as a Plotly Figure. + """ + complex_signal = rec.data[0] + center_frequency = int(rec.metadata.get("center_frequency", 0)) + sample_rate = int(rec.metadata.get("sample_rate", 1)) + + # Compute FFT + fft_result = fftshift(fft(complex_signal)) + freqs = fftshift(np.fft.fftfreq(len(complex_signal), 1 / sample_rate)) + center_frequency + + # Convert to magnitude in dB + magnitude = np.abs(fft_result) + magnitude_db = 20 * np.log10(magnitude + 1e-10) + + fig = go.Figure() + fig.add_trace(go.Scatter(x=freqs, y=magnitude_db, mode="lines", name="FFT", line=dict(width=0.6, color="#FF6B9D"))) + + fig.update_layout( + title="FFT Magnitude", + xaxis_title="Frequency [Hz]", + yaxis_title="Magnitude [dB]", + template="plotly_dark", + height=300, + width=800, + showlegend=False, + ) + + return fig + + +def spectrogram_3d(rec: Recording) -> Figure: + """Create a 3D spectrogram plot from the recording. + + :param rec: Input signal to plot. + :type rec: ria_toolkit_oss.datatypes.Recording + + :return: 3D Spectrogram, as a Plotly Figure. + """ + complex_signal = rec.data[0] + sample_rate = int(rec.metadata.get("sample_rate", 1)) + plot_length = len(complex_signal) + + # Determine FFT size + if plot_length < 2000: + fft_size = 64 + elif plot_length < 10000: + fft_size = 256 + elif plot_length < 1000000: + fft_size = 1024 + else: + fft_size = 2048 + + frequencies, times, Sxx = signal.spectrogram( + complex_signal, + fs=sample_rate, + nfft=fft_size, + nperseg=fft_size, + noverlap=fft_size // 8, + scaling="density", + mode="complex", + return_onesided=False, + ) + + # Convert complex values to amplitude and then to log scale + Sxx_magnitude = np.abs(Sxx) + Sxx_log = 10 * np.log10(Sxx_magnitude + 1e-10) + + # Shift frequency bins for proper visualization + frequencies_shifted = np.fft.fftshift(frequencies) + Sxx_shifted = np.fft.fftshift(Sxx_log, axes=0) + + fig = go.Figure( + data=[ + go.Surface( + z=Sxx_shifted, + x=times, + y=frequencies_shifted, + colorscale="Viridis", + showscale=True, + colorbar=dict(title="Power [dB]"), + ) + ] + ) + + fig.update_layout( + title="3D Spectrogram", + scene=dict( + xaxis_title="Time [s]", + yaxis_title="Frequency [Hz]", + zaxis_title="Power [dB]", + camera=dict(eye=dict(x=1.5, y=1.5, z=1.3)), + ), + template="plotly_dark", + height=600, + width=900, + ) + + return fig