如何让PyTorch CNN求解器适配变网格尺寸与非均匀加密？

核心问题分析 #

• 卷积核权重错误：均匀网格下的固定1.0权重不适用于非均匀网格，热传导的邻域权重需与网格步长成反比。
• 边界条件缺失：未定义边界约束，导致边界区域数值异常发散。
• 迭代逻辑偏离离散格式：当前更新公式不符合热传导方程的有限差分离散形式，导致数值结果失真。
• 非均匀网格适配不足：未生成对应非均匀网格的步长张量，系数计算仍沿用均匀网格的固定hx/hy。

修正方案 #

1. 生成非均匀网格：通过非线性采样实现局部加密，同时为每个网格点生成对应步长张量。
2. 调整卷积核权重：基于邻域步长动态计算权重，替代固定值。
3. 修复迭代更新逻辑：严格遵循非均匀网格下的热传导离散格式更新温度场。
4. 添加边界条件：设置Dirichlet边界（边界温度固定为0），约束边界数值。

修正后的代码 #

import torch 
import torch.nn as nn
import sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator
from matplotlib import cm
import copy

def plot_solution(X,Y,Z):
    fig = plt.figure(figsize=(10,10))
    ax = plt.axes(projection='3d')
    surf=ax.plot_surface(X,Y,Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)
    ax.set_title("2D heat Variation, CNN (Non-uniform Grid)")
    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter('{x:.02f}')
    fig.colorbar(surf, shrink=0.5)
    plt.show()

# 参数设置
E = 1e-4
x = 1.0
y = 1.0
Nx = 50  # x方向总点数
Ny = 50  # y方向总点数
Px = Nx + 1
Py = Ny + 1
max_iterations = 2000
source_strength = 10.0  # 调整源项强度避免数值过高

# 生成非均匀网格：中心区域加密
x_nodes = np.linspace(0, 1, Nx+1)
x_nodes = x_nodes + 0.2 * np.sin(np.pi * x_nodes)  # 中心区域步长缩小
x_nodes = np.sort(x_nodes)
hx = np.diff(x_nodes)
hx = np.concatenate([[hx[0]], hx, [hx[-1]]])  # 扩展至节点数长度，边界复用相邻步长

y_nodes = np.linspace(0, 1, Ny+1)
y_nodes = y_nodes + 0.2 * np.sin(np.pi * y_nodes)
y_nodes = np.sort(y_nodes)
hy = np.diff(y_nodes)
hy = np.concatenate([[hy[0]], hy, [hy[-1]]])

# 生成网格矩阵
X, Y = np.meshgrid(x_nodes, y_nodes)
Z = np.zeros_like(X)

# 转换为PyTorch张量
hx_tensor = torch.tensor(hx, dtype=torch.float32).view(1, 1, Px, 1).repeat(1,1,1,Py)
hy_tensor = torch.tensor(hy, dtype=torch.float32).view(1, 1, 1, Py).repeat(1,1,Px,1)

# 定义卷积核：仅保留上下左右邻域连接
weights = torch.zeros(1,1,3,3)
with torch.no_grad():
    weights[0,0,1,0] = 1.0  # 左邻域
    weights[0,0,1,2] = 1.0  # 右邻域
    weights[0,0,0,1] = 1.0  # 上邻域
    weights[0,0,2,1] = 1.0  # 下邻域

temperature_conv = nn.Conv2d(1, 1, 3, stride=1, padding=1, bias=False)
with torch.no_grad():
    temperature_conv.weight.copy_(weights)

# 初始化温度场、源项、扩散系数
solution = torch.zeros(1, 1, Px, Py, dtype=torch.float32)
# 设置Dirichlet边界：边界温度固定为0
solution[:,:,0,:] = 0.0
solution[:,:,-1,:] = 0.0
solution[:,:,:,0] = 0.0
solution[:,:,:,-1] = 0.0

src = torch.zeros(1, 1, Px, Py, dtype=torch.float32)
# 中心区域设置源项
center_x_start = int(Px*0.3)
center_x_end = int(Px*0.7)
center_y_start = int(Py*0.3)
center_y_end = int(Py*0.7)
src[:,:,center_x_start:center_x_end, center_y_start:center_y_end] = source_strength

diff = torch.ones(1, 1, Px, Py, dtype=torch.float32)
# 中心区域调整扩散系数（模拟非均匀介质）
diff[:,:,center_x_start:center_x_end, center_y_start:center_y_end] = 0.5

# 计算非均匀网格的离散系数
hx_left = hx_tensor[:,:,:-2,:]
hx_right = hx_tensor[:,:,2:,:]
hx_avg = (hx_left + hx_right)/2
hy_top = hy_tensor[:,:,:,:-2]
hy_bottom = hy_tensor[:,:,:,2:]
hy_avg = (hy_top + hy_bottom)/2

diag_coeff = torch.zeros_like(solution)
off_diag_coeff = torch.zeros_like(solution)

with torch.no_grad():
    # 内部点系数计算
    diag_coeff[:,:,1:-1,1:-1] = 1.0 / ( diff[:,:,1:-1,1:-1] * ( (1/hx_left + 1/hx_right)/hx_avg + (1/hy_top + 1/hy_bottom)/hy_avg ) )
    # 边界系数复用相邻内部点值
    diag_coeff[:,:,0,:] = diag_coeff[:,:,1,:]
    diag_coeff[:,:,-1,:] = diag_coeff[:,:,-2,:]
    diag_coeff[:,:,:,0] = diag_coeff[:,:,:,1]
    diag_coeff[:,:,:,-1] = diag_coeff[:,:,:,-2]

# 迭代求解
uncertainty = float('inf')
for i in range(max_iterations):
    prev_solution = solution.clone()
    # 计算邻域加权和
    neighbor_sum = temperature_conv(solution)
    # 按离散格式更新内部点温度
    solution[:,:,1:-1,1:-1] = (neighbor_sum[:,:,1:-1,1:-1] + src[:,:,1:-1,1:-1]) * diag_coeff[:,:,1:-1,1:-1]
    # 保持边界条件不变
    solution[:,:,0,:] = 0.0
    solution[:,:,-1,:] = 0.0
    solution[:,:,:,0] = 0.0
    solution[:,:,:,-1] = 0.0
    # 计算收敛误差
    uncertainty = torch.max(torch.abs(solution - prev_solution)).item()
    if uncertainty < E:
        print(f"Converged at iteration {i}")
        break

if uncertainty >= E:
    print("Solution doesn't converge within the maximum number of iterations") 

# 可视化结果
with torch.no_grad():
    Z[:,:] = solution[0,0,:,:].detach().numpy()
plot_solution(X,Y,Z)

关键修正说明 #

• 非均匀网格生成：通过正弦函数扰动均匀网格，实现中心区域加密，同时为每个节点生成对应步长张量。
• 自适应卷积逻辑：卷积核仅保留上下左右邻域连接，权重通过步长系数动态调整，适配非均匀网格。
• 边界约束：显式设置Dirichlet边界，避免边界数值发散。
• 迭代格式修正：严格遵循非均匀网格下的有限差分格式更新温度场，保证物理计算的一致性。

内容的提问来源于stack exchange，提问作者Egandoue Peroumal