共轭梯度法代码输出NaN求助：rnorm/actRnorm无法绘图

我编写了如下共轭梯度法（conjugatGradient）代码，但rnorm和actRnorm列表始终输出NaN，导致无法绘制对应曲线，仅estimate曲线可正常显示，恳请提供解决建议。

代码：

def conjugatGradient(A,b,n):
  x = b
  r = b - A*x
  p = r
  r_old = r.conj().T*r

  rnorm = []
  actRnorm = []
  estimate = []
  
  for i in range(n):
    Ap = A*p
    alpha = r_old/(p.conj().T*Ap)
    x = x + alpha*p
    r = r - alpha*Ap
    r_new = r.conj().T*r
    if np.sqrt(r_new).any() < 1e-10:
      break
    p = r + (r_new/r_old)*p
    r_old = r_new
    
    rnorm.append(np.linalg.norm(r))
    actRnorm.append(np.linalg.norm(b-A*x))
    estimate.append(2*((np.sqrt(np.linalg.cond(A))-1)/(np.sqrt(np.linalg.cond(A))+1))**i)

  plt.plot(rnorm)
  plt.plot(actRnorm)
  plt.plot(estimate)
  plt.show()

矩阵A和向量b定义如下：

import numpy as np
import matplotlib.pyplot as plt

A = np.diag(range(1,101))
for i in range(1,100):
  A[i-1,i] = 1
  A[i,i-1] = 1

b = np.ones(100)

• 修正矩阵乘法方式：代码中用*做矩阵与向量的乘法是元素级运算，而非线性代数中的矩阵乘法，这会导致Ap、alpha等核心变量计算完全错误，最终产生NaN。需将所有矩阵-向量、向量内积的乘法替换为@（numpy标准矩阵乘法运算符）：

  • 替换A*x为A@x
  • 替换A*p为A@p
  • 替换r.conj().T*r为r.conj().T @ r
  • 替换p.conj().T*Ap为p.conj().T @ Ap

• 将内积结果转为标量：r.conj().T @ r得到的是1x1的矩阵，直接参与除法会引发广播错误，需用.item()提取标量值：

  • r_old = (r.conj().T @ r).item()
  • r_new = (r.conj().T @ r).item()

• 优化初始值设置：共轭梯度法通常将初始迭代值x设为全零向量，而非b，这能保证初始残差r = b - A@x的计算符合标准逻辑：

  • x = np.zeros_like(b)

• 修正终止条件判断：np.sqrt(r_new).any()无意义（r_new是标量），直接判断np.sqrt(r_new) <= 1e-10即可，避免无效的数组操作。

修复后的完整代码：

import numpy as np
import matplotlib.pyplot as plt

def conjugatGradient(A,b,n):
  x = np.zeros_like(b)  # 初始化为全零向量
  r = b - A@x
  p = r
  r_old = (r.conj().T @ r).item()  # 转为标量

  rnorm = []
  actRnorm = []
  estimate = []
  
  for i in range(n):
    Ap = A@p
    pAp = (p.conj().T @ Ap).item()
    if pAp == 0:  # 避免除以零
        break
    alpha = r_old / pAp
    x = x + alpha*p
    r = r - alpha*Ap
    r_new = (r.conj().T @ r).item()
    
    if np.sqrt(r_new) <= 1e-10:
      break
    
    p = r + (r_new/r_old)*p
    r_old = r_new
    
    rnorm.append(np.linalg.norm(r))
    actRnorm.append(np.linalg.norm(b - A@x))
    estimate.append(2*((np.sqrt(np.linalg.cond(A))-1)/(np.sqrt(np.linalg.cond(A))+1))**i)

  plt.plot(rnorm, label='r norm')
  plt.plot(actRnorm, label='actual r norm')
  plt.plot(estimate, label='estimate')
  plt.legend()
  plt.show()

# 定义矩阵A和向量b
A = np.diag(range(1,101))
for i in range(1,100):
  A[i-1,i] = 1
  A[i,i-1] = 1

b = np.ones(100)

# 调用函数
conjugatGradient(A, b, 100)

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