123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
from utils.data.recording import Recording
|
|
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}", overwrite=True
|
|
)
|
|
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", overwrite=True)
|
|
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, seed=num)
|
|
recording.to_npy(filename=f"wb{i + 1}", overwrite=True)
|
|
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}", overwrite=True)
|
|
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}", overwrite=True)
|
|
print(f"birdie{i + 1} file saved.")
|
|
|
|
def generate_zeros(self, length: int = 8192):
|
|
data = np.zeros(shape=length, dtype=np.complex64)
|
|
recording = Recording(data=data)
|
|
recording.to_npy(filename="zero", overwrite=True)
|
|
|
|
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()
|