神经网络拟合高频信号实验

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq
import warnings
warnings.filterwarnings('ignore')

# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)

# 生成高频信号
def generate_high_frequency_signal(n_samples=1000, freq=20):
    """生成高频正弦信号"""
    x = torch.linspace(0, 2*np.pi, n_samples).reshape(-1, 1)
    y = torch.sin(freq * x)  # 高频成分
    y += 0.5 * torch.sin(3 * freq * x)  # 更高频成分
    y += 0.3 * torch.sin(5 * freq * x)  # 超高频成分
    y += 0.1 * torch.randn_like(y)  # 加噪声
    
    return x, y

# 定义神经网络模型
class SimpleNN(nn.Module):
    def __init__(self, hidden_size=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )
    
    def forward(self, x):
        return self.net(x)

# 定义带傅里叶特征的神经网络
class FourierFeatureNN(nn.Module):
    def __init__(self, hidden_size=256, fourier_features=20):
        super().__init__()
        self.fourier_features = fourier_features
        
        # 随机初始化傅里叶特征映射的频率
        self.freqs = nn.Parameter(torch.randn(fourier_features) * 10.0)
        
        # 傅里叶特征映射后的网络
        self.net = nn.Sequential(
            nn.Linear(2 * fourier_features, hidden_size),  # 2* 因为包括sin和cos
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )
    
    def forward(self, x):
        # 傅里叶特征映射
        fourier_encoded = []
        for freq in self.freqs:
            fourier_encoded.append(torch.sin(freq * x))
            fourier_encoded.append(torch.cos(freq * x))
        
        fourier_features = torch.cat(fourier_encoded, dim=1)
        return self.net(fourier_features)

def train_and_evaluate(model, x_train, y_train, x_test, y_test, epochs=5000, lr=0.001):
    """训练和评估模型"""
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    train_losses = []
    test_losses = []
    
    for epoch in range(epochs):
        # 训练
        model.train()
        optimizer.zero_grad()
        y_pred = model(x_train)
        loss = criterion(y_pred, y_train)
        loss.backward()
        optimizer.step()
        
        # 评估
        model.eval()
        with torch.no_grad():
            y_test_pred = model(x_test)
            test_loss = criterion(y_test_pred, y_test)
        
        if epoch % 500 == 0:
            train_losses.append(loss.item())
            test_losses.append(test_loss.item())
            
            if epoch % 1000 == 0:
                print(f'Epoch {epoch:4d}: Train Loss = {loss.item():.6f}, Test Loss = {test_loss.item():.6f}')
    
    return train_losses, test_losses, y_pred.detach()

