import numpy as np from ria_toolkit_oss.datatypes import Annotation, Recording from ria_toolkit_oss.io.recording import ( from_npy, from_sigmf, load_rec, to_npy, to_sigmf, ) complex_data_1 = np.array([0.5 + 0.5j, 0.1 + 0.1j, 0.3 + 0.3j, 0.4 + 0.4j, 0.5 + 0.5j], dtype=np.complex64) real_data_1 = np.array([[0.5, 0.1, 0.3, 0.4, 0.5], [0.5, 0.1, 0.3, 0.4, 0.5]]) sample_metadata = {"source": "test", "timestamp": 1723472227.698788} nd_complex_data_1 = np.array( [ [0.5 + 0.5j, 0.1 + 0.1j, 0.3 + 0.3j, 0.4 + 0.4j, 0.5 + 0.5j], [0.5 + 0.5j, 0.1 + 0.1j, 0.3 + 0.3j, 0.4 + 0.4j, 0.5 + 0.5j], ] ) nd_real_data_1 = np.array([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]) complex_data_out_of_range_1 = np.array([1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j]) def test_npy_save_1(tmp_path): # Create test recording recording1 = Recording(data=complex_data_1, metadata=sample_metadata) # Save to tmp_path filename = tmp_path / "test" to_npy(filename=filename.name, path=tmp_path, recording=recording1) # Reload recording2 = from_npy(filename) # Verify assert np.array_equal(recording1.data, recording2.data) assert recording1.metadata == recording2.metadata def test_npy_save_2(tmp_path): # Create test recording recording1 = Recording(data=nd_complex_data_1, metadata=sample_metadata) # Save to tmp_path filename = tmp_path / "test" to_npy(filename=filename.name, path=tmp_path, recording=recording1) # Reload recording2 = from_npy(filename) # Verify assert np.array_equal(recording1.data, recording2.data) assert recording1.metadata is not None assert recording1.metadata == recording2.metadata # Check that metadata is loaded properly as a dict assert recording1.metadata.get("source") == recording2.metadata.get("source") def test_npy_save_3(tmp_path): # Create test recording without metadata recording1 = Recording(data=nd_complex_data_1) # Save to tmp_path filename = tmp_path / "test" to_npy(filename=filename.name, path=tmp_path, recording=recording1) # Reload recording2 = from_npy(filename) # Verify assert np.array_equal(recording1.data, recording2.data) assert recording1.metadata == recording2.metadata def test_npy_annotations(tmp_path): # Create annotations annotation1 = Annotation(sample_start=0, sample_count=100, freq_lower_edge=0, freq_upper_edge=100) annotation2 = Annotation(sample_start=1, sample_count=101, freq_lower_edge=1, freq_upper_edge=101) annotations = [annotation1, annotation2] # Create test recording with annotations recording1 = Recording(data=nd_complex_data_1, metadata=sample_metadata, annotations=annotations) # Save to tmp_path filename = tmp_path / "test" to_npy(filename=filename.name, path=tmp_path, recording=recording1) # Reload recording2 = from_npy(filename) # Verify annotations assert recording1.annotations == recording2.annotations def test_load_recording_npy(tmp_path): # test to_npy and load_recording methods with npy annotation1 = Annotation(sample_start=0, sample_count=1, freq_lower_edge=0, freq_upper_edge=1) annotation2 = Annotation(sample_start=1, sample_count=2, freq_lower_edge=1, freq_upper_edge=2) annotations = [annotation1, annotation2] recording1 = Recording(data=complex_data_1, metadata=sample_metadata, annotations=annotations) # Save to tmp_path filename = tmp_path / "test.npy" recording1.to_npy(path=tmp_path, filename=filename.name) # Load from tmp_path recording2 = load_rec(filename) assert recording1.annotations == recording2.annotations # Check that original metadata was preserved assert all( key in recording2.metadata and recording2.metadata[key] == value for key, value in recording1.metadata.items() ) assert np.array_equal(recording1.data, recording2.data) def test_sigmf_1(tmp_path): # Create annotations annotation1 = Annotation(sample_start=0, sample_count=1, freq_lower_edge=0, freq_upper_edge=1) annotation2 = Annotation(sample_start=1, sample_count=2, freq_lower_edge=1, freq_upper_edge=2) annotations = [annotation1, annotation2] # Create test recording with annotations recording1 = Recording(data=complex_data_1, metadata=sample_metadata, annotations=annotations) # Save to tmp_path in SigMF format filename = tmp_path / "test" to_sigmf(recording=recording1, path=tmp_path, filename=filename.name) # Reload recording2 = from_sigmf(filename) # Verify annotations assert recording1.annotations == recording2.annotations # Verify metadata (original keys preserved) assert all( key in recording2.metadata and recording2.metadata[key] == value for key, value in recording1.metadata.items() ) # Verify data assert np.array_equal(recording1.data, recording2.data) def test_sigmf_2(tmp_path): # checks that recording can be saved to sigmf and then retrieved without data loss annotation1 = Annotation(sample_start=0, sample_count=1, freq_lower_edge=0, freq_upper_edge=1) annotation2 = Annotation(sample_start=1, sample_count=2, freq_lower_edge=1, freq_upper_edge=2) annotations = [annotation1, annotation2] recording1 = Recording(data=complex_data_1, metadata=sample_metadata, annotations=annotations) # Save to tmp_path using the base name filename = tmp_path / "test" to_sigmf(recording=recording1, path=tmp_path, filename=filename.name) # Load from tmp_path; from_sigmf expects the base name recording2 = from_sigmf(filename) assert recording1.annotations == recording2.annotations # checks that the original metadata was preserved (although some sigmf specific metadata may have been added) assert all( key in recording2.metadata and recording2.metadata[key] == value for key, value in recording1.metadata.items() ) assert np.array_equal(recording1.data, recording2.data)