227 lines
9.4 KiB
Python
227 lines
9.4 KiB
Python
import numpy as np
|
|
|
|
from ria_toolkit_oss.datatypes 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.26516288 - 0.36655702j, -2.44693984 + 1.27294267j, 3 + 3j, 4.1583403 - 0.96625365j]]),
|
|
)
|
|
|
|
|
|
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
|