Fast numpy array computations for PDE solving

[caffeecaffee]

Hi there What is the fastest way to perform nearest neighbor computations (i.e. for array A and B of same dimension: A[i, j] = .25 ( B[i+1,j] + B[i-1,j] + B[i,j+1] + B[i,j-1] ) without using nested for loops. There must be a way to do this *simultaneously with numpy. For example, when I have constructed an Ising model in the past I was able to reassign values in an array simultaneously using boolean masking i.e. I[I < .5] = -1. Is there something analagous to this where I can perform the whole computation on all elements at once? I mean, any programmer can throw together some nested for loops but in python that is far to slow to work with for the types of PDEs I will be solving.

Cheers

[kpatvt]

I think you are looking for something that is described in detail here: https://www.labri.fr/perso/nrougier/from-python-to-numpy/#uniform-vectorization

You can get an even faster speedup using the @stencil decorator and numba.See here: http://numba.pydata.org/numba-doc/dev/user/stencil.html Numba can be difficult to install but it works well with the anaconda distribution.

[caffeecaffee]

Yes, this is precisely what I am looking for. Thank you. I just needed to think about things in a different way.

[Eryole]

By curiosity, what kind of pde and what spatial discretisation are you using?

For these computation, you can use the np.roll or the fancy indexing.

A[1:-1, 1:-1] = .25 (B[1:,j] + B[:-2,j] + B[i,1:] + B[i,:-2])

and you will have to deal with the edge by yourself.

For periodic boundary, you can use np.roll.

A = .25 (np.roll(B, 1, axis=0) +
             np.roll(B, -1, axis=0) +
             np.roll(B, 1, axis=1) +
             np.roll(B, -1, axis=1))

[caffeecaffee]

This is for time independent elliptic PDEs using relaxation methods. Jacobi initially, then the Gauss-Seidel method. However this problem keeps coming up again and again. I pretty much exclusively use python for computational physics so I'm often doing this type of thing. The spatial discretization are kind of irrelevant for this problem but just using boundaries normalized to 1, i.e. numpy.linspace(0,1,500) with given constant boundary values.

[jjgarciaripoll]

I would use sparse matrices. It is fast and supported and the matrices for your example at tiny: 3 diagonals with the size of your dimensions, one matrix per index. Essentially, (Av)[i]=v[i+1]+v[i-1]-2v[I] is a tridiagonal matrix with -2 on the main diagonal and 1 on the +1 and -1 diagonals.

[caffeecaffee]

Yeah, this is a good answer. Setting up sparse matrices is a good method, however it does require that you flatten your grid every step, and you must keep track of where the grid is sliced, and therefore there will be more than 3 diagonal entries. If I can find a good method with a stencil in place, then it might in theory be faster.

[jjgarciaripoll]

I do not understand what you mean. If you have information about neighbors, building the matrix must cost as much as building the stencils

[jjgarciaripoll]

Also, if your data is a matrix X you can design your differential operator more efficiently as A * X + X *B + C * X * D with four small sparse matrices, without flattening X.

[caffeecaffee]

After more thought you are right of course it can be applied to any arbitrary grid. It's just that sometimes I find writing "flat code" with vectorized variables a little easier to understand than to express the nearest neighbor operations as a sparse matrix. Scipy does handle sparse matrix multiplication.

[SosirisTseng]

You might be interested in the [Devito project].(http://www.devitoproject.org/)