Python求解稀疏矩阵Au=f时输出出现NaN问题排查

问题描述 #

求解稀疏矩阵方程Au=f的近似解时，所有输入矩阵均为np.float64类型，但输出结果出现NaN，无法定位问题根源。

原代码 #

import numpy as np
import scipy
from scipy.sparse import coo_matrix
from scipy.sparse import random

#Defining a function that takes N as an input

def our_func (N) :

  #Taking the given values of parameters h and k
  h =  1/N
  k = (29*(np.pi))/2
   
  #Defining initial input matrix with domensions N+1
  I = np.empty([N+1,N+1],  dtype=np.float64) 
  #[REF :- https://www.geeksforgeeks.org/how-to-create-an-empty-matrix-with-numpy-in-python/]
  print("datatype of I is :-")
  print(I.dtype)

  #Generating the matrix A in coo sparse format
  A_coo = coo_matrix(I).reshape((N+1,N+1))
  print("datatype of A_coo is :-")
  print(A_coo.dtype)

  #Converting the matrix A from coo sparse format to CSR sparse format 
  A = A_coo.tocsr().reshape((N+1,N+1))
  print("datatype of A is :-")
  print(A.dtype)
 
  # Creating an empty storage array for the right side of our equation Au=f
  f = np.empty((N+1), dtype=np.float64)

  #The rows of matrix A and f are defined as :-
  A[0,0] = 1    #for i==j==0
  A[N,N] = 1    #for i==j==N
  f[N] = 1      #for i==N
  f[0] = 0      #for i==0
  for i in range (int(N)) :
    for j in range (int(N)) :
        A[i,1] = 2 - ((h**2)*(k**2)) 
        A[i,i+1] = -1
        A[i,i-1] = -1
        A[i,j] = 0      #for i!==j
  
  #Our function Returns matrix A and f, with matrix A in sparse storage format
  return A, f

# Defining a new function using the function scipy.sparse.linalg.spsolve 
# to solve the sparse matrix-vector problem: Au=f

def our_solution(N):
  A,f = our_func(N)
  our_sparse_matrix = scipy.sparse.linalg.spsolve(A,f)
  return our_sparse_matrix

# Using this defined function to compute the approximate solution for our 
# problem, for N=10, N=100, and N=1000.
approx_sol_1 = our_solution(10)
print("approx_sol_1 is :- ")
print(approx_sol_1)

问题根源 #

1. 未初始化的矩阵与数组：
   • 用np.empty创建的I矩阵和f数组包含随机垃圾值，转成稀疏矩阵后这些值会保留，后续未覆盖的位置仍为无效值，导致矩阵结构异常或求解时出现NaN。
2. 循环逻辑完全错误：
   • 双重循环中，每次循环都对A[i,1]重复赋值，且执行A[i,j] = 0会覆盖之前设置的A[i,i+1]和A[i,i-1]，最终中间行的有效元素被清空，矩阵变为奇异矩阵，求解时必然出现NaN。
3. 稀疏矩阵赋值方式不当：直接对CSR矩阵进行索引赋值效率低，且容易遗漏或覆盖元素，不符合稀疏矩阵的构建规范。

修复后的代码 #

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

def our_func(N):
    h = 1 / N
    k = (29 * np.pi) / 2
    size = N + 1
    main_diag = np.zeros(size, dtype=np.float64)
    upper_diag = np.zeros(size - 1, dtype=np.float64)
    lower_diag = np.zeros(size - 1, dtype=np.float64)
    
    # 设置边界条件
    main_diag[0] = 1
    main_diag[-1] = 1
    
    # 设置中间行的对角线元素
    middle_val = 2 - (h ** 2) * (k ** 2)
    main_diag[1:-1] = middle_val
    upper_diag[1:] = -1  # 上对角线（i, i+1）
    lower_diag[:-1] = -1 # 下对角线（i, i-1）
    
    # 构建CSR格式的稀疏矩阵
    A = csr_matrix(([main_diag, upper_diag, lower_diag], [0, 1, -1]), shape=(size, size))
    
    # 初始化f数组，默认全0，再设置边界值
    f = np.zeros(size, dtype=np.float64)
    f[-1] = 1
    
    return A, f

def our_solution(N):
    A, f = our_func(N)
    return scipy.sparse.linalg.spsolve(A, f)

# 测试N=10的情况
approx_sol_1 = our_solution(10)
print("approx_sol_1 is :- ")
print(approx_sol_1)

修复说明 #

• 改用对角线数组构建稀疏矩阵，避免未初始化的垃圾值，同时符合稀疏矩阵的高效构建方式。
• 修正循环逻辑，直接对对角线数组赋值，准确设置矩阵的三对角结构。
• 用np.zeros初始化f数组，确保所有元素初始为0，仅修改边界值，避免无效值干扰求解。

内容的提问来源于stack exchange，提问作者Anushka Tilekar