"back substitution" method for solving linear system

Question

I'm using numpy to write the "back substitution" method for solving linear system where "A" is a nonsingular upper triangular matrix.

import numpy as np
def upperTriSol(A, b):
 n = np.size(b)
 x = np.zeros_like(b)
 x[-1] = 1. / A[-1, -1] * b[-1]
 for i in xrange(n-2, -1, -1):
 x[i] = 1. / A[i, i] * (b[i] - np.sum(A[i, i+1:] * x[i+1:]))
 return x

I know that "for loops" are slow, so I wonder if there is any way to avoid them in this case? If not, is there a "correct" way to write efficient "for loops" in numpy?

• Gareth Rees: What was wrong with numpy.linalg.solve?
• Oliver: @GarethRees Nothing wrong with it. It's just some college stuff we need to write by ourselves. It works already, just want to improve a little if possible.
• Gareth Rees: Can you edit the question to explain that, please?

Answer 1

A few quick suggestions that I may flesh out in further edits:

If looks like x[i] is a function of x[i+1] and selected items from A and b. Then it could also be seen as forward itertion on a reversed list, xr[i+1] dependent on xr[i].

np.cumsum and related accumulated functions are often useful in problems like this.

There's a set of tri functions that give you indices and values on upper and lower triangular matrices.

Typically 'vectorizing' in numpy requires stepping back a bit from the problem and imagining in terms of operations on a table or 2d array of values.

(this question would have gotten more attention on SO).