recording-generation/recording_generation.py

import os
import random

import numpy as np
from utils.io.recording import from_npy

from signal_generation import (create_birdie_recording, create_ctnb_recording,
                               create_lfm_recording, create_modulated_signal,
                               create_noise_recording)


class RecordingGenerator:
    def __init__(self, sample_rate: int | float = 10e6):
        self.sample_rate = int(sample_rate)

    def generate_collision(
        self,
        mod_choices: list = ["qam16", "qam64", "qam256"],
        roll_offs: list = [0.15, 0.35],
        sps_choices: list = [4, 5, 6],
        length: int = 8192,
    ):
        for modulation in mod_choices:
            for roll_off in roll_offs:
                roll_off_str = str(roll_off)[2:]
                sps = random.choice(sps_choices)
                length = 8192
                recording = create_modulated_signal(
                    modulation=modulation, sps=sps, beta=roll_off, length=length
                )
                recording.to_npy(filename=f"{modulation}_{roll_off_str}")
                print(f"{modulation}_{roll_off_str} file saved.")

    def generate_lfm(
        self, period_choices: list = [0.25, 0.3, 0.35], width_choices: list = [10]
    ):
        for chirp_type in ["up", "down", "up_down"]:
            chirp_period = random.choice(period_choices)
            width_factor = random.choice(width_choices)
            width = self.sample_rate // width_factor

            recording = create_lfm_recording(
                sample_rate=self.sample_rate,
                width=width,
                chirp_type=chirp_type,
                chirp_period=chirp_period,
                length=int(self.sample_rate * chirp_period),
            )

            print(f"LFM chirp length: {int(self.sample_rate * chirp_period)}")
            recording.to_npy(filename=f"{chirp_type}_chirp")
            print(f"{chirp_type}_chirp file saved.")

    def generate_wb(self, num: int = 2, length: int = 8192):
        for i in range(num):
            recording = create_noise_recording(
                length=length,
                rms_power=0.2,
            )
            recording.to_npy(filename=f"wb{i + 1}")
            print(f"wb{i + 1} file saved.")

    def generate_ctnb(self, num: int = 2, length: int = 8192):
        for i in range(num):
            recording = create_ctnb_recording(length=length)
            recording.to_npy(filename=f"ctnb{i + 1}")
            print(f"ctnb{i + 1} file saved.")

    def generate_birdie(self, num: int = 2, length: int = 8192, wave_num: int = 5):
        for i in range(num):
            recording = create_birdie_recording(
                sample_rate=int(self.sample_rate),
                length=length,
                wave_number=int(wave_num + i),
            )
            recording.to_npy(filename=f"birdie{i + 1}")
            print(f"birdie{i + 1} file saved.")

    def convert_to_dat(
        self,
        source_directory: str = "recordings",
        save_directory: str = "dat_recordings",
    ):
        os.makedirs(save_directory, exist_ok=True)
        for root, _, files in os.walk(source_directory):
            for name in files:
                filename = os.path.join(root, name)
                savename = save_directory + "/" + name[:-4] + ".dat"

                recording = from_npy(file=filename)
                data = recording.data[0]

                # Convert complex128 -> float64 -> int16 after scaling
                real = np.real(data)
                imag = np.imag(data)

                # Scale down the float values to fit in int16 range [-32768, 32767]
                # Adjust the scaling factor depending on your signal's dynamic range
                scale_factor = 32767 / np.max(np.abs(np.concatenate((real, imag))))
                real_scaled = (real * scale_factor).astype(np.int16)
                imag_scaled = (imag * scale_factor).astype(np.int16)

                # Interleave real and imag
                interleaved = np.empty((real_scaled.size * 2,), dtype=np.int16)
                interleaved[0::2] = real_scaled
                interleaved[1::2] = imag_scaled

                interleaved.tofile(savename)
                print(f"Saved {savename}")


if __name__ == "__main__":
    generator = RecordingGenerator()