Трехдиагональный матричный алгоритм с режимом Numba nopython - PullRequest
Новая
       17

Трехдиагональный матричный алгоритм с режимом Numba nopython

1 голос
/ 25 марта 2019

Я пытаюсь написать алгоритм TDMA с numba в режиме nopython. Вот мой код: 

@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. Можно ли использовать функцию map в режиме nopython?

Я понимаю, что моя TDMA может все еще работать медленно. Так есть ли самый быстрый код, использующий язык Python 3 для реализации алгоритма трехдиагональной матрицы?

1 Ответ

1 голос
/ 25 марта 2019

Проблема в том, что у вас есть двумерные массивы, но индексируйте и присваивайте их, как будто они являются одномерными. Таким образом, вы можете просто ravel() их перед передачей в функцию numba. Я не уверен, что это действительно правильно - но для целей этого ответа я предполагаю, что это так.

Также вам не нужно копировать 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 (как правило, невероятные быстрые математические библиотеки) реализуют его. И NumPy использует BLAS.

Однако SciPy предоставляет некоторые (очень быстрые) специальные решатели линейной алгебры для специальных типов матриц:

Основы

  • 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, …]) Решите уравнение a x = b для x, предполагая, что a - матрица с полосами.
  • solveh_banded(ab, b[, overwrite_ab, …]) Решить уравнение a x = b.
  • solve_circulant(c, b[, singular, tol, …]) Решите C x = b для x, где C - циркулянтная матрица.
  • solve_triangular(a, b[, trans, lower, …]) Решите уравнение a x = b для x, предполагая, что a является треугольной матрицей.
  • solve_toeplitz(c_or_cr, b[, check_finite]) Решите систему Теплица, используя Levinson Recursion

[...]

Относительно вашего вопроса с map: Текущий официальный список поддерживаемых встроенных функций не включает map. Поэтому вы не можете использовать map в режиме Numbas Nopython.

...