numpy np.apply_along_axis function speed up?

np.apply_along_axis is not for speed.

There is no way to apply a pure Python function to every element of a Numpy array without calling it that many times, short of AST rewriting…

Fortunately, there are solutions:

• Vectorizing

  Although this is often hard, it’s normally the easy solution. Find some way to express your calculation in a way that generalizes over the elements, so you can work on the whole matrix at once. This will result in the loops being hoisted out of Python and in to heavily optimised C and Fortran routines.

• JITing: Numba and Parakeet, and to a lesser extent PyPy with NumPyPy

  Numba and Parakeet both deal with JITing loops over Numpy data structures, so if you inline the looping into a function (this can be a wrapper function), you can get massive speed boosts for almost-free. This depends on the data structures used, though.

• Symbolic evaluators like Theano and numexpr

  These allow you to use embedded languages to express calculations, which can end up much faster than even the vectorized versions.

• Cython and C extensions

  If all else is lost, you can always dig down manually to C. Cython hides a lot of the complexity and has a lot of lovely magic too, so it’s not always that bad (although it helps to know what you’re doing).

---

Here you go.

This is my testing “environment” (you should really have provided this :P):

import itertools
import numpy

a = numpy.arange(200).reshape((200,1)) ** 2

def my_func(a, i,j):
    b = numpy.zeros((2,2))
    b[0,0] = a[i]
    b[1,0] = a[i]
    b[0,1] = a[i]
    b[1,1] = a[j]
    return  numpy.linalg.eigh(b)

eigvals = {}
eigvecs = {}

for i, j in itertools.combinations(range(a.size), 2):
    eigvals[i, j], eigvecs[i, j] = my_func(a,i,j)

Now, it’s far easier to get all the permutations instead of the combinations, because you can just do this:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

This might seem wasteful, but there are only twice as many permutations so it’s not a big deal.

So we want to use these indexes to get the relevant elements:

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

and then we can make our matrices:

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

We need the matrices to be in the last two dimensions:

target.shape
#>>> (2, 2, 200, 200)

target = numpy.swapaxes(target, 0, 2)
target = numpy.swapaxes(target, 1, 3)

target.shape
#>>> (200, 200, 2, 2)

And we can check that it works:

target[10, 20]
#>>> array([[100, 100],
#>>>        [100, 400]])

Yay!

So then we just run numpy.linalg.eigh:

values, vectors = numpy.linalg.eigh(target)

And look, it works!

values[10, 20]
#>>> array([  69.72243623,  430.27756377])

eigvals[10, 20]
#>>> array([  69.72243623,  430.27756377])

So then I’d imagine you might want to concatenate these:

numpy.concatenate([values[row, row+1:] for row in range(len(values))])
#>>> array([[  0.00000000e+00,   1.00000000e+00],
#>>>        [  0.00000000e+00,   4.00000000e+00],
#>>>        [  0.00000000e+00,   9.00000000e+00],
#>>>        ..., 
#>>>        [  1.96997462e+02,   7.78160025e+04],
#>>>        [  3.93979696e+02,   7.80160203e+04],
#>>>        [  1.97997475e+02,   7.86070025e+04]])

numpy.concatenate([vectors[row, row+1:] for row in range(len(vectors))])
#>>> array([[[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        ..., 
#>>>        [[-0.70890372,  0.70530527],
#>>>         [ 0.70530527,  0.70890372]],
#>>> 
#>>>        [[-0.71070503,  0.70349013],
#>>>         [ 0.70349013,  0.71070503]],
#>>> 
#>>>        [[-0.70889463,  0.7053144 ],
#>>>         [ 0.7053144 ,  0.70889463]]])

It’s also possible to do this concatenate loop just after numpy.mgrid to halve the amount of work:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

# Convert to all *combinations* and reduce the dimensionality
indexes = numpy.concatenate([indexes[:, row, row+1:] for row in range(indexes.shape[1])], axis=1)

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

target = numpy.rollaxis(target, 2)

values, vectors = numpy.linalg.eigh(target)

Yeah, that last sample is all you need.