如何用SciPy生成方格上薛定谔算子的稀疏矩阵？

你的原代码用稠密矩阵生成N²×N²的邻接矩阵，当N稍大时必然触发内存不足——比如N=100时，矩阵规模是10000×10000，仅存储浮点数就需要约800MB内存，N=200更是超过3GB。用SciPy的sparse包可以完美解决这个问题，只存储非零元素，内存占用会大幅降低。

以下是分步实现方案：

1. 生成稀疏邻接矩阵A #

邻接矩阵A的非零元素仅对应网格中相邻的点（上下左右），我们可以用**坐标格式（COO）**来构建它：

• 遍历每个网格点(i,j)，计算其线性索引row_idx = i*N + j
• 检查四个方向的邻居：(i+1,j)、(i-1,j)、(i,j+1)、(i,j-1)
• 若邻居在网格范围内（0 ≤ 坐标 < N），计算邻居的线性索引col_idx，并记录这一组(row_idx, col_idx)作为非零位置
• 最后用scipy.sparse.coo_matrix生成稀疏矩阵

代码示例：

import numpy as np
from scipy import sparse

def build_adjacency_matrix(N):
    rows = []
    cols = []
    # 遍历每个网格点
    for i in range(N):
        for j in range(N):
            row_idx = i * N + j
            # 检查四个方向的邻居
            for di, dj in [(1,0), (-1,0), (0,1), (0,-1)]:
                ni, nj = i + di, j + dj
                if 0 <= ni < N and 0 <= nj < N:
                    col_idx = ni * N + nj
                    rows.append(row_idx)
                    cols.append(col_idx)
    # 构建COO格式的稀疏矩阵，非零值全为1
    A = sparse.coo_matrix((np.ones(len(rows)), (rows, cols)), shape=(N*N, N*N))
    # 转成CSR格式，方便后续运算
    return A.tocsr()

2. 生成稀疏对角矩阵D #

不需要用np.diag生成稠密对角矩阵，直接用scipy.sparse.diags生成稀疏对角矩阵：

def build_diagonal_matrix(arr):
    # arr是长度为N²的一维数组，对应对角线上的元素
    return sparse.diags(arr, format='csr')

3. 组合得到薛定谔算子H #

稀疏矩阵支持直接加法运算，最终得到H：

# 示例使用：假设N=5，arr是长度为25的对角元素数组
N = 5
arr = np.random.rand(N*N)  # 替换成你的实际对角元素

A = build_adjacency_matrix(N)
D = build_diagonal_matrix(arr)
H = A + D

关键说明 #

• 内存效率：COO/CSR格式仅存储非零元素，对于N×N网格，邻接矩阵A的非零元素数量是4N² - 4N（边界点少两个邻居），相比稠密矩阵的N⁴元素，内存占用几乎可以忽略
• 运算兼容性：SciPy稀疏矩阵支持大部分线性代数运算（比如求解特征值、矩阵向量乘法），不需要转成稠密矩阵就能使用
• 代码简化：去掉了原代码的四重循环，用二重循环+方向遍历生成非零位置，速度更快且逻辑更清晰

内容的提问来源于stack exchange，提问作者Keen-ameteur