def analyze_frequency_spectrum(signal, sampling_rate=1000):
    """分析信号的频率谱"""
    n = len(signal)
    signal_np = signal.numpy().flatten()
    
    # 计算FFT
    yf = fft(signal_np)
    xf = fftfreq(n, 1/sampling_rate)[:n//2]
    
    # 计算幅度谱
    magnitude = 2.0/n * np.abs(yf[:n//2])
    
    return xf, magnitude

def visualize_results(x, y_true, y_pred_nn, y_pred_fourier, freq_spectrum_nn, freq_spectrum_fourier):
    """可视化结果"""
    fig, axes = plt.subplots(2, 4, figsize=(15, 10))
    
    # 原始信号 vs NN预测
    axes[0, 0].plot(x.numpy(), y_true.numpy(), 'b-', alpha=0.6, label='Original Signal')
    axes[0, 0].plot(x.numpy(), y_pred_nn.numpy(), 'r--', alpha=0.8, label='NN Prediction')
    axes[0, 0].set_title('Standard NN: Time Domain')
    axes[0, 0].set_xlabel('Time')
    axes[0, 0].set_ylabel('Amplitude')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    
    # 原始信号 vs Fourier NN预测
    axes[0, 1].plot(x.numpy(), y_true.numpy(), 'b-', alpha=0.6, label='Original Signal')
    axes[0, 1].plot(x.numpy(), y_pred_fourier.numpy(), 'g--', alpha=0.8, label='Fourier NN Prediction')
    axes[0, 1].set_title('Fourier Feature NN: Time Domain')
    axes[0, 1].set_xlabel('Time')
    axes[0, 1].set_ylabel('Amplitude')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)
    
    # 两个预测的对比
    axes[0, 2].plot(x.numpy(), y_true.numpy(), 'b-', alpha=0.3, label='Original')
    axes[0, 2].plot(x.numpy(), y_pred_nn.numpy(), 'r--', alpha=0.6, label='Standard NN')
    axes[0, 2].plot(x.numpy(), y_pred_fourier.numpy(), 'g--', alpha=0.6, label='Fourier NN')
    axes[0, 2].set_title('Comparison of Predictions')
    axes[0, 2].set_xlabel('Time')
    axes[0, 2].set_ylabel('Amplitude')
    axes[0, 2].legend()
    axes[0, 2].grid(True, alpha=0.3)

    # 误差对比
    error_nn = torch.abs(y_true - y_pred_nn)
    error_fourier = torch.abs(y_true - y_pred_fourier)
    axes[0, 3].plot(x.numpy(), error_nn.numpy(), 'r-', alpha=0.6, label='Standard NN Error')
    axes[0, 3].plot(x.numpy(), error_fourier.numpy(), 'g-', alpha=0.6, label='Fourier NN Error')
    axes[0, 3].set_title('Prediction Errors')
    axes[0, 3].set_xlabel('Time')
    axes[0, 3].set_ylabel('Absolute Error')
    axes[0, 3].legend()
    axes[0, 3].grid(True, alpha=0.3)
    
    # 频率谱对比
    truncated_freq_index = 101
    axes[1, 0].plot(freq_spectrum_nn[0][:truncated_freq_index ], freq_spectrum_nn[1][:truncated_freq_index ], 'r-', label='NN Prediction Spectrum')
    axes[1, 0].set_title('Standard NN: Frequency Spectrum')
    axes[1, 0].set_xlabel('Frequency (Hz)')
    axes[1, 0].set_ylabel('Magnitude')
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)
    
    axes[1, 1].plot(freq_spectrum_fourier[0][:truncated_freq_index ], freq_spectrum_fourier[1][:truncated_freq_index ], 'g-', label='Fourier NN Spectrum')
    axes[1, 1].set_title('Fourier Feature NN: Frequency Spectrum')
    axes[1, 1].set_xlabel('Frequency (Hz)')
    axes[1, 1].set_ylabel('Magnitude')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)
    
    axes[1, 2].plot(freq_spectrum_original[0][:truncated_freq_index], freq_spectrum_original[1][:truncated_freq_index], 'g-', label='Fourier NN Spectrum',
              linewidth=1.5)
    axes[1, 2].set_title('Original Frequency Spectrum')
    axes[1, 2].set_xlabel('Frequency (Hz)')
    axes[1, 2].set_ylabel('Magnitude')
    axes[1, 2].legend()
    axes[1, 2].grid(True, alpha=0.3)

    error_nn = freq_spectrum_original[1] - freq_spectrum_nn[1]
    error_fourier = freq_spectrum_original[1] - freq_spectrum_fourier[1]
    axes[1, 3].plot(freq_spectrum_original[0][:truncated_freq_index], error_nn[:truncated_freq_index], 'r-', alpha=0.6, label='Standard NN  Error', linewidth=1.5)
    axes[1, 3].plot(freq_spectrum_original[0][:truncated_freq_index], error_fourier[:truncated_freq_index], 'g-', alpha=0.6, label='Fourier NN Error', linewidth=1.5)
    axes[1, 3].set_title('Prediction Errors')
    axes[1, 3].set_xlabel('Frequency (Hz)')
    axes[1, 3].set_ylabel('Magnitude Error')
    axes[1, 3].legend()
    axes[1, 3].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()

# 计算频谱偏差
def calculate_spectral_bias(original_spectrum, predicted_spectrum, n_freqs=50):
        """计算频率谱的偏差"""
        bias = []
        for i in range(min(len(original_spectrum[1]), n_freqs)):
            orig_mag = original_spectrum[1][i]
            pred_mag = predicted_spectrum[1][i]
            if orig_mag > 0.01:  # 只考虑显著频率
                relative_error = abs(orig_mag - pred_mag) / orig_mag
                bias.append((original_spectrum[0][i], relative_error))
        return bias

