傅里叶神经算子简化湍流预测：高效模型与案例分析

近年来，研究者们开始探索神经网络在模拟数值模型方面的应用，以突破传统数值模型计算资源消耗大、计算效率低等问题。其中，一个引人注目的方法是傅里叶神经算子（Fourier Neural Operator，简称FNO）。 四傅里叶神经算子（FNO） 基本原理 FNO的核心在于通过离散傅里叶变换（Discrete Fourier Transform，简称DFT）将输入从物理空间转换到频域空间，再通过滤波器去除高频噪声，最后通过逆傅里叶变换恢复到物理空间。这种方法能够更有效地捕捉数据中的模式，尤其适用于高度复杂的流体动力学问题，如湍流预测。 数据处理流程 FNO处理流程包括特征提取、傅里叶变换、滤波和逆傅里叶变换几个主要步骤。对于图像数据，输入通常是五维的，即批次数（Batch）、时间步（Timestep）、通道数（Channel）、高度（Height）和宽度（Width）。经过特征提取后，输入数据的形状会调整为适合FNO块处理的形式。 滤波 在FNO中，滤波不仅仅是指零除高频部分，还会根据模型的权重调整低频部分的重要性。如果训练过程中某个频率对预测没有影响，其对应的权重值会较低或接近于0。具体实现中，通过两次乘法操作来选择和调整频域内的模式。 ```python class SpectralConv2d_fast(nn.Module): def init(self, in_channels, out_channels, modes1, modes2): super(SpectralConv2d_fast, self).init() self.in_channels = in_channels self.out_channels = out_channels self.modes1 = modes1 self.modes2 = modes2 self.scale = (1 / (in_channels * out_channels)) self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) def forward(self, x): batchsize = x.shape[0] x_ft = torch.fft.rfft2(x, norm='ortho') out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat, device=x.device) out_ft[:, :, :self.modes1, :self.modes2] = compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1) out_ft[:, :, -self.modes1:, :self.modes2] = compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2) x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)), norm='ortho') return x ``` 湍流预测案例 本文使用了一个实际的湍流预测案例，数据集来自rsync命令下载的一个文件库。数据集中包含密度、压力和速度等多个参数，总时间步为1000。实验中，仅预测密度，使用前10个时间步的密度数据作为输入。 ```python 数据集类 class DensityTurbulenceDataset(Dataset): def init(self, dataset_dir, seq_len_total): self.file_list = sorted(glob(os.path.join(dataset_dir, "*.npz"))) self.seq_len_total = seq_len_total self.data_cache = [] for f_path in tqdm(self.file_list, desc=f"加载 {os.path.basename(dataset_dir)} 数据"): self.data_cache.append(np.load(f_path)['arr_0']) def __len__(self): return len(self.file_list) - self.seq_len_total + 1 def __getitem__(self, idx): sequence_data = np.stack([self.data_cache[idx + i] for i in range(self.seq_len_total)]) inputs_full_sequence = sequence_data[:SEQ_LEN_FEATURE, :, :, :] targets_density_sequence = sequence_data[SEQ_LEN_FEATURE:, 0, :, :] return torch.from_numpy(inputs_full_sequence).float(), torch.from_numpy(targets_density_sequence).float() 模型构建与训练 model = Net2d(16, 20).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma) dataset = DensityTurbulenceDataset(dataset_dir, seq_len_total=SEQ_LEN_FEATURE + SEQ_LEN_ROLLOUT) loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True) 训练循环 model.train() for epoch in range(TOTAL_EPOCHS): if epoch <= last_epoch: continue total_loss = 0 for inputs_seq_norm, targets_seq_norm in tqdm(loader, desc=f"Epoch {epoch+1}/{TOTAL_EPOCHS}"): optimizer.zero_grad() current_input_sequence = inputs_seq_norm predicted_density_norm = [] for r_step in range(SEQ_LEN_ROLLOUT): x, y = create_embd_dim(len(current_input_sequence)) used_input_sequence = torch.cat((current_input_sequence, x, y), dim=1).permute(0, 2, 3, 4, 1) pred_norm = model(used_input_sequence).squeeze().unsqueeze(1) predicted_density_norm.append(pred_norm) current_input_sequence = torch.cat([ current_input_sequence[:, 1:, :, :], # 删除最旧的帧 pred_norm # 添加新的预测帧 ], dim=1) predicted_density_norm = torch.cat(predicted_density_norm, dim=1) loss = l2loss(predicted_density_norm, targets_seq_norm) + 0.2 * h1_loss(predicted_density_norm, targets_seq_norm) loss.backward() optimizer.step() total_loss += loss.item() scheduler.step() avg_loss = total_loss / len(loader) if avg_loss < best_loss: best_loss = avg_loss torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': best_loss }, checkpoint_path) ``` 实验结果 训练完成后，FNO模型的表现显著优于U-net模型。特别地，FNO模型生成的图像是173 MB，而U-net模型生成的结果文件大小仅为20 MB。在预测准确性方面，FNO模型能够更好地捕捉湍流中的运动模式，而U-net模型则出现了运动过快的问题。 行业评价与公司背景 FNO作为一种新型的神经网络架构，在处理高维度、复杂物理现象的数据时展现出了独特的优势。与传统的数值模型相比，FNO不仅能够大幅减少计算资源的消耗，还能提高预测的精度和效率。这一技术在气象预报、流体动力学等领域具有广阔的应用前景。研究表明，FNO在处理大规模数据集时的性能尤为突出，未来可能会有更多的应用在这些领域出现。 该研究的代码和数据集可以从GitHub仓库获得，有兴趣的读者可以进一步探索和实验。FNO的开发者来自多个研究机构，旨在推动深度学习在科学计算中的应用。