Merge pull request 'Render sample_spectrogram for split-I/Q RadioDataset samples' (#36) from fix-spectrogram-split-iq into main

Reviewed-on: #36
Reviewed-by: gillian <gillian@qoherent.ai>

src/ria_toolkit_oss/viz/radio_dataset.py

@@ -59,6 +59,30 @@ def create_styled_error_figure(title: str, message: str, suggestion: str = None)
     return fig
 
 
+def _to_complex_1d(sample) -> "np.ndarray | None":
+    """Normalize a dataset sample to a 1-D complex signal.
+
+    Handles the layouts RadioDataset samples appear in:
+      * already-complex (any shape) -> flattened
+      * split I/Q with a length-2 axis first  (2, ...)  -> row0 + 1j*row1   (I, Q)
+      * split I/Q with a length-2 axis last   (..., 2)   -> [...,0] + 1j*[...,1]
+      * real 1-D                              -> real signal (imag = 0)
+    Returns None if it can't produce a usable array.
+    """
+    if sample is None:
+        return None
+    arr = np.asarray(sample)
+    if arr.size == 0:
+        return None
+    if np.iscomplexobj(arr):
+        return arr.reshape(-1)
+    if arr.ndim >= 2 and arr.shape[0] == 2:  # (2, ...)  I/Q rows
+        return arr[0].reshape(-1) + 1j * arr[1].reshape(-1)
+    if arr.ndim >= 2 and arr.shape[-1] == 2:  # (..., 2)  I/Q columns
+        return arr[..., 0].reshape(-1) + 1j * arr[..., 1].reshape(-1)
+    return arr.reshape(-1).astype(complex)  # real 1-D signal
+
+
 def _check_dataset_compatibility(dataset, plot_type: str) -> tuple[bool, str]:
     """Check if dataset is compatible with a specific plot type.
     Returns (is_compatible, error_message)
@@ -85,14 +109,16 @@ def _check_dataset_compatibility(dataset, plot_type: str) -> tuple[bool, str]:
             if len(metadata) < 1:
                 return False, "No samples available for spectrogram"
 
-            # Check if we can access sample data (basic test)
+            # Check if we can access sample data (basic test). Normalize to a 1-D
+            # complex signal first so split-I/Q samples (shape (2, T)) report their
+            # true length T, not the size of the I/Q axis.
             try:
-                sample_data = dataset[0] if hasattr(dataset, "__getitem__") else None
+                sig = _to_complex_1d(dataset[0]) if hasattr(dataset, "__getitem__") else None
-                if sample_data is None or len(sample_data) < 32:
-                    return False, "Insufficient sample data for spectrogram (need at least 32 points)"
             except Exception:
-                # If we can't access data, we'll rely on synthetic data generation
+                # If we can't access data, we'll rely on synthetic data generation.
-                pass
+                sig = None
+            if sig is not None and sig.size < 32:
+                return False, "Insufficient sample data for spectrogram (need at least 32 points)"
 
         return True, ""
 
@@ -337,7 +363,7 @@ def _calculate_spectrogram_params(n_samples: int) -> tuple[int, int, int, int]:
 
 def _compute_spectrogram(sample_data, nperseg: int, hop_length: int, n_frames: int, freq_bins: int):
     """Compute spectrogram using FFT."""
-    n_samples = len(sample_data)
+    n_samples = np.asarray(sample_data).reshape(-1).shape[0]
     Sxx = np.zeros((freq_bins, n_frames))
 
     for i in range(n_frames):
@@ -409,13 +435,17 @@ def sample_spectrogram_plot(dataset, class_key: str = "modulation", sample_idx:
             sample_idx = random.randint(0, len(metadata) - 1)
         sample_metadata = metadata.iloc[sample_idx]
 
-        # Get sample data and ensure it's complex
+        # Normalize the sample to a 1-D complex signal (combines split I/Q, etc.)
-        sample_data = _get_sample_data(dataset, sample_idx)
+        sample_data = _to_complex_1d(_get_sample_data(dataset, sample_idx))
-        if not np.iscomplexobj(sample_data):
+        if sample_data is None or sample_data.size < 32:
-            sample_data = sample_data.astype(complex)
+            return create_styled_error_figure(
+                "Spectrogram Not Available",
+                "This sample doesn't have enough signal data for a spectrogram.",
+                "Spectrograms need at least 32 complex samples.",
+            )
 
         # Calculate spectrogram parameters and compute spectrogram
-        n_samples = len(sample_data)
+        n_samples = sample_data.size
         nperseg, hop_length, n_frames, freq_bins = _calculate_spectrogram_params(n_samples)
         Sxx = _compute_spectrogram(sample_data, nperseg, hop_length, n_frames, freq_bins)
 

tests/viz/test_radio_dataset.py

@@ -0,0 +1,88 @@
+"""Tests for spectrogram visualization across RadioDataset sample layouts.
+
+Regression: split-I/Q samples with shape ``(2, T)`` previously reported a length
+of 2 (the I/Q axis) instead of ``T``, so every such dataset was rejected with
+"doesn't have sufficient signal data for spectrogram visualization".
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from ria_toolkit_oss.viz.radio_dataset import _to_complex_1d, sample_spectrogram_plot
+
+
+class _FakeDataset:
+    """Minimal RadioDataset stand-in: a one-row metadata frame + a fixed sample."""
+
+    def __init__(self, sample):
+        self._sample = np.asarray(sample)
+        self.metadata = pd.DataFrame({"modulation": ["bpsk"]})
+
+    def __getitem__(self, idx):
+        return self._sample
+
+
+def _has_spectrogram(fig):
+    """True when fig is a real spectrogram (a Heatmap trace), not an error card."""
+    return any(getattr(tr, "type", None) == "heatmap" for tr in fig.data)
+
+
+# --- _to_complex_1d ---------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "sample, expected_len",
+    [
+        (np.exp(1j * np.linspace(0, 1, 1024)), 1024),  # complex 1-D
+        (np.ones((2, 1024)), 1024),  # split I/Q rows (2, T)
+        (np.ones((1024, 2)), 1024),  # split I/Q cols (T, 2)
+        (np.ones(1024), 1024),  # real 1-D
+        (np.ones((2, 4, 256)), 1024),  # multi-channel I/Q rows
+    ],
+)
+def test_to_complex_1d_normalizes(sample, expected_len):
+    sig = _to_complex_1d(sample)
+    assert sig is not None
+    assert sig.ndim == 1
+    assert sig.size == expected_len
+    assert np.iscomplexobj(sig)
+
+
+def test_to_complex_1d_combines_iq_rows():
+    arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])  # I row, Q row
+    assert np.allclose(_to_complex_1d(arr), np.array([1 + 4j, 2 + 5j, 3 + 6j]))
+
+
+def test_to_complex_1d_combines_iq_cols():
+    arr = np.array([[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]])  # (T, 2)
+    assert np.allclose(_to_complex_1d(arr), np.array([1 + 4j, 2 + 5j, 3 + 6j]))
+
+
+@pytest.mark.parametrize("sample", [None, np.array([])])
+def test_to_complex_1d_returns_none_for_empty(sample):
+    assert _to_complex_1d(sample) is None
+
+
+# --- sample_spectrogram_plot ------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "sample",
+    [
+        np.exp(1j * np.linspace(0, 10, 1024)),  # complex 1-D
+        np.random.randn(2, 1024),  # split I/Q rows  <-- the reported bug
+        np.random.randn(1024, 2),  # split I/Q cols
+        np.random.randn(1024),  # real 1-D
+        np.random.randn(2, 4, 256),  # multi-channel I/Q
+    ],
+)
+def test_sample_spectrogram_renders(sample):
+    fig = sample_spectrogram_plot(_FakeDataset(sample), sample_idx=0)
+    assert _has_spectrogram(fig), "expected a real spectrogram, got an error/unavailable figure"
+
+
+@pytest.mark.parametrize("sample", [np.random.randn(2), np.array([])])
+def test_sample_spectrogram_too_short_returns_error(sample):
+    fig = sample_spectrogram_plot(_FakeDataset(sample), sample_idx=0)
+    assert not _has_spectrogram(fig), "expected the 'Not Available' figure for too-short/empty samples"