gpt4 book ai didi

python - Numba nopython 模式的三对角矩阵算法

转载 作者:行者123 更新时间:2023-12-01 08:09:42 25 4
gpt4 key购买 nike

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

@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语言实现三对角矩阵算法的代码呢?

最佳答案

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

此外,您也不需要复制 ac,因为您不修改它们,实际上只需要复制 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]) Compute the inverse of a matrix.
  • solve(a, b[, sym_pos, lower, overwrite_a, …]) Solves the linear equation set a * x = b for the unknown x for square a matrix.
  • solve_banded(l_and_u, ab, b[, overwrite_ab, …]) Solve the equation a x = b for x, assuming a is banded matrix.
  • solveh_banded(ab, b[, overwrite_ab, …]) Solve equation a x = b.
  • solve_circulant(c, b[, singular, tol, …]) Solve C x = b for x, where C is a circulant matrix.
  • solve_triangular(a, b[, trans, lower, …]) Solve the equation a x = b for x, assuming a is a triangular matrix.
  • solve_toeplitz(c_or_cr, b[, check_finite]) Solve a Toeplitz system using Levinson Recursion

[...]

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

关于python - Numba nopython 模式的三对角矩阵算法,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/55330433/

25 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com