visualize.py

# -*- coding: utf-8 -*-
"""
可视化脚本 - 信号波形与特征可视化
包含详细的坐标轴参数标注和图像说明
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from config import config
from dataset import ModulationDataGenerator

plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def plot_signal_waveforms(generator, save_path='./results/waveforms.png'):
    """绘制各种调制类型的时域波形（含详细坐标标注）"""
    snr = 20  # 高信噪比
    
    fig = plt.figure(figsize=(16, 12))
    gs = GridSpec(4, 3, figure=fig, hspace=0.5, wspace=0.35)
    
    for idx, mod_type in enumerate(config.MODULATION_TYPES):
        row = idx // 3
        col = idx % 3
        
        iq = generator.generate_signal(mod_type, snr_db=snr)
        i_component = iq[0]
        q_component = iq[1]
        
        ax = fig.add_subplot(gs[row, col])
        # 时间轴：采样点 -> 秒 -> ms，范围0~1ms（200个采样点 / 200kHz）
        t = np.arange(len(i_component)) / config.SAMPLE_RATE * 1000  # ms
        
        ax.plot(t[:200], i_component[:200], 'b-', alpha=0.8, linewidth=0.8, label='I (同相)')
        ax.plot(t[:200], q_component[:200], 'r-', alpha=0.8, linewidth=0.8, label='Q (正交)')
        
        ax.set_title(f'{mod_type}', fontsize=10, fontweight='bold')
        ax.set_xlabel('时间 (ms)\n[采样率: 200 kHz, 范围: 0~1 ms]', fontsize=7)
        ax.set_ylabel('幅度 (归一化)', fontsize=7)
        ax.legend(loc='upper right', fontsize=6)
        ax.grid(True, alpha=0.3, linestyle='--')
        # 设置x轴范围0~1ms，y轴自适应
        ax.set_xlim([0, 1.0])
    
    plt.suptitle('各类调制信号时域波形 (I/Q分量)\n'
                 '说明：蓝色=I路同相分量，红色=Q路正交分量；SNR=20dB（高信噪比）',
                 fontsize=12, fontweight='bold', y=0.98)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"波形图已保存到: {save_path}")


def plot_constellation_diagrams(generator, save_path='./results/constellation.png'):
    """绘制星座图（含坐标标注）"""
    snr = 20
    
    fig = plt.figure(figsize=(16, 12))
    gs = GridSpec(4, 3, figure=fig, hspace=0.5, wspace=0.35)
    
    # 按类别排列：PSK/QAM, FSK, AM, FM
    psk_qam_types = ['BPSK', 'QPSK', '8PSK', '16QAM', '64QAM']
    fsk_types = ['FSK2', 'FSK4']
    am_types = ['AM', 'AM-DSB', 'AM-SSB']
    fm_types = ['WBFM']
    
    all_types = psk_qam_types + fsk_types + am_types + fm_types
    if len(all_types) < 11:
        all_types = config.MODULATION_TYPES[:11]
    
    for idx, mod_type in enumerate(all_types):
        row = idx // 3
        col = idx % 3
        
        iq = generator.generate_signal(mod_type, snr_db=snr)
        i_component = iq[0]
        q_component = iq[1]
        
        ax = fig.add_subplot(gs[row, col])
        ax.scatter(i_component[::10], q_component[::10], alpha=0.5, s=5, c='blue')
        
        ax.set_title(f'{mod_type}', fontsize=10, fontweight='bold')
        ax.set_xlabel('I (同相分量)\n[归一化幅度]', fontsize=7)
        ax.set_ylabel('Q (正交分量)\n[归一化幅度]', fontsize=7)
        ax.grid(True, alpha=0.3, linestyle='--')
        ax.set_aspect('equal', adjustable='box')
        ax.axhline(y=0, color='k', linewidth=0.5)
        ax.axvline(x=0, color='k', linewidth=0.5)
        ax.set_xlim([-2.5, 2.5])
        ax.set_ylim([-2.5, 2.5])
        ax.set_xticks([-2, -1, 0, 1, 2])
        ax.set_yticks([-2, -1, 0, 1, 2])
    
    plt.suptitle('各类调制信号星座图 (I/Q平面)\n'
                 '说明：横轴I为同相分量，纵轴Q为正交分量；SNR=20dB，采样间隔显示',
                 fontsize=12, fontweight='bold', y=0.98)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"星座图已保存到: {save_path}")


def plot_spectrograms(generator, save_path='./results/spectrograms.png'):
    """绘制频谱图（含频率/时间轴标注）"""
    snr = 20
    
    fig = plt.figure(figsize=(16, 12))
    gs = GridSpec(4, 3, figure=fig, hspace=0.5, wspace=0.35)
    
    # 频率范围：-Fs/2 到 +Fs/2，映射到显示范围
    freq_range = config.SAMPLE_RATE / 1000 / 2  # 转换为kHz = 100 kHz
    
    for idx, mod_type in enumerate(config.MODULATION_TYPES[:11]):
        row = idx // 3
        col = idx % 3
        
        iq = generator.generate_signal(mod_type, snr_db=snr)
        complex_signal = iq[0] + 1j * iq[1]
        
        ax = fig.add_subplot(gs[row, col])
        # specgram: NFFT=128, noverlap=64
        # 时间分辨率: (NFFT - noverlap) / Fs = 64/200000 = 0.32ms per bin
        # 频率分辨率: Fs/NFFT = 200000/128 ≈ 1562.5 Hz
        spec, freqs, t_axis, im = ax.specgram(
            complex_signal, Fs=config.SAMPLE_RATE, NFFT=128, 
            noverlap=64, cmap='viridis'
        )
        
        ax.set_title(f'{mod_type}', fontsize=10, fontweight='bold')
        ax.set_xlabel('时间 (ms)\n[NFFT=128, 重叠=64, 步长≈0.32ms]', fontsize=7)
        ax.set_ylabel('频率 (kHz)\n[分辨率≈1.56 kHz/bin]', fontsize=7)
    
    plt.suptitle('各类调制信号频谱图 (STFT)\n'
                 '说明：横轴为时间，纵轴为频率，颜色代表功率谱密度；SNR=20dB',
                 fontsize=12, fontweight='bold', y=0.98)
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"频谱图已保存到: {save_path}")


def plot_signal_comparison(save_path='./results/comparison.png'):
    """对比同一种调制在不同SNR下的信号（含详细标注）"""
    generator = ModulationDataGenerator()
    mod_type = 'QPSK'
    
    snr_levels = [-5, 10, 30]  # 低、中、高信噪比
    
    fig, axes = plt.subplots(3, 2, figsize=(14, 10))
    
    for row, snr in enumerate(snr_levels):
        iq = generator.generate_signal(mod_type, snr_db=snr)
        
        # 时域波形（I分量）
        t_samples = np.arange(200) / config.SAMPLE_RATE * 1000  # ms
        axes[row, 0].plot(t_samples, iq[0, :200], 'b-', alpha=0.8, linewidth=0.8)
        axes[row, 0].set_ylabel(f'SNR={snr}dB\n幅度', fontsize=9)
        if row == 0:
            axes[row, 0].set_title(f'时域波形 (I分量)\n[0~1ms, 采样率200kHz]', fontsize=10)
        if row == 2:
            axes[row, 0].set_xlabel('时间 (ms)\n[采样点0~200]', fontsize=9)
        axes[row, 0].grid(True, alpha=0.3, linestyle='--')
        axes[row, 0].set_xlim([0, 1.0])
        
        # 频谱图
        complex_signal = iq[0] + 1j * iq[1]
        axes[row, 1].specgram(complex_signal, Fs=config.SAMPLE_RATE, 
                               NFFT=64, noverlap=32, cmap='viridis')
        if row == 0:
            axes[row, 1].set_title('频谱图 (STFT)\n[NFFT=64, 重叠=32]', fontsize=10)
        if row == 2:
            axes[row, 1].set_xlabel('时间 (采样点)\n[NFFT=64]', fontsize=9)
        axes[row, 1].set_ylabel('频率 (kHz)', fontsize=9)
    
    plt.suptitle(f'{mod_type}调制信号在不同信噪比下的对比\n'
                 '说明：上行=低SNR(-5dB,噪声大) → 中行=中SNR(10dB) → 下行=高SNR(30dB,清晰)',
                 fontsize=12, fontweight='bold', y=0.98)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"对比图已保存到: {save_path}")


def plot_training_history(history_path='./results/training_history.npy', 
                          save_path='./results/training_curve.png'):
    """绘制训练曲线（含坐标标注）"""
    try:
        history = np.load(history_path, allow_pickle=True).item()
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        epochs = range(1, len(history['train_loss']) + 1)
        
        # 损失曲线
        axes[0].plot(epochs, history['train_loss'], 'b-', 
                     label='训练损失 (Train Loss)', linewidth=2)
        axes[0].plot(epochs, history['val_loss'], 'r-', 
                     label='验证损失 (Val Loss)', linewidth=2)
        axes[0].set_xlabel('Epoch (训练轮次)', fontsize=11)
        axes[0].set_ylabel('损失值 (Loss)', fontsize=11)
        axes[0].set_title('模型损失曲线\n[越低越好，验证损失贴近训练损失表示泛化良好]',
                          fontsize=11)
        axes[0].legend(fontsize=10)
        axes[0].grid(True, alpha=0.3, linestyle='--')
        # 标注起始和结束值
        axes[0].annotate(f'初始: {history["train_loss"][0]:.4f}', 
                         xy=(1, history['train_loss'][0]), 
                         xytext=(5, history['train_loss'][0] * 1.1),
                         fontsize=8, color='blue')
        axes[0].annotate(f'最终: {history["train_loss"][-1]:.4f}', 
                         xy=(len(epochs), history['train_loss'][-1]), 
                         xytext=(len(epochs)-10, history['train_loss'][-1] * 1.2),
                         fontsize=8, color='blue')
        
        # 准确率曲线
        axes[1].plot(epochs, history['train_acc'], 'b-', 
                     label='训练准确率 (Train Acc)', linewidth=2)
        axes[1].plot(epochs, history['val_acc'], 'r-', 
                     label='验证准确率 (Val Acc)', linewidth=2)
        axes[1].set_xlabel('Epoch (训练轮次)', fontsize=11)
        axes[1].set_ylabel('准确率 (%)', fontsize=11)
        axes[1].set_title('模型准确率曲线\n[越高越好，验证准确率贴近训练准确率表示无过拟合]',
                          fontsize=11)
        axes[1].legend(fontsize=10)
        axes[1].grid(True, alpha=0.3, linestyle='--')
        axes[1].set_ylim([0, 105])
        # 标注最佳值
        best_epoch = np.argmax(history['val_acc']) + 1
        best_acc = history['val_acc'][best_epoch - 1]
        axes[1].annotate(f'最佳: {best_acc:.1f}%\n(Epoch {best_epoch})', 
                         xy=(best_epoch, best_acc), 
                         xytext=(best_epoch + 3, best_acc - 5),
                         fontsize=9, color='red',
                         arrowprops=dict(arrowstyle='->', color='red', lw=1))
        
        plt.suptitle('训练过程监控\n'
                     f'总共 {len(epochs)} 个 Epoch | Batch Size={config.BATCH_SIZE} | '
                     f'学习率={config.LEARNING_RATE}', 
                     fontsize=12, fontweight='bold', y=1.02)
        plt.tight_layout()
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
        print(f"训练曲线已保存到: {save_path}")
    except Exception as e:
        print(f"无法加载训练历史: {e}")


def plot_data_distribution(save_path='./results/data_distribution.png'):
    """绘制数据分布（含详细标注）"""
    generator = ModulationDataGenerator()
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # 调制类型分布
    mod_counts = [config.SAMPLES_PER_CLASS] * config.NUM_CLASSES
    colors = plt.cm.tab20(np.linspace(0, 1, config.NUM_CLASSES))
    
    bars = axes[0].bar(config.MODULATION_TYPES, mod_counts, color=colors)
    axes[0].set_xlabel('调制类型', fontsize=11)
    axes[0].set_ylabel('样本数量', fontsize=11)
    axes[0].set_title('调制类型样本分布\n[11类均衡分布，每类2000样本]',
                      fontsize=11)
    axes[0].tick_params(axis='x', rotation=45)
    axes[0].grid(axis='y', alpha=0.3, linestyle='--')
    # 标注每类样本数
    for bar in bars:
        height = bar.get_height()
        axes[0].text(bar.get_x() + bar.get_width()/2, height + 20,
                     f'{int(height)}', ha='center', va='bottom', fontsize=7)
    
    # SNR分布
    snr_list = list(config.SNR_RANGE)
    snr_counts = [config.SAMPLES_PER_CLASS // len(snr_list)] * len(snr_list)
    snr_labels = [str(s) for s in snr_list]
    
    bars2 = axes[1].bar(snr_labels, snr_counts, color='steelblue')
    axes[1].set_xlabel('信噪比 SNR (dB)', fontsize=11)
    axes[1].set_ylabel('样本数量', fontsize=11)
    axes[1].set_title(f'信噪比分布\n[共{len(snr_list)}个等级，从{snr_list[0]}dB到{snr_list[-1]}dB，步长5dB]',
                      fontsize=11)
    axes[1].grid(axis='y', alpha=0.3, linestyle='--')
    # 标注每级样本数
    for bar in bars2:
        height = bar.get_height()
        axes[1].text(bar.get_x() + bar.get_width()/2, height + 20,
                     f'{int(height)}', ha='center', va='bottom', fontsize=8)
    
    # 总样本数
    total_samples = config.SAMPLES_PER_CLASS * config.NUM_CLASSES
    total_per_snr = config.SAMPLES_PER_CLASS
    fig.suptitle(f'数据集分布统计\n'
                 f'调制类别: {config.NUM_CLASSES}种 | SNR等级: {len(snr_list)}级 | '
                 f'每类每SNR: {total_per_snr // len(snr_list)}样本 | 总计: {total_samples:,}样本',
                 fontsize=12, fontweight='bold', y=1.02)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"数据分布图已保存到: {save_path}")


def main():
    """主函数"""
    generator = ModulationDataGenerator(seed=config.RANDOM_SEED)
    
    print("正在生成可视化图表...")
    print(f"  采样率: {config.SAMPLE_RATE/1000} kHz")
    print(f"  信号长度: {config.SIGNAL_LENGTH} 点")
    print(f"  载波频率: {config.CARRIER_FREQ/1000} kHz")
    print(f"  调制类型: {config.NUM_CLASSES} 种")
    print(f"  SNR范围: {list(config.SNR_RANGE)} dB")
    print()
    
    plot_signal_waveforms(generator)
    plot_constellation_diagrams(generator)
    plot_spectrograms(generator)
    plot_signal_comparison()
    plot_data_distribution()
    plot_training_history()
    
    print("\n所有可视化图表已生成！")
    print("结果保存在 ./results/ 目录下")


if __name__ == '__main__':
    main()