ria-toolkit-oss/tests/transforms/test_iq_augmentations.py

import numpy as np
import pytest

from ria_toolkit_oss.data import Recording
from ria_toolkit_oss.transforms import iq_augmentations

TEST_DATA1 = [[1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j]]
TEST_DATA2 = [[1 + 42j, 54 - 34j, -234 - 7j]]
TEST_DATA3 = [[1 + 1j, 10 + 3j, -1 + 4j, 7 + 3j]]
TEST_DATA4 = [[1 + 2j, 3 + 4j, 5 + 5j, 4 + 4j, 2 - 6j]]
TEST_METADATA = {"apple": 1, "watermelon": 2, "mango": 3}


def test_rms_power_1():
    # Testing rms power calculation
    signal = [1, 1, 1, 1, 1, 1]
    signal_rms_power = np.sqrt(np.mean(np.abs(signal) ** 2))
    assert np.allclose(signal_rms_power, [1])


def test_rms_power_2():
    # Testing rms power calculation
    signal = [1, -1, 1, -1, 1, -1]
    signal_rms_power = np.sqrt(np.mean(np.abs(signal) ** 2))
    assert np.allclose(signal_rms_power, [1])


def test_awgn():
    # Testing awgn calculations
    rms_power = 0.5
    channels, length = 10, 10000
    variance = rms_power**2
    magnitude = np.random.normal(loc=0, scale=np.sqrt(variance), size=(channels, length))
    phase = np.random.uniform(low=0, high=2 * np.pi, size=(channels, length))
    noise = magnitude * np.exp(1j * phase)
    noise_power = np.sqrt(np.mean(np.abs(noise) ** 2))
    assert (rms_power - noise_power) < 0.01
    assert noise.shape == (channels, length)


def test_generate_awgn():
    # Testing generate_awgn() with array_like input
    length = 1000
    sample_rate = 1000000
    frequency = 1000
    amplitude = 1
    baseband_phase = 0
    rf_phase = 0
    dc_offset = 0
    snr = 0

    total_time = length / sample_rate
    t = np.linspace(0, total_time, length, endpoint=False)
    sine_wave = amplitude * np.sin(2 * np.pi * frequency * t + baseband_phase) + dc_offset
    complex_sine_wave = sine_wave * np.exp(1j * rf_phase)
    signal = complex_sine_wave.reshape(1, complex_sine_wave.shape[0])
    noisy_sine_wave = iq_augmentations.generate_awgn(signal, snr)
    assert np.sqrt(np.mean(np.abs(noisy_sine_wave) ** 2)) / np.sqrt(np.mean(np.abs(signal) ** 2)) - 2 < 0.01


def test_time_reversal():
    # Testing time_reversal() with array_like input
    transformed_data = iq_augmentations.time_reversal(TEST_DATA1)
    assert np.array_equal(transformed_data, np.asarray([[4 + 4j, 3 + 3j, 2 + 2j, 1 + 1j]]))


def test_time_reversal_rec():
    # Testing time_reversal() with Recording input
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.time_reversal(rec)
    assert np.array_equal(transformed_rec.data, np.asarray([[4 + 4j, 3 + 3j, 2 + 2j, 1 + 1j]]))
    assert rec.metadata == transformed_rec.metadata


def test_spectral_inversion():
    # Testing spectral_inversion() with array_like input
    transformed_data = iq_augmentations.spectral_inversion(TEST_DATA1)
    assert np.array_equal(transformed_data, np.asarray([[1 - 1j, 2 - 2j, 3 - 3j, 4 - 4j]]))


def test_spectral_inversion_rec():
    # Testing spectral_inversion() with Recording input
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.spectral_inversion(rec)
    assert np.array_equal(transformed_rec.data, np.asarray([[1 - 1j, 2 - 2j, 3 - 3j, 4 - 4j]]))
    assert rec.metadata == transformed_rec.metadata


def test_channel_swap():
    # Testing channel_swap() with array_like input
    transformed_data = iq_augmentations.channel_swap(TEST_DATA2)
    assert np.array_equal(transformed_data, np.asarray([[42 + 1j, -34 + 54j, -7 - 234j]]))


