Merge pull request 'PSD, FFT, and Spectrogram 3D' (#14) from new-recording-widgets into main

Reviewed-on: #14
Reviewed-by: madrigal <madrigal@qoherent.ai>

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