基于PINNS物理神经网络的简单N-S方程逆参数求解代码

📅 2026/7/11 22:27:56 👁️ 阅读次数 📝 编程学习
基于PINNS物理神经网络的简单N-S方程逆参数求解代码

本文介绍了利用深度学习网络,求解N-S方程参数的python代码。

先给出N-S方程如下:

本项目,即是通过神经网络训练的方法,求解这个参数ν。

1,首先是生成测试数据:

import torch import numpy as np import matplotlib.pyplot as plt def generate_ns_data(nu_param=0.70, N_f=10000, N_b=500): """ 生成满足 NS 方程的合成数据 Args: nu_param: 粘度系数真值 N_f: 内部配点数量 N_b: 边界/初始时刻采样数量 (用于训练,这里暂未用到) Returns: data_dict: 包含坐标和物理量的字典 """ # 1. 定义空间域和时间域 # x, y 范围 [-pi, pi], t 范围 [0, 1] x_min, x_max = -np.pi, np.pi y_min, y_max = -np.pi, np.pi t_min, t_max = 0.0, 1.0 # 2. 随机采样配点 # 这里的点是用来计算 PDE Loss 的 x_f = np.random.uniform(x_min, x_max, N_f) y_f = np.random.uniform(y_min, y_max, N_f) t_f = np.random.uniform(t_min, t_max, N_f) # 3. 计算解析解 # 这里使用的流场函数,能够完美闭合 NS 方程 # u = -cos(x)sin(y) # v = sin(x)cos(y) # p = -0.25(cos(2x) + cos(2y)) u_true = -np.cos(x_f) * np.sin(y_f) v_true = np.sin(x_f) * np.cos(y_f) p_true = -0.25 * (np.cos(2 * x_f) + np.cos(2 * y_f)) # 4. 验证 NS 方程残差 # 我们手动算一下导数,确认这个数据集是对的 # u_x = sin(x)sin(y), u_xx = cos(x)sin(y) # u_y = -cos(x)cos(y), u_yy = -cos(x)sin(y) # u_t = 0 (定常流场) # 代入 u 的动量方程: # u_t + u*u_x + v*u_y = -p_x + nu * (u_xx + u_yy) # 左边 = 0 + (-cos*sin)*(sin*sin) + (sin*cos)*(-cos*cos) # = -cos*sin*(sin^2 + cos^2) = -cos*sin # 右边 p_x = 0.5*sin(2x) = sin(x)cos(x) (根据压力公式求导) # nu*(u_xx+u_yy) = nu*(cos*sin - cos*sin) = 0 # 等等,这好像不平衡... # --- 修正解析解公式 --- # 为了让方程平衡且包含 nu 项,我们需要更复杂的解析解,或者接受这是一个特殊情况。 # 让我们用最经典的 Navier-Stokes 解析解库中的例子: # u = -cos(x)sin(y)exp(-2nu t) # v = sin(x)cos(y)exp(-2nu t) # p = -0.25(cos(2x) + cos(2y))exp(-4nu t) # 重新计算带有时间衰减的解析解 (Taylor-Green Vortex 形式) decay = np.exp(-2 * nu_param * t_f) u_true = -np.cos(x_f) * np.sin(y_f) * decay v_true = np.sin(x_f) * np.cos(y_f) * decay p_true = -0.25 * (np.cos(2 * x_f) + np.cos(2 * y_f)) * (decay ** 2) # 压力保持量级 # 这组解代入 NS 方程是完全平衡的! # u_t = 2*nu*u # (u*grad)u ... # 最终残差 f_u = 0, f_v = 0 绝对成立。 print(f"生成数据完成,粘度真值 nu = {nu_param}") print(f"配点数量: {N_f}") x_min, x_max = -np.pi, np.pi y_min, y_max = -np.pi, np.pi t_min, t_max = 0.0, 1.0 x_norm = 2 * (x_f - x_min) / (x_max - x_min) - 1 y_norm = 2 * (y_f - y_min) / (y_max - y_min) - 1 t_norm = 2 * (t_f - t_min) / (t_max - t_min) - 1 lambda_x = 2.0 / (x_max - x_min) lambda_y = 2.0 / (y_max - y_min) lambda_t = 2.0 / (t_max - t_min) norm_info = { # 边界值 (用于反归一化或可视化) 'x_min': x_min, 'x_max': x_max, 'y_min': y_min, 'y_max': y_max, 't_min': t_min, 't_max': t_max, # 导数修正系数 (用于 PDE Loss 计算) 'lambda_x': lambda_x, 'lambda_y': lambda_y, 'lambda_t': lambda_t, # 其他物理参数 'nu_true': nu_param } # 5. 打包数据 # 注意:通常训练时 x,y,t 需要归一化,这里先输出原始物理量 data_dict = { 'x': x_norm, 'y': y_norm, 't': t_norm, 'u': u_true, 'v': v_true, 'p': p_true, 'norm_info': norm_info }

