How to access index value of the axis in question with numpy apply_along_axis()

Question

I have a problem that can be fully vectorised, but I don't have enough space, so I am trying for a half and half solution, using numpy's apply_along_axis().

(NOTE: This is a toy example that illustrates the nub of the problem. In other words, I am not looking for a numpy or scipy function that does what the function here is doing - it is not the real function, just one that is simple to illustrate with.)

What I would like to do is figure out a way to access the index of the axis being passed at each iteration.

Say we took a 4 x 4 matrix:

    M = np.array(([0,0,1,1], [1,1,0,1], [1,0,1,0], [0,0,1,1]))
    M 
   
   array([[0, 0, 1, 1],
          [1, 1, 0, 1],
          [1, 0, 1, 0],
          [0, 0, 1, 1]])

And wanted to compute the pairwise bitwise logical and of each column against every other column, but to save (a lot of) time, we only compute for columns i,j where index of j > i (so that we end up with a triangular matrix).

In pandas I can do this quite easily with apply(), but it is too slow for my purposes.

I know there are pairwise functions in scikit-learn, but please assume these do not suit my purposes (my function is more complicated than this toy one)

If I were to use numpy's apply_long_axis(), I can only work out how to compare all i,j and j,i, not the smaller problem as described before.

This is my solution:

def intersections_np(col, M):
    col = col[:,np.newaxis]
    intersection = (M & col).sum(0)
    return(intersection)

result_np = np.apply_along_axis(intersections_np, arr = M, axis = 0,  M = M)
result_np

array([[2, 1, 1, 1],
       [1, 1, 0, 1],
       [1, 0, 3, 2],
       [1, 1, 2, 3]], dtype=int32)

But what I would really like to do is something like:

def intersections_np(col, M):
    col = col[:,np.newaxis]
    start_index =  + 1
    other_cols = M[:,start_index:]
    intersection = (other_cols & col).sum(0)
    
    return(intersection)

result_np = np.apply_along_axis(intersections_np, arr = M, axis = 0,  M = M)

and return:

result_np

array([[nan, nan, nan, nan],
       [1, nan, nan, nan],
       [1, 0, nan, nan],
       [1, 1, 2, nan]], dtype=int32)

Does any one know if something like this can be done?

Thanks

user_901925 · Answer

Let's do a bit of timing.

Your basic apply:

In [142]: timeit np.apply_along_axis(intersections_np, arr = M, axis = 0,  M = M)                    
158 µs ± 3.97 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

And equivalent iteration (technically it may require a transpose, the the result is symmetric, so that doesn't matter):

In [143]: timeit np.array([intersections_np(M[:,i],M) for i in range(M.shape[1])])                   
65.4 µs ± 1.93 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

and @jfahne suggestion:

In [144]: %%timeit  
     ...: np.reshape(np.array([(M.T[i] & M.T[j]).sum(0) if j>i else 0 \ 
     ...: for i in range(len(M.T)) for j in range(len(M.T))]),(M.T).shape).T 
     ...:  
     ...:                                                                                            
95.2 µs ± 2.99 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Notice that apply is slower than plain iteration. That's consistent with my past testing. apply is only a help if the array is 3d or more, and iteration is an 'ugly' double nested one. There it's prettier, though still not faster. It's a convenience tool, not a speed one.

A fully 'vectorized' solution (with numpy broadcasting etc):

In [148]: (M[:,:,None] & M[:,None,:]).sum(0)                                                         
Out[148]: 
array([[2, 1, 1, 1],
       [1, 1, 0, 1],
       [1, 0, 3, 2],
       [1, 1, 2, 3]])
In [149]: timeit (M[:,:,None] & M[:,None,:]).sum(0)                                                  
14.9 µs ± 182 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

It does make an intermediate (4,4,4) array, and doesn't avoid the duplicates, but because there's no iteration at a Python level, it is quite fast. Trying to restrict the calculations to the lower (or upper) triangle isn't usually worth it.

But if you really want the lower triangle, and speed, consider using numba. It can be quite fast for iterative problems (but at some cost in flexibility).

Here's a version of your intersections restricted to the lower triangle

In [159]: def foo(M): 
     ...:     m = M.shape[0] 
     ...:     res = np.full((m,m), np.nan) 
     ...:     for i in range(m-1): 
     ...:         temp = (M[:,i,None] & M[:,(i+1):]).sum(0) 
     ...:         res[-temp.shape[0]:,i] = temp 
     ...:     return res 
     ...:      
     ...:                                                                                            
In [160]: foo(M)                                                                                     
Out[160]: 
array([[nan, nan, nan, nan],
       [ 1., nan, nan, nan],
       [ 1.,  0., nan, nan],
       [ 1.,  1.,  2., nan]])
In [161]: timeit foo(M)                                                                              
59.3 µs ± 2.42 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Basically the same timing as my [143] - it has fewer calculations in the & step, but more indexing, so there's little speed change.

user_14087478 · Answer

This has a rather nice pythonic solution:

np.reshape(np.array([(M.T[i] & M.T[j]).sum(0) if j>i else 0 \
for i in range(len(M.T)) for j in range(len(M.T))]),(M.T).shape).T

The use of M.T is to access the columns. The resultant vector is reshaped into the same shape as the transposed array. The array is then transposed back to the original array shape and yields the desired output.

How to access index value of the axis in question with numpy apply_along_axis()

2 Answers