Speeding up a linear transform using Python

Question

In an optical wavefront propagation problem, I need to do excessive Fourier-type computations:

import numpy as np

N=10000
input=np.random.random(N)

x=np.linspace(-1,1,N)
y=np.linspace(-2,2,N)
X,Y=np.meshgrid(x,y,sparse=True)

output=np.dot(np.exp(1j*2*np.pi*X*Y),input)

However, the true kernel (here, the complex exponential) is more complicated. In particular, it is not periodic, so no FFT-type reduction is possible! So, briefly, I need to evaluate trigonometric functions on a very large grid and then do a matrix-vector multiplication.

Doing this with pure numpy takes more than a week for one execution of my computations on a middle class desktop machine (with eight cores). Using numexpr reduces the runtime by a factor of roughly 1/3. Employing parallel Cython instead further reduces runtime by a factor of about 8/10.

Nevertheless, this is still too slow. I'm not an expert on scientific computing at all. So, for me, it is hard to judge how to speed this up further. Hence, my questions are:

1. Do there exist other software frameworks (numba, weave, OpenCL, ...) that are better suited to this problem?

2. What kind of hardware upgrade would make sense? (Obviously, the main computational costs are the function evaluations.)

EDIT: The continuous problem is a generalized one-dimensional Fresnel propagator

$$ u(x)= \int_{-\infty}^\infty \exp\left(i\pi\frac{(x-y)^2}{\lambda z(x,y)}\right) \frac{u_0(y)}{\sqrt{i\lambda z(x,y)}} dy,$$ where $\lambda$ is the wavelength of light and z is a positive function which is smooth, but more or less arbitrary

Answer 1

Assuming that your kernel is somewhat smooth, use low-rank approximation.

Here's a naive example:

import numpy as np

N=2000
input=np.random.random(N)

x=np.linspace(-1,1,N)
y=np.linspace(-2,2,N)
X,Y=np.meshgrid(x,y,sparse=True)
A = np.exp(1j*2*np.pi*X*Y)

output = np.dot(A, input)

U,S,V = np.linalg.svd(A)

# find truncation rank for given tolerance
k = np.nonzero(S < 1e-12)[0][0]
print('Truncation rank:', k)
Xk = U[:,:k].dot(np.diag(S[:k]))
Yk = V[:k,:]

output2 = Xk.dot(Yk.dot(input))

# check quality of approximation
print(np.linalg.norm(output - output2, np.inf))

%timeit np.dot(A, input)
%timeit Xk.dot(Yk.dot(input))

Output:

Truncation rank: 32
6.3535593536e-12
100 loops, best of 3: 3.9 ms per loop
10000 loops, best of 3: 61.4 µs per loop

What the above does is compute a truncated singular value decomposition of your matrix and use the resulting low-rank approximation of A to speed up the matrix-vector product. So, instead of using $A$, you use an approximation $$ A \approx X_k Y_k^T, \qquad X_k,Y_k \in \mathbb R^{N\times k} $$ with $k \ll N$.

In the example, this leads to a speedup of more than 60x, and this will only get better as N gets larger. That's because the full matrix-vector product has complexity $\mathcal O(N^2)$, whereas the low-rank matrix-vector product has complexity $\mathcal O(Nk)$. The result is not exact, but you can influence the accuracy by playing with the truncation rank.

Of course, computing the SVD is very slow. Instead, there are faster alternatives for computing low-rank approximations to a matrix, for instance: fast randomized SVD or, what I would recommend, Adaptive Cross Approximation (ACA). The latter is a black-box algorithm for computing a low-rank approximation without ever even computing the entire matrix A and thus is very well suited for very large problems. It's also easy to implement.

The original literature on ACA (in particular, by M. Bebendorf from 2000 on) is somewhat inaccessible. Instead, here's a random paper which describes ACA in a basic form: (link).

If your kernel is relatively simple, it might even be possible to derive a separable approximation to it analytically, which directly leads to a low-rank approximation instead of having to compute it algorithmically. Check the literature on "multipole expansion" and "fast multipole method" for many examples of this.

Finally, the bottleneck in your original program was actually computing A, not the matrix-vector product. ACA will eliminate that, since it doesn't need the full matrix.

UPDATE:

I just realized that scipy already includes a fast method for approximating the SVD in the little-known scipy.linalg.interpolative module. Here's an example:

import scipy.linalg.interpolative as sli

U,S,V = sli.svd(A, 1e-12)
print('Rank:', U.shape[1])

Xk = U.dot(np.diag(S))
Yk = V.conj().T

output3 = Xk.dot(Yk.dot(input))

# check quality of approximation
print(np.linalg.norm(output - output3, np.inf))

And the output:

Rank: 37
2.62815545927e-10

The difference in timing between exact SVD and approximate SVD (again, this will get even better with larger N due to the involved asymptotics):

%timeit np.linalg.svd(A)
%timeit sli.svd(A, 1e-12)

1 loop, best of 3: 11.4 s per loop
1 loop, best of 3: 297 ms per loop

However, this still requires you to compute the entire matrix A first and doesn't get rid of that bottleneck. Nevertheless, if you need to compute A only once and then have many matrix-vector products, it may be a suitable option.

Answer 2

Matrix-vector multiplication is RAM-bound. You make two arithmetic operations for each float that you read. So the major cost is going to be storing the whole matrix in slow, non-cached core memory and reading it back.

You should consider evaluating your "complicated kernel" on the fly instead of saving it to memory:

output = np.zeros(n)
for j in range(n):
  for i in range(n):
    output[i] += kernel(i, j) * input[j]

(Warning: this loop is going to be extremely slow until you cythonize it).

If you have to make a single matrix-vector product with each matrix, like in your example, this is likely going to be a major speedup.

If you have to make many products with the same matrix, then it depends: try both approaches and compare them. (Also, if you have all the vectors available at the same time, store them in a matrix and turn this operation into a matrix-matrix multiplication). In the hardware department, the first thing to check is to make sure that you have enough RAM to store the whole matrix -- if you have to hit the disk, then things are going to get s..l..o..w.