89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
|
"""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"
|