CODEWITHSUNDEEP
pythonpython-3.xperformancenumpynumba
N N
N N

Reputation: 71

Tridiagonal Matrix Algorithm with Numba nopython mode

I am trying to write a TDMA algorithm with numba in nopython mode. Here is my code:

@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)

Then test this solver:

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)

However, my TDMA function fails with a 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]
  ^

It works properly with @jit decorator, but fails with in nopython mode. How should I modify this TDMA function to make it compatible with nopyhon?

The line I commented:

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

Is not compatible to nopython either. Is it possible to use the map function in nopython mode?

I understand that my TDMA may be still slow. So is there a fastest code using python 3 language to implement Tridiagonal Matrix Algorithm?

Upvotes: 1

Views: 675

Answers (1)

MSeifert
MSeifert

Reputation: 152765

The problem is that you have 2D arrays but index and assign them like they are 1D arrays. So you could just ravel() them before passing them to the numba function. I'm not sure if that's actually correct - but for the purpose of this answer I assume that it is.

Also you don't need to copy a and c because you don't modify them and you actually only need to copy the first element of b and d.

So a working function might look like this:

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

And you call it like this:

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

Because I used ravel() the result doesn't have the same shape as the 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]

However I really wouldn't re-implement NumPy functions except when you can utilize some structure in your data that the NumPy function doesn't know. NumPy is a high-performance library that already uses really fine-tuned implementations so a casual re-implementation might only be faster for extremely small data-sets or in case you can exploit some facts about your data (that allow an extremely more performant algorithm).

I have to admit that I don't know "Tridiagonal Matrix Algorithm" but I know that some BLAS libraries (generally incredible fast math libraries) implement it. And NumPy uses BLAS.

However SciPy provides some (very fast) special linear algebra solvers for special matrix types:

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

[...]

Regarding your question with map: The current official list of supported built-in functions does not include map. So you cannot use map in Numbas nopython mode.

Upvotes: 1

Related Questions