记一次fft模块重构(一):基础仿真篇

@dataclass
class ToneParams:
#单路信号参数
    freq_hz: float
    amplitude : float #归一化幅度
    waveform : str = "sine" #波形, sine | traiangle | square,默认sine
    phase: float = 0 #相位
@dataclass
class ChannelParams:
#信道参数
    snr_db: float | None = None
    seed: int = 42

@dataclass
class ADCParams:
    adc_scale: float #adc量程
    quantized_acc: int #量化精度
    fs_hz: float = 1e6 #采样率
    n_samples: int = 1024 #采样点数
	sampled_time: int = 1000 #采样时间,单位ms

def generate_test_signal(
    tones: list[ToneParams],
    channel: ChannelParams,
    adc: ADCParams 
) -> tuple[np.ndarray, dict]:
#返回(信号数组,元信息字典)

_WAVEFORM_GENERATORS = {
    "sine":    lambda t, f, A: A * np.sin(2 * np.pi * f * t),
    "triangle": lambda t, f, A: A * (2/np.pi) * np.arcsin(np.sin(2*np.pi*f*t)),
    "square":  lambda t, f, A: A * np.sign(np.sin(2*np.pi*f*t)),
}

"dc": lambda b: b

def generate_test_signal(
    tones: list[ToneParams],
    channel: ChannelParams,
    adc: ADCParams 
) -> tuple[np.ndarray, dict]:
#返回(信号数组,元信息字典)
    Ts = 1 / adc.fs_hz
    t_seq = np.arange(0,adc.sampled_time,Ts)
    sumed_signal = np.zeros(len(t_seq))
    for i in tones: #返回信号加和
        sumed_signal += _WAVEFORM_GENERATORS[i.waveform](
            t = t_seq,
            f = i.freq_hz,
            A = i.amplitude,
            phi = i.phase
        )
    #添加信道噪声
    rng = np.random.default_rng(channel.seed)
    P_signal = np.mean(sumed_signal ** 2)
    if channel.snr_db:
        P_noise = P_signal / (10 ** (channel.snr_db / 10))
        noise = np.sqrt(P_noise) * rng.standard_normal(len(sumed_signal))
        sumed_signal += noise
    #ADC采样部分
    #数据切片
    if len(sumed_signal) < adc.n_samples:
        for i in range(adc.n_samples - len(sumed_signal)):
            np.append(sumed_signal,0) #如果信号数小于采样点数,默认填充0
    sumed_signal = sumed_signal[0:adc.n_samples]
    #处理超过采样范围的幅值,超过采样范围的复制设定为峰值(模拟削顶削底效应)
    for i in range(len(sumed_signal)):
        if sumed_signal[i] >= 1:
            sumed_signal[i] = 1
        elif sumed_signal[i] <= -1:
            sumed_signal[i] = -1
    
    #ADC量化采样
    quantized_signal = np.zeros(len(t_seq))
    for i in range(len(sumed_signal)):
        quantized_signal[i] = int(sumed_signal[i] * ((adc.quantized_acc - 1) // 2))

    #信号元数据字典
    signal_data = {
        "tone": tones,
        "channel": channel,
        "adc": adc
    }
    return (quantized_signal,signal_data)

import numpy as np
#1.标量误差统计
def rmse( #均方根误差
    estimates: np.ndarray,
    truth: float | np.ndarray 
) -> float:
    return float(np.sqrt(np.mean((estimates - truth) ** 2)))

def mae( #平均绝对误差
    estimates: np.ndarray,
    truth: float | np.ndarray
) -> float:
    return float(np.mean(np.abs(estimates - truth)))

def bias(
    estimates: np.ndarray,
    truth: float | np.ndarray
) -> float:
    return float(np.mean(estimates - truth))

def std_dev(#标准差
    estimates: np.ndarray
) -> float:
    return float(np.std(estimates, ddof=1))

def reltave_error_pct(#相对误差百分比
    estimates: np.ndarray,
    truth: float
) -> np.ndarray:
    return np.abs(estimates - truth) / np.abs(truth) * 100.0

def max_abs_error(#最大绝对误差
    estimates: np.ndarray,
    truth: float
) -> float:
    return np.max(np.abs(estimates - truth))


#2.波形误差估计统计
def signal_rmse(#两路信号的均方根误差
    sig_est: np.ndarray,
    sig_true: np.ndarray
) -> float:
    return np.sqrt(np.mean((sig_est - sig_true) ** 2))

def signal_correlation(#信号的皮尔逊相关系数
    sig_est: np.ndarray,
    sig_true: np.ndarray
) -> float:
    return np.corrcoef(sig_est, sig_true)[0, 1]

def signal_nmse_db(#归一化均方误差(dB)
    sig_est: np.ndarray,
    sig_true: np.ndarray
) -> float:
    mse = np.mean((sig_est - sig_true) ** 2)
    power = np.mean(sig_true ** 2)
    return 10.0 * np.log10(mse / power) if power > 0 else -np.inf

#3.辅助工具函数
def compute_snr(#计算信噪比
    sig_est: np.ndarray,
    sig_clean: np.ndarray
) -> float:
    noise = sig_est - sig_clean
    p_noise = np.mean(noise ** 2)
    p_signal = np.mean(sig_clean ** 2)
    return 10 * np.log10(p_signal / p_noise) if p_noise > 0 else np.inf

def estimate_noise_floor(#计算噪声底
    spectrum: np.ndarray,
    signal_bins: list[int]
) -> float:
    mask = np.ones(len(spectrum), dtype=bool)
    for b in signal_bins:
        mask[max(0, b-3): min(len(spectrum), b+4)] = False
    return float(np.mean(spectrum[mask])) if np.any(mask) else 0.0

_WINDOWS_GENERATORS = { #窗函数生成字典
        "rect": lambda pts: np.ones(pts),
        "hann": lambda pts: scipy.signal.windows.hann(pts, sym=False),
        "hamming": lambda pts: scipy.signal.windows.hamming(pts, sym=False),
        "blackmanharris": lambda pts: scipy.signal.windows.blackmanharris(pts, sym=False),
        "flattop": lambda pts: scipy.signal.windows.flattop(pts, sym=False)
        }

def measure_coherent_gain (window: np.ndarray) -> float: #求相干增益，用于矫正幅值
    return float(np.mean(window))

def measure_scalloping_loss (window: np.ndarray,#求扇贝损失，用于检查数据的误差上限
                             N_fft: int,
                             n_sweep=500#扫频步长，越小越精确
                             ) -> float:
    N = len(window)
    n = np.arange(N)
    k0 = N_fft // 4 #远离DC分量
    s = np.sin(2 * np.pi * k0 / N_fft * n)
    X = np.fft.fft(s * window, n=N_fft)
    peak_ref = np.max(np.abs(X))
    min_peak = float('inf')
    peak = 0
    for offset in np.linspace(0, 1.0, n_sweep):
        s = np.sin(2 * np.pi * (k0 + offset) / N_fft * n)
        X = np.fft.fft(s * window,n=N_fft)
        peak = np.max(np.abs(X))
        if peak <= min_peak:
            min_peak = peak
    scalloping_loss = -20 * np.log10(min_peak / peak_ref)
    return float(scalloping_loss)

def parabolic_interp(
                   spectrum: np.ndarray, #幅度谱,必须是线性幅度谱，不能是分贝谱
                    peak_bin: int, #峰值bin
                    fs_hz: float,
                    n_pts: int,
                    coherent_gain: float
                    ) -> EstimateData:
   if peak_bin < 1 or peak_bin > len(spectrum) - 2:
       print("[ERROR] peak bin is too small or too large")#简单的错误处理
       return EstimateData(freq_estimated=0,amplitude_estimated=0)
   y_minus1 = spectrum[peak_bin - 1]
   y_0 = spectrum[peak_bin]
   y_plus1 = spectrum[peak_bin + 1]
   freq_delta = 0.5 * (y_minus1 - y_plus1) / (y_minus1 - 2 * y_0 + y_plus1)
   a = (y_minus1 + y_plus1 - 2 * y_0)#抛物线方程中的a
   y_peak = y_0 - a * (freq_delta ** 2)
   freq_estimated_bin = (peak_bin + freq_delta) * (fs_hz / n_pts)
   amplitude_estimated = y_peak / (n_pts * coherent_gain)
   output_EstimateData = EstimateData(
           freq_estimated=freq_estimated_bin,
           amplitude_estimated=amplitude_estimated
           )
   
   return output_EstimateData

@dataclass
class EstimateData:
    freq_estimated: float
    amplitude_estimated: float

def get_db_spectrum(spectrum: np.ndarray) -> np.ndarray: #获得窗函数dB谱
    dB_spectrum = 20 * np.log10(np.maximum(np.abs(spectrum) / np.max(spectrum), 1e-300)) #转化为dB谱
    return dB_spectrum

def estimate_noise_floor(
        spectrum: np.ndarray, #信号幅度谱
        k = 1.5, #σ裁剪系数
        ) -> float:
    #中位数估计法
    mu_init = np.median(spectrum) / np.log(2)
    #mask_init = spectrum < mu_init * k #剔除明显信号
    #noice_floor = np.mean(spectrum[mask_init])
    noice_floor = mu_init
    noice_floor_db = 20 * np.log10(np.abs(np.maximum(noice_floor / np.max(spectrum), 1e-300)))
    return float(noice_floor_db)

def find_peaks_above_noise(
        spectrum: np.ndarray,  #过滤频谱峰值,要求线性谱
        margin_dB=6.0 #大于噪声底x dB才会判定为有效峰值
                           ) -> list[int]:
    peaks_index = []
    noice_floor = estimate_noise_floor(spectrum)
    spectrum_db = get_db_spectrum(spectrum)
    for i in range(len(spectrum_db)-1):
        if (
        spectrum_db[i-1] <= spectrum_db[i] #找局部极大值
        and spectrum_db[i+1] <= spectrum_db[i]
        and spectrum_db[i] >= noice_floor + margin_dB #防止筛选噪声
        ):
            peaks_index.append(i)

    return peaks_index

def peak_dominant_filtering( #支配滤波,用于排除假峰(假局部极大值)
        peaks_index: list[int],
        spectrum_db: np.ndarray,
        dominance_radius = 10, #寻找范围
        dominance_ratio = 10 #衰减大于10dB则不算做新信号
        ) -> list[int]:
    peaks_merged = []
    current_peak_index = 0 #当前峰值索引
    peak_mask = [False for _ in range(len(peaks_index))]
    peak_mask[current_peak_index] = True
    for i in range(1,len(peaks_index)):
        if peaks_index[i] - peaks_index[current_peak_index] <= dominance_radius:
            #检测到主瓣半径内有比自己更大的值,以该值为中心进行计算
            test1 = spectrum_db[peaks_index[i]]
            test2 = spectrum_db[peaks_index[current_peak_index]]
            
            if spectrum_db[peaks_index[i]] >= spectrum_db[peaks_index[current_peak_index]]:
                peak_mask[i] = True
                #如果后面有一个幅值差距很大的峰,则记当前谱线为假峰
                if spectrum_db[peaks_index[i]] - spectrum_db[peaks_index[current_peak_index]] >= dominance_ratio:
                    peak_mask[current_peak_index] = False
                current_peak_index = i

        else:
            peak_mask[i] = True
            current_peak_index = i

    for i in range(len(peaks_index)):
        if peak_mask[i]:
            peaks_merged.append(peaks_index[i])

    return peaks_merged