2,原始数据,要进行转换,转成pytorch能处理的tensor类型,代码如下:(注意,这里在原始基础上,添加了5%的高斯噪声)

import torch import torch.nn as nn import numpy as np # ========================================== # 1. 数据加载与预处理 (整合前面的步骤) # ========================================== # 假设 data 是我们之前生成的字典 # data = generate_normalized_ns_data(...)
import numpy as np def add_noise(data, noise_level=0.05): """ data: 原始数据,形状 noise_level: 噪声强度,例如 0.05 代表 5% 的噪声 """ # 计算信号的标准差作为基准,或者用最大值幅度 # 方案一:基于信号标准差 (更稳健) signal_std = np.std(data) noise = np.random.normal(0, signal_std * noise_level, data.shape) # 方案二:基于信号最大值 (通信背景常用,类似 SINR) # noise = np.random.normal(0, np.max(np.abs(data)) * noise_level, data.shape) noisy_data = data + noise return noisy_data
# 转换为 Tensor 并归一化 def data_to_tensor(data, device): # 输入 x = torch.from_numpy(data['x'].astype(np.float32)).unsqueeze(1).to(device) y = torch.from_numpy(data['y'].astype(np.float32)).unsqueeze(1).to(device) t = torch.from_numpy(data['t'].astype(np.float32)).unsqueeze(1).to(device) # 输出 u_noisy = add_noise(data['u'], noise_level=0.05) # 加入 5% 的高斯噪声 v_noisy = add_noise(data['v'], noise_level=0.05) p_noisy = add_noise(data['p'], noise_level=0.05) u_true = torch.from_numpy(u_noisy.astype(np.float32)).unsqueeze(1).to(device) v_true = torch.from_numpy(v_noisy.astype(np.float32)).unsqueeze(1).to(device) p_true = torch.from_numpy(p_noisy.astype(np.float32)).unsqueeze(1).to(device) # 启用梯度追踪 (为了计算 PDE Loss) x.requires_grad_(True) y.requires_grad_(True) t.requires_grad_(True) # 组合输入 (Batch, 3) X = torch.cat([x, y, t], dim=1) return X, u_true, v_true, p_true, x, y, t # 执行转换 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("device:" + str(device.type)); import generateData1 norm_info = generateData1.data['norm_info'] X_train, u_train, v_train, p_train, x_c, y_c, t_c = data_to_tensor(generateData1.data, device) print(f"训练数据准备完毕: X形状={X_train.shape}, u形状={u_train.shape}")

3, 定义pytorch进行训练的损失函数如下:

