如何用Scipy稀疏矩阵实现含二维数组的隐式方程？

我太懂从一维隐式格式跳到二维时的困惑了——一维三对角矩阵的构造简直手到擒来，但二维里的上下左右邻域关系一下子就把复杂度拉满了。不过你提到的把二维数组重塑为一维向量确实是这类问题的标准解法，下面我一步步给你拆解怎么落地：

先回顾下你整理好的核心方程：

-alphau[i+1, j] + (1+4alpha)u[i, j] - alphau[i-1, j] - alphau[i, j+1] - alpha*u[i, j-1] = u0[i,j] - h[i,j]

核心思路：二维索引转一维 #

首先得把二维网格点(i,j)映射到一维向量的索引k，假设你的网格是M行N列（i从0到M-1，j从0到N-1），我们可以用行优先的规则：k = i * N + j（当然列优先也可以，只要全程保持一致就行）。

分步实现（用Scipy稀疏矩阵） #

因为二维问题的矩阵会非常大（比如100x100的网格对应10000x10000的矩阵），必须用稀疏矩阵来节省内存，这里用Scipy的csr_matrix来构造：

1. 定义索引映射函数

先写个小工具函数帮我们快速转换索引：

def idx(i, j, num_cols):
    # num_cols是网格的列数，返回(i,j)对应的一维索引
    return i * num_cols + j

2. 构造稀疏矩阵A和右端向量b

遍历每个二维网格点，给稀疏矩阵填充非零元素，同时组装右端向量：

import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import spsolve

# 网格参数（示例值，你可以改成自己的）
M = 10  # 行数
N = 10  # 列数
alpha = 0.1  # 你定义的delta t/delta x²常数

# 初始化示例数据（替换成你自己的u0和h）
u0 = np.random.rand(M, N)
h = np.random.rand(M, N)

# 准备稀疏矩阵的存储容器
data = []
row_indices = []
col_indices = []
b = np.zeros(M * N)  # 右端向量

# 遍历每个网格点
for i in range(M):
    for j in range(N):
        k = idx(i, j, N)
        # 主对角线元素：对应u[i,j]，值为1+4*alpha
        data.append(1 + 4 * alpha)
        row_indices.append(k)
        col_indices.append(k)
        
        # 左边邻点(i,j-1)：如果不是左边界
        if j > 0:
            k_left = idx(i, j-1, N)
            data.append(-alpha)
            row_indices.append(k)
            col_indices.append(k_left)
        # 右边邻点(i,j+1)：如果不是右边界
        if j < N - 1:
            k_right = idx(i, j+1, N)
            data.append(-alpha)
            row_indices.append(k)
            col_indices.append(k_right)
        # 上边邻点(i-1,j)：如果不是上边界
        if i > 0:
            k_up = idx(i-1, j, N)
            data.append(-alpha)
            row_indices.append(k)
            col_indices.append(k_up)
        # 下边邻点(i+1,j)：如果不是下边界
        if i < M - 1:
            k_down = idx(i+1, j, N)
            data.append(-alpha)
            row_indices.append(k)
            col_indices.append(k_down)
        
        # 填充右端向量b
        b[k] = u0[i, j] - h[i, j]

# 构造稀疏矩阵A
A = csr_matrix((data, (row_indices, col_indices)), shape=(M*N, M*N))

# 求解Ax = b
u_flat = spsolve(A, b)

# 把一维结果重塑回二维数组
u = u_flat.reshape(M, N)

额外小贴士 #

• 边界条件处理：上面的代码自动跳过了不存在的边界邻点，如果你的问题有特定边界条件（比如Dirichlet固定值、Neumann导数条件），只需要在循环里对边界点单独修改对应的矩阵元素和b值就行。
• 用spdiags替代循环：如果你想沿用一维的spdiags思路，二维的矩阵会有5条主要对角线：主对角线（偏移0）、上下邻点对应的偏移±N、左右邻点对应的偏移±1。不过要注意边界位置的对角线会有缺失值，需要手动补0，相对来说循环构造csr矩阵更直观，调试起来更方便。

内容的提问来源于stack exchange，提问作者Kristian T