pytorch在矩形间断区间求二维poisson方程的代码

下面是使用PyTorch求解二维Poisson方程在矩形间断区间上的一个简单实现。在这个例子中，我们使用的是有限差分方法来离散化Poisson方程，并使用PyTorch的自动微分功能来计算梯度。

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 定义Poisson方程的边界条件
def boundary(x, y):
    return torch.sin(x) + torch.sin(y)

# 定义Poisson方程的真实解
def true_solution(x, y):
    return torch.sin(x) + torch.sin(y)

# 定义Poisson方程的初始猜测
def initial_guess(x, y):
    return torch.zeros_like(x)

# 定义有限差分算子
def diff2d(u, h):
    u_xx = (u[2:, 1:-1] - 2 * u[1:-1, 1:-1] + u[:-2, 1:-1]) / (h ** 2)
    u_yy = (u[1:-1, 2:] - 2 * u[1:-1, 1:-1] + u[1:-1, :-2]) / (h ** 2)
    return u_xx + u_yy

# 定义Poisson方程的损失函数
def poisson_loss(u, h, f):
    u_xx = (u[2:, 1:-1] - 2 * u[1:-1, 1:-1] + u[:-2, 1:-1]) / (h ** 2)
    u_yy = (u[1:-1, 2:] - 2 * u[1:-1, 1:-1] + u[1:-1, :-2]) / (h ** 2)
    return ((u_xx + u_yy - f[1:-1, 1:-1]) ** 2).mean()

# 定义矩形区域的大小和分辨率
Lx, Ly = 4, 4
nx, ny = 100, 100
dx, dy = Lx / (nx - 1), Ly / (ny - 1)

# 生成网格点
x = torch.linspace(0, Lx, nx)
y = torch.linspace(0, Ly, ny)
X, Y = torch.meshgrid(x, y)

# 初始化解决方案
u = initial_guess(X, Y)

# 定义边界条件
u[0, :] = boundary(x, 0)
u[-1, :] = boundary(x, Ly)
u[:, 0] = boundary(0, y)
u[:, -1] = boundary(Lx, y)

# 在矩形区域上添加间断
u[int(ny / 4):int(3 * ny / 4), int(nx / 4):int(3 * nx / 4)] = -1

# 定义Poisson方程的右侧函数
f = torch.zeros((nx, ny))
f[1:-1, 1:-1] = -1

# 定义优化器和迭代次数
optimizer = optim.Adam([u], lr=0.01)
n_iters = 10000

# 迭代求解Poisson方程
for i in range(n_iters):
    loss = poisson_loss(u, dx, f)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # 重新设置边界条件
    u[0, :] = boundary(x, 0)
    u[-1, :] = boundary(x, Ly)
    u[:, 0] = boundary(0, y)
    u[:, -1] = boundary(Lx, y)
    
    # 输出当前迭代次数和损失函数值
    if i % 1000 == 0:
        print(f"Iteration {i}: Loss {loss.item()}")
        
# 计算真实解和误差
true_u = true_solution(X, Y)
error = torch.abs(u - true_u)

# 可视化结果
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 10))
plt.subplot(2, 2, 1)
plt.imshow(u.T, cmap='jet', origin='lower', extent=[0, Lx, 0, Ly])
plt.xlabel('x')
plt.ylabel('y')
plt.title('Numerical solution')

plt.subplot(2, 2, 2)
plt.imshow(true_u.T, cmap='jet', origin='lower', extent=[0, Lx, 0, Ly])
plt.xlabel('x')
plt.ylabel('y')
plt.title('True solution')

plt.subplot(2, 2, 3)
plt.imshow(error.T, cmap='jet', origin='lower', extent=[0, Lx, 0, Ly])
plt.xlabel('x')
plt.ylabel('y')
plt.title('Error')

plt.subplot(2, 2, 4)
plt.plot(x, u[:, int(ny / 2)], label='Numerical')
plt.plot(x, true_u[:, int(ny / 2)], label='True')
plt.legend()
plt.xlabel('x')
plt.ylabel('u')
plt.title('Solution along y=2')

plt.show()

在这个例子中，我们在矩形区域的中心添加了一个间断，并使用有限差分算子离散化Poisson方程。在迭代求解过程中，我们使用Adam优化器来最小化损失函数，同时在每次迭代后重新设置边界条件。最终，我们计算了真实解和误差，并将结果可视化。