Python: vectorizing a structured linear system solve

Question

Overview

I am looking for a way to solve a structured linear system in Python without using a for loop (preferably using vectorization, if possible).

Background

Consider the following linear system: \begin{align} \begin{pmatrix} E_0 \\ F_1 & E_1 \\ & F_2 & E_2 \\ && \ddots & \ddots \\ &&& F_{K-1} & E_{K-1} \end{pmatrix} \begin{pmatrix} x_0 \\ x_1 \\ x_2 \\ \vdots \\ x_{K-1} \end{pmatrix} = \begin{pmatrix} b_0 \\ b_1 \\ b_2 \\ \vdots \\ b_{K-1} \end{pmatrix} \end{align} where $E_i, F_i \in \mathbb{R}^{n \times n}$, and $x_i, b_i \in \mathbb{R}^n$ for $i = 0, \ldots, K-1$

Further, the $E_i$ are invertible for $i = 0, \ldots, K-1$.

Then this system can be solved through forward substitution:

Solve $E_0 x_0 = b_0$

for $i = 1, \ldots, K-1$: Solve $E_i x_i = b_i - F_i x_{i-1}$

My Current Implementation

The block matrices $E_i$ and $F_i$ are available by calling Ek(i) and Fk(i).

Currently $x$ and $b$ are shaped as a numpy arrays with shape $K \times n$ so that x[k] gives $x_k$, and so forth.

import numpy as np
from scipy.sparse.linalg import spsolve

# define K and n, create b and initialize x

x[0] = spsolve(Ek(0), b[0])
for i in range(1, K):
    x[i] = spsolve(Ek(i), b[i] - Fk(i) @ x[i-1])    

Can this be vectorized? I would like to not use a for loop here since they are quite slow in Python.

Answer 1

In your explanation, you solve the large problem using forward substitution. This implies that you are solving your large problem successively: you first need $x_{i-1}$ before you can solve for $x_{i}$. This means that you need to loop, there is no way to avoid this.

However, for each subproblem, you use the spsolve routine, dedicated to solving sparse systems. Why? Are the submatrices $E_{i}$ really sparse? What is the order of magnitude of n? What sparse storage format do you use for the $E$ and $F$ matrices?

Another approach would be to consider your full matrix as a banded matrix and use scipy.linalg.solve_banded. If you have numpy linked to an optimized BLAS/LAPACK like MKL, this might also prove to be quite fast.

Could you elaborate a bit more on your problem (sizes, structure of the $E$ and $F$ matrices)? And if you feel up to it, make the comparison between the forward substitution approach (using a for loop) and the banded matrix approach and post it here...

Answer 2

I think the way to go is to use some library code at that point. Vectorization is quite a low-level (hardware near) optimization which seems a bit unnatural to tackle via python. Search for Lapack or BLAS vectorized operations. including them into your code might end up in a couple of extra lines, but will potentially spare you the misery of doing it by hand.scipy.linalg.lapack)

Also, before you think about vectorization, consider whether there are operations which you can run thread-parallel. Depending on your machine, that might already give you significant performance increase.

I'm no expert on performance in python, is there a chance of calling your substitution recursively might give a performance boost?