def test_channel_swap_rec():
    # Testing channel_swap() with Recording input
    rec = Recording(data=TEST_DATA2, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.channel_swap(rec)
    assert np.array_equal(transformed_rec.data, np.asarray([[42 + 1j, -34 + 54j, -7 - 234j]]))
    assert rec.metadata == transformed_rec.metadata


def test_amplitude_reversal():
    # Testing amplitude_reversal() with array_like input
    transformed_data = iq_augmentations.amplitude_reversal(TEST_DATA2)
    assert np.array_equal(transformed_data, np.asarray([[-1 - 42j, -54 + 34j, 234 + 7j]]))


def test_amplitude_reversal_rec():
    # Testing amplitude_reversal() with array_like input
    rec = Recording(data=TEST_DATA2, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.amplitude_reversal(rec)
    assert np.array_equal(transformed_rec.data, np.asarray([[-1 - 42j, -54 + 34j, 234 + 7j]]))
    assert rec.metadata == transformed_rec.metadata


def test_drop_samples_back_fill():
    # Testing drop_samples() when fill_type='back-fill'
    np.random.seed(0)
    transformed_data = iq_augmentations.drop_samples(TEST_DATA1, max_section_size=2, fill_type="back-fill")
    assert np.array_equal(transformed_data, np.asarray([[1 + 1j, 1 + 1j, 3 + 3j, 3 + 3j]]))


def test_drop_samples_front_fill():
    # Testing drop_samples() when fill_type='front-fill'
    np.random.seed(0)
    transformed_data = iq_augmentations.drop_samples(TEST_DATA1, max_section_size=2, fill_type="front-fill")
    assert np.array_equal(transformed_data, np.asarray([[3 + 3j, 3 + 3j, 3 + 3j, 4 + 4j]]))


def test_drop_samples_mean():
    # Testing drop_samples() when fill_type='mean'
    np.random.seed(0)
    transformed_data = iq_augmentations.drop_samples(TEST_DATA1, max_section_size=2, fill_type="mean")
    assert np.array_equal(transformed_data, np.asarray([[2.5 + 2.5j, 2.5 + 2.5j, 3 + 3j, 2.5 + 2.5j]]))


def test_drop_samples_zeros():
    # Testing drop_samples() when fill_type='zeros'
    np.random.seed(0)
    transformed_data = iq_augmentations.drop_samples(TEST_DATA1, max_section_size=2, fill_type="zeros")
    assert np.array_equal(transformed_data, np.asarray([[0, 0, 3 + 3j, 0]]))


def test_quantize_tape_floor():
    # Testing quantize_tape() with array_like input when rounding = 'floor'
    transformed_data = iq_augmentations.quantize_tape(TEST_DATA3, bin_number=2, rounding_type="floor")
    assert np.array_equal(transformed_data, np.asarray([[-1 - 1j, 4.5 - 1j, -1 - 1j, 4.5 - 1j]]))


def test_quantize_tape_rec_ceiling():
    # Testing quantize_tape() with Recording input when rounding = 'ceiling'
    rec = Recording(data=TEST_DATA3, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.quantize_tape(rec, bin_number=2, rounding_type="ceiling")
    assert np.array_equal(transformed_rec.data, np.asarray([[4.5 + 4.5j, 10 + 4.5j, 4.5 + 4.5j, 10 + 4.5j]]))
    assert rec.metadata == transformed_rec.metadata


def test_quantize_parts_ceiling():
    # Testing quantize_parts() with array_like input when rounding = 'ceiling'
    np.random.seed(0)
    transformed_data = iq_augmentations.quantize_parts(
        TEST_DATA3, max_section_size=2, bin_number=2, rounding_type="ceiling"
    )
    assert np.array_equal(transformed_data, np.asarray([[4.5 + 4.5j, 10 + 4.5j, -1 + 4j, 10 + 4.5j]]))


def test_quantize_parts_rec_floor():
    # Testing quantize_parts() with Recording input when rounding = 'floor'
    np.random.seed(0)
    rec = Recording(data=TEST_DATA3, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.quantize_parts(rec, max_section_size=2, bin_number=2, rounding_type="floor")
    assert np.array_equal(transformed_rec.data, np.asarray([[-1 - 1j, 4.5 - 1j, -1 + 4j, 4.5 - 1j]]))
    assert rec.metadata == transformed_rec.metadata


def test_magnitude_rescale():
    # Testing magnitude_rescale with array_like input
    np.random.seed(0)
    transformed_data = iq_augmentations.magnitude_rescale(TEST_DATA1, starting_bounds=(2, 3), max_magnitude=2)
    assert np.allclose(
        transformed_data, np.asarray([[1 + 1j, 2 + 2j, 3.55706771 + 3.55706771j, 4.74275695 + 4.74275695j]])
    )


def test_magnitude_rescale_rec():
    # Testing magnitude_rescale with Recording input
    np.random.seed(0)
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.magnitude_rescale(rec, starting_bounds=(2, 3), max_magnitude=2)
    assert np.allclose(
        transformed_rec.data, np.asarray([[1 + 1j, 2 + 2j, 3.55706771 + 3.55706771j, 4.74275695 + 4.74275695j]])
    )
    assert rec.metadata == transformed_rec.metadata


def test_cut_out_ones():
    # Testing cut_out() with array_like input when fill_type = 'ones'
    np.random.seed(0)
    transformed_data = iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="ones")
    assert np.array_equal(transformed_data, np.asarray([[1 + 1j, 1 + 1j, 3 + 3j, 1 + 1j]]))


def test_cut_out_avg_snr_1():
    # Testing cut_out() with array_like input when fill_type = 'avg-snr'
    np.random.seed(0)
    transformed_data = iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="avg-snr")
    assert np.allclose(
        transformed_data,
        np.asarray([[1.04504475 - 3.19650874j, 2.18835276 + 1.87922077j, 3 + 3j, 3.38706877 - 0.53958902j]]),
    )


def test_patch_shuffle():
    # Testing patch_shuffle() with array_like input
    np.random.seed(0)
    transformed_data = iq_augmentations.patch_shuffle(TEST_DATA4, max_patch_size=3)
    assert np.array_equal(transformed_data, np.asarray([[3 + 2j, 1 + 4j, 5 + 5j, 2 - 6j, 4 + 4j]]))


def test_patch_shuffle_rec():
    # Testing patch_shuffle() with Recording input
    np.random.seed(0)
    rec = Recording(data=TEST_DATA4, metadata=TEST_METADATA)
    transformed_rec = iq_augmentations.patch_shuffle(rec, max_patch_size=3)
    assert np.array_equal(transformed_rec.data, np.asarray([[3 + 2j, 1 + 4j, 5 + 5j, 2 - 6j, 4 + 4j]]))
    assert rec.metadata == transformed_rec.metadata


# ---------------------------------------------------------------------------
# Additional coverage: error paths and missing Recording variants
# ---------------------------------------------------------------------------


# --- generate_awgn ---


def test_generate_awgn_recording_input():
    # generate_awgn() with a Recording should return a Recording with same metadata.
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    result = iq_augmentations.generate_awgn(rec, snr=10)
    assert isinstance(result, Recording)
    assert result.metadata == rec.metadata


def test_generate_awgn_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.generate_awgn(np.array([[1.0, 2.0, 3.0]]))


def test_generate_awgn_invalid_1d_raises():
    with pytest.raises(ValueError):
        iq_augmentations.generate_awgn(np.array([1 + 1j, 2 + 2j]))


# --- time_reversal ---


def test_time_reversal_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.time_reversal(np.array([[1.0, 2.0, 3.0]]))


def test_time_reversal_multi_channel_raises():
    with pytest.raises(NotImplementedError):
        iq_augmentations.time_reversal([[1 + 1j, 2 + 2j], [3 + 3j, 4 + 4j]])


# --- spectral_inversion ---


def test_spectral_inversion_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.spectral_inversion(np.array([[1.0, 2.0, 3.0]]))


def test_spectral_inversion_multi_channel_raises():
    with pytest.raises(NotImplementedError):
        iq_augmentations.spectral_inversion([[1 + 1j, 2 + 2j], [3 + 3j, 4 + 4j]])


# --- channel_swap ---


def test_channel_swap_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.channel_swap(np.array([[1.0, 2.0, 3.0]]))


# --- amplitude_reversal ---


def test_amplitude_reversal_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.amplitude_reversal(np.array([[1.0, 2.0, 3.0]]))


# --- drop_samples ---


def test_drop_samples_rec_input():
    # drop_samples() with a Recording should return a Recording.
    np.random.seed(0)
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    result = iq_augmentations.drop_samples(rec, max_section_size=2, fill_type="zeros")
    assert isinstance(result, Recording)
    assert result.metadata == rec.metadata


def test_drop_samples_invalid_max_section_size_zero():
    # max_section_size < 1 must raise ValueError.
    with pytest.raises(ValueError):
        iq_augmentations.drop_samples(TEST_DATA1, max_section_size=0)


def test_drop_samples_invalid_max_section_size_too_large():
    # max_section_size >= n must raise ValueError.
    with pytest.raises(ValueError):
        iq_augmentations.drop_samples(TEST_DATA1, max_section_size=len(TEST_DATA1[0]))


def test_drop_samples_invalid_fill_type_raises():
    with pytest.raises(ValueError):
        iq_augmentations.drop_samples(TEST_DATA1, max_section_size=2, fill_type="unknown")


def test_drop_samples_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.drop_samples(np.array([[1.0, 2.0, 3.0]]))


# --- quantize_tape ---


def test_quantize_tape_invalid_rounding_type_raises():
    # An unrecognised rounding_type must raise UserWarning.
    with pytest.warns(UserWarning):
        iq_augmentations.quantize_tape(TEST_DATA1, rounding_type="round")


def test_quantize_tape_invalid_real_raises():
    with pytest.raises(ValueError):
        iq_augmentations.quantize_tape(np.array([[1.0, 2.0, 3.0]]))


# --- quantize_parts ---


def test_quantize_parts_invalid_rounding_type_raises():
    with pytest.warns(UserWarning):
        iq_augmentations.quantize_parts(TEST_DATA1, rounding_type="round")


# --- magnitude_rescale ---


def test_magnitude_rescale_invalid_bounds_negative_raises():
    with pytest.raises(ValueError):
        iq_augmentations.magnitude_rescale(TEST_DATA1, starting_bounds=(-1, 2))


def test_magnitude_rescale_invalid_bounds_too_large_raises():
    n = len(TEST_DATA1[0])
    with pytest.raises(ValueError):
        iq_augmentations.magnitude_rescale(TEST_DATA1, starting_bounds=(0, n))


# --- cut_out ---


def test_cut_out_zeros():
    # cut_out() with fill_type='zeros' must fill the section with 0+0j.
    np.random.seed(0)
    result = iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="zeros")
    assert result.dtype == np.asarray(TEST_DATA1).dtype or np.iscomplexobj(result)


def test_cut_out_low_snr():
    # cut_out() with 'low-snr' should change the signal.
    np.random.seed(0)
    result = iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="low-snr")
    assert result.shape == np.asarray(TEST_DATA1).shape