import torch import torch.nn as nn import numpy as np def compute_pde_loss(net, x, y, t, nu_param, norm_info): """ 计算 PDE 残差 net: 神经网络模型 x, y, t: 坐标 Tensor (requires_grad=True) nu_param: 粘度系数 norm_info: 归一化参数字典 (用于修正导数) """ # 1. 前向传播,得到网络预测值 # 假设网络输出 X_input = torch.cat([x, y, t], dim=1) outputs = net(X_input) u_pred = outputs[:, 0:1] v_pred = outputs[:, 1:2] p_pred = outputs[:, 2:3] # 2. 计算导数 # 注意:导数必须对归一化坐标求,然后乘以缩放系数 # 一阶导数 u_x = torch.autograd.grad(u_pred, x, grad_outputs=torch.ones_like(u_pred), create_graph=True)[0] u_y = torch.autograd.grad(u_pred, y, grad_outputs=torch.ones_like(u_pred), create_graph=True)[0] u_t = torch.autograd.grad(u_pred, t, grad_outputs=torch.ones_like(u_pred), create_graph=True)[0] v_x = torch.autograd.grad(v_pred, x, grad_outputs=torch.ones_like(v_pred), create_graph=True)[0] v_y = torch.autograd.grad(v_pred, y, grad_outputs=torch.ones_like(v_pred), create_graph=True)[0] v_t = torch.autograd.grad(v_pred, t, grad_outputs=torch.ones_like(v_pred), create_graph=True)[0] p_x = torch.autograd.grad(p_pred, x, grad_outputs=torch.ones_like(p_pred), create_graph=True)[0] p_y = torch.autograd.grad(p_pred, y, grad_outputs=torch.ones_like(p_pred), create_graph=True)[0] # 二阶导数 u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0] u_yy = torch.autograd.grad(u_y, y, grad_outputs=torch.ones_like(u_y), create_graph=True)[0] v_xx = torch.autograd.grad(v_x, x, grad_outputs=torch.ones_like(v_x), create_graph=True)[0] v_yy = torch.autograd.grad(v_y, y, grad_outputs=torch.ones_like(v_y), create_graph=True)[0] # 3. 导数修正 (至关重要!) # 因为输入是归一化到 [-1, 1] 的,实际物理导数 = 网络导数 * 缩放系数 lambda_x = norm_info['lambda_x'] lambda_y = norm_info['lambda_y'] lambda_t = norm_info['lambda_t'] # 修正一阶导 u_x = u_x * lambda_x u_y = u_y * lambda_y u_t = u_t * lambda_t v_x = v_x * lambda_x v_y = v_y * lambda_y v_t = v_t * lambda_t p_x = p_x * lambda_x p_y = p_y * lambda_y # 修正二阶导 u_xx = u_xx * (lambda_x ** 2) u_yy = u_yy * (lambda_y ** 2) v_xx = v_xx * (lambda_x ** 2) v_yy = v_yy * (lambda_y ** 2) # 4. 计算 NS 方程残差 # 连续性方程 f_e = u_x + v_y # 动量方程 f_u = u_t + (u_pred * u_x + v_pred * u_y) + p_x - nu_param * (u_xx + u_yy) f_v = v_t + (u_pred * v_x + v_pred * v_y) + p_y - nu_param * (v_xx + v_yy) # 计算 MSE Loss loss_f = torch.mean(f_u ** 2) + torch.mean(f_v ** 2) + torch.mean(f_e ** 2) return loss_f

5,定义训练模型,在主函数执行训练:

import torch import torch.nn as nn import numpy as np import geneData from loss_function import compute_pde_loss # ========================================== # 模型定义 (简单示例) # ========================================== class PINN_Net(nn.Module): def __init__(self): super(PINN_Net, self).__init__() self.net = nn.Sequential( nn.Linear(3, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 3) # 输出 u, v, p ) def forward(self, x): return self.net(x) # 初始化 model = PINN_Net().to(geneData.device) # 定义可训练的粘度参数 nu (这是我们要反演的目标) # 初始猜测值设为 0.1 nu_param = torch.nn.Parameter(torch.tensor(0.1, device=geneData.device, dtype=torch.float32)) # 优化器 (同时优化网络权重和 nu) optimizer = torch.optim.Adam(list(model.parameters()) + [nu_param], lr=1e-4) # ========================================== # 训练循环 # ========================================== iterations = 10000 import geneData for i in range(iterations): # 1. 清空梯度 optimizer.zero_grad() # ------------------------- # A. 数据损失 # 这里用到了生成的 u, v, p! # ------------------------- u_pred, v_pred, p_pred = model(geneData.X_train).split(1, dim=1) loss_u = torch.mean((u_pred - geneData.u_train) ** 2) loss_v = torch.mean((v_pred - geneData.v_train) ** 2) loss_p = torch.mean((p_pred - geneData.p_train) ** 2) loss_data = loss_u + loss_v + loss_p # ------------------------- # B. 物理损失 # 这里不需要 u, v, p 真值,只需要坐标 # ------------------------- loss_pde = compute_pde_loss(model, geneData.x_c, geneData.y_c, geneData.t_c, nu_param, geneData.norm_info) # ------------------------- # C. 总损失 # ------------------------- # 权重可以调整,通常 Data Loss 权重较大 lambda_data = 1.0 lambda_pde = 1.0 total_loss = lambda_data * loss_data + lambda_pde * loss_pde # 2. 反向传播 total_loss.backward() # 3. 更新参数 optimizer.step() # 4. 打印日志 if i % 100 == 0: print( f"Iter {i:05d} | Loss: {total_loss.item():.4e} | Data: {loss_data.item():.4e} | PDE: {loss_pde.item():.4e} | Nu: {nu_param.item():.4f}")

在笔者自己的电脑上,用cpu(pytorch版本太高未找到合适的cuda配置)跑了二十多分钟,得到结果如下:

总结,这里训练批量大小是10000,最终学习到的参数为0.5987,与实际的0.6相差不大。当然,可以进一步的加大训练数据。此外,也可以加大噪声比例,以期更接近工业场景。

后面博主会进一步对此进行优化和深入,比如用实际工业场景数据进行训练等。尽请期待。