def main():
    print("=" * 60)
    print("神经网络拟合高频信号实验")
    print("=" * 60)
    
    # 生成数据
    print("\n1. 生成高频信号...")
    n_samples = 100000
    train_test_ratio = 0.8
    n_train = int(n_samples * train_test_ratio)
    x, y = generate_high_frequency_signal(n_samples, freq=20)
    x_train, y_train = x[:n_train], y[:n_train]
    x_test, y_test = x[n_train:], y[n_train:]
    
    print(f"信号形状: x={x.shape}, y={y.shape}")
    print(f"信号频率: 包含20Hz, 60Hz, 100Hz的高频成分")
    
    # 训练标准神经网络
    print("\n2. 训练标准神经网络...")
    model_nn = SimpleNN(hidden_size=256)
    train_losses_nn, test_losses_nn, y_pred_nn = train_and_evaluate(
        model_nn, x_train, y_train, x_test, y_test, epochs=5000, lr=0.001
    )
    
    # 训练傅里叶特征神经网络
    print("\n3. 训练傅里叶特征神经网络...")
    model_fourier = FourierFeatureNN(hidden_size=256, fourier_features=20)
    train_losses_fourier, test_losses_fourier, y_pred_fourier = train_and_evaluate(
        model_fourier, x_train, y_train, x_test, y_test, epochs=5000, lr=0.001
    )
    
    # 分析频率谱
    print("\n4. 分析频率谱...")
    freq_spectrum_original = analyze_frequency_spectrum(y, sampling_rate=1000)
    freq_spectrum_nn = analyze_frequency_spectrum(y_pred_nn, sampling_rate=1000)
    freq_spectrum_fourier = analyze_frequency_spectrum(y_pred_fourier, sampling_rate=1000)
    
    # 打印频率分析结果
    print("\n频率分析:")
    print(f"原始信号主要频率: {freq_spectrum_original[0][np.argmax(freq_spectrum_original[1][:50])]:.1f} Hz")
    print(f"标准NN预测主要频率: {freq_spectrum_nn[0][np.argmax(freq_spectrum_nn[1][:50])]:.1f} Hz")
    print(f"傅里叶NN预测主要频率: {freq_spectrum_fourier[0][np.argmax(freq_spectrum_fourier[1][:50])]:.1f} Hz")
    
    # 计算和显示最终误差
    final_error_nn = torch.mean(torch.abs(y_test - model_nn(x_test))).item()
    final_error_fourier = torch.mean(torch.abs(y_test - model_fourier(x_test))).item()
    
    print("\n" + "=" * 60)
    print("实验结果总结:")
    print("=" * 60)
    print(f"标准神经网络最终测试误差: {final_error_nn:.6f}")
    print(f"傅里叶特征神经网络最终测试误差: {final_error_fourier:.6f}")
    print(f"误差改进: {(final_error_nn - final_error_fourier) / final_error_nn * 100:.1f}%")
    
    print("\n关键观察:")
    print("1. 标准神经网络难以拟合高频信号（频谱偏差）")
    print("2. 傅里叶特征变换显著改善了高频信号的拟合能力")
    print("3. 高频信号需要网络学习快速变化的函数，这超出了ReLU网络的频率容量")
    print("4. 傅里叶特征映射将输入投影到高频空间，帮助网络捕捉快速变化")
    
    
    
    # 额外实验：展示频谱偏差
    print("\n" + "=" * 60)
    print("频谱偏差分析:")
    print("=" * 60)
    
    
    
    bias_nn = calculate_spectral_bias(freq_spectrum_original, freq_spectrum_nn)
    bias_fourier = calculate_spectral_bias(freq_spectrum_original, freq_spectrum_fourier)
    
    print("\n频率相对误差（越高表示拟合越差）:")
    print("频率(Hz) | 标准NN误差 | 傅里叶NN误差")
    print("-" * 40)
    
    for i in range(min(10, len(bias_nn))):
        freq = bias_nn[i][0]
        nn_error = bias_nn[i][1]
        fourier_error = bias_fourier[i][1] if i < len(bias_fourier) else float('nan')
        print(f"{freq:7.1f} | {nn_error:10.3f} | {fourier_error:10.3f}")
    
    # 可视化
    visualize_results(x, y, y_pred_nn, y_pred_fourier, freq_spectrum_nn, freq_spectrum_fourier)

    return {
        'model_nn': model_nn,
        'model_fourier': model_fourier,
        'errors': (final_error_nn, final_error_fourier),
        'spectral_bias': (bias_nn, bias_fourier)
    }

if __name__ == "__main__":
    # 检查环境
    print(f"Python版本: 3.12.10")
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA版本: {torch.version.cuda}")
    
    # 运行实验
    results = main()