def test_cut_out_high_snr():
    # cut_out() with 'high-snr' should return data with same shape.
    np.random.seed(0)
    result = iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="high-snr")
    assert result.shape == np.asarray(TEST_DATA1).shape


def test_cut_out_rec_input():
    # cut_out() with Recording should return Recording with preserved metadata.
    np.random.seed(0)
    rec = Recording(data=TEST_DATA1, metadata=TEST_METADATA)
    result = iq_augmentations.cut_out(rec, max_section_size=2, fill_type="zeros")
    assert isinstance(result, Recording)
    assert result.metadata == rec.metadata


def test_cut_out_invalid_fill_type_raises():
    with pytest.warns(UserWarning):
        iq_augmentations.cut_out(TEST_DATA1, max_section_size=2, fill_type="bad")


def test_cut_out_invalid_max_section_size_raises():
    with pytest.raises(ValueError):
        iq_augmentations.cut_out(TEST_DATA1, max_section_size=0)


# --- patch_shuffle ---


def test_patch_shuffle_max_patch_size_leq_1_raises():
    with pytest.raises(ValueError):
        iq_augmentations.patch_shuffle(TEST_DATA1, max_patch_size=1)


def test_patch_shuffle_max_patch_size_too_large_raises():
    n = len(TEST_DATA1[0])
    with pytest.raises(ValueError):
        iq_augmentations.patch_shuffle(TEST_DATA1, max_patch_size=n + 1)