ria-toolkit-oss/tests/io/test_recording_io.py

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)