写一个求解 偏微分方程python 代码

在Python中，我们可以使用一些库来解决偏微分方程（PDE），其中最常用的是`scipy`和`numpy`库。这里我们将以最基础的二维泊松方程为例，展示如何编写一个求解器。泊松方程是一个典型的线性PDE，形式为：

∇²u = f(x, y)

其中，u是未知函数，f是给定的源函数，∇²是拉普拉斯算子。

import numpy as np
from scipy.sparse import spdiags, csc_matrix
from scipy.sparse.linalg import spsolve

def laplacian_matrix(n, dx, dy):
    # 创建一个二维网格
    x = np.linspace(0, 1, n+1)[:-1] + dx/2
    y = np.linspace(0, 1, n+1)[:-1] + dy/2
    X, Y = np.meshgrid(x, y)

    # 网格点的行和列索引
    rows = np.ravel_multi_index(np.dstack((Y.flatten(), X.flatten())), (n, n))
    cols = rows.copy()
    data = np.ones_like(rows)

    # 拉普拉斯矩阵的第一、二对角线元素
    laplacian = spdiags([-data, data, -data], [-1, 0, 1], n**2, n**2, format='csc')

    return laplacian, rows, cols

def solve_poisson(f, n=100, dx=0.01, dy=0.01):
    # 计算拉普拉斯矩阵
    laplacian, rows, cols = laplacian_matrix(n, dx, dy)
    
    # 将边界条件转换为内点矩阵
    b = np.zeros(n**2)
    b = f(0, 0)  # 左上角点
    b[-1] = f(1, 1)  # 右下角点
    b[rows == 0] = f(x, y)  # 左边
    b[rows == (n**2-1)] = f(x[-1], y[-1])  # 右边
    b[cols == 0] = f(x, y)  # 底边
    b[cols == (n**2-1)] = f(x, y[-1])  # 顶边

    # 解方程
    u = spsolve(laplacian, b)

    return u.reshape((n, n))

# 示例源函数
def source_function(x, y):
    return np.sin(2*np.pi*x) * np.sin(2*np.pi*y)

# 解决并可视化结果
solution = solve_poisson(source_function, n=100, dx=0.01)