具有 Numba nopython 模式的三对角矩阵算法答案

【问题标题】：Tridiagonal Matrix Algorithm with Numba nopython mode具有 Numba nopython 模式的三对角矩阵算法
【发布时间】：2019-08-15 06:13:36
【问题描述】：

我正在尝试在 nopython 模式下使用 numba 编写 TDMA algorithm。这是我的代码：

@jit(nopython=True)
def TDMA(a,b,c,d):
  n = len(d)
  x = np.zeros(n)
  w = np.zeros(n)
  #   ac, bc, cc, dc = map(np.copy, (a, b, c, d)) # copy arrays
  ac = np.copy(a)
  bc = np.copy(b)
  cc = np.copy(c)
  dc = np.copy(d)
  for i in range(1,n):
    w[i] = ac[i-1]/bc[i-1]
    bc[i] = bc[i] - w[i]*cc[i-1]
    dc[i] = dc[i] - w[i]*dc[i-1]

  x[n-1] = dc[n-1]/bc[n-1]
  for k in range(n-2,-1,-1):
    x[k] = (dc[k]-cc[k]*x[k+1])/bc[k]
  return np.array(x)

然后测试这个求解器：

A = np.array([[5, 2, 0, 0],[1, 5, 2, 0],[0, 1, 5, 2],[0, 0, 1, 5]],float)
B = np.array([[15],[2],[7],[20]],float)
a = A.diagonal(-1)
b = A.diagonal()
c = A.diagonal(1)
x1 = np.linalg.solve(A,B)
x2 = TDMA(a,b,c,B)
print('by default solver, x1 = ',x1)
print('by TDMA, x2 = ',x2)

但是，我的 TDMA 函数失败并显示 TypingError：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot resolve setitem: array(float64, 1d, C)[int64] = array(float64, 1d, C)

File "<ipython-input-20-e25cda7246bd>", line 16:
def TDMA(a,b,c,d):
    <source elided>

  x[n-1] = dc[n-1]/bc[n-1]
  ^

它与@jit 装饰器一起正常工作，但在nopython 模式下失败。我应该如何修改这个 TDMA 函数以使其与nopyhon 兼容？

我评论的那一行：

ac, bc, cc, dc = map(np.copy, (a, b, c, d)) # copy arrays

也不兼容nopython。 nopython模式下是否可以使用map功能？

我了解我的 TDMA 可能仍然很慢。那么有没有最快的使用python 3语言实现三对角矩阵算法的代码呢？

【问题讨论】：

标签： python python-3.x performance numpy numba

【解决方案1】：

问题是您有二维数组，但索引和分配它们就像它们是一维数组一样。因此，您可以在将它们传递给 numba 函数之前将它们 ravel() 。我不确定这是否真的正确 - 但出于这个答案的目的，我假设它是正确的。

另外你不需要复制a和c，因为你没有修改它们，你实际上只需要复制b和d的第一个元素。

所以一个工作函数可能看起来像这样：

import numba as nb
import numpy as np

@nb.njit
def TDMA(a,b,c,d):
    n = len(d)
    x = np.zeros(n)
    bc = np.zeros(len(b))
    bc[0] = b[0]
    dc = np.zeros(len(d))
    dc[0] = d[0]
    
    for i in range(1, n):
        w = a[i - 1] / bc[i - 1]
        bc[i] = b[i] - w * c[i - 1]
        dc[i] = d[i] - w * dc[i - 1]

    x[n - 1] = dc[n - 1] / bc[n - 1]
    for k in range(n - 2, -1, -1):
        x[k] = (dc[k] - c[k] * x[k + 1]) / bc[k]
    return x

你这样称呼它：

TDMA(a.ravel(), b.ravel(), c.ravel(), B.ravel())

因为我使用了ravel()，所以结果与np.linalg.solve的形状不同：

by default solver, x1 =  [[ 3.05427975]
 [-0.13569937]
 [-0.18789144]
 [ 4.03757829]]
by TDMA, x2 =  [ 3.05427975 -0.13569937 -0.18789144  4.03757829]

但是，我真的不会重新实现 NumPy 函数，除非您可以在数据中使用 NumPy 函数不知道的某些结构。 NumPy 是一个高性能库，它已经使用了真正经过微调的实现，因此对于极小的数据集或者您可以利用有关数据的某些事实（允许性能更高的算法），随意重新实现可能会更快)。

我不得不承认我不知道“三对角矩阵算法”，但我知道一些 BLAS libraries（通常是令人难以置信的快速数学库）实现了它。 NumPy 使用 BLAS。

然而，SciPy 为特殊矩阵类型提供了一些（非常快的）特殊线性代数求解器：

Basics

inv(a[, overwrite_a, check_finite]) 计算矩阵的逆矩阵。

solve(a, b[, sym_pos, lower, overwrite_a, …]) 求解线性方程组 a * x = b 对于未知 x 的平方矩阵。

solve_banded(l_and_u, ab, b[, overwrite_ab, …]) 对 x 求解方程 a x = b，假设 a 是带状矩阵。

solveh_banded(ab, b[, overwrite_ab, …]) 求解方程 a x = b。

solve_circulant(c, b[, singular, tol, …]) 为 x 求解 C x = b，其中 C 是循环矩阵。

solve_triangular(a, b[, trans, lower, …]) 对 x 求解方程 a x = b，假设 a 是一个三角矩阵。

solve_toeplitz(c_or_cr, b[, check_finite]) 使用 Levinson 递归求解 Toeplitz 系统

[...]

关于您对map的问题：目前官方list of supported built-in functions不包括map。所以你不能在 Numbas nopython 模式下使用map。

【讨论】：