CODEWITHSUNDEEP
pythonnumpy
slaw
slaw

Reputation: 6869

Numpy Apply Along Axis and Get Row Index

I have a 2D array (it is actually very large and a view of another array):

x = np.array([[0, 1, 2],
          [1, 2, 3],
          [2, 3, 4],
          [3, 4, 5]]
        )

And I have a function that processes each row of the array:

def some_func(a):
    """
    Some function that does something funky with a row of numbers
    """
    return [a[2], a[0]]  # This is not so funky

np.apply_along_axis(some_func, 1, x)

What I'm looking for is some way to call the np.apply_along_axis function so that I have access to the row index (for the row being processed) and then be able to process each row with this function:

def some_func(a, idx):
    """
    I plan to use the index for some logic on which columns to
    return. This is only an example
    """
    return [idx, a[2], a[0]]  # This is not so funky

Upvotes: 13

Views: 6614

Answers (3)

Nathaniel Mahowald
Nathaniel Mahowald

Reputation: 21

I had this problem for a 3 dimensional tensor, so I thought it'd be worthwhile to post a solution that generalizes, which was to use np.ndenumerate.

    f = lambda indices: #(whatever you'd like to do)

    output = np.empty(M.shape)
    for i, x in np.ndenumerate(M):
        output(i) = f(i)

Upvotes: 2

nanangarsyad
nanangarsyad

Reputation: 700

Here's an alternative solution, while waiting for the real feature got implemented. This is gonna be a bit untidy. But enough to tackle your problem at the moment, maybe. :)

# create global variable
In [247]: global counter  

# Initialize it to your need
In [248]: counter = 0 

# create your function callback, lambda also could be used
In [252]: def funct(row): 
     ...:     # reference to global variable created before hand 
     ...:     global counter   
     ...:     counter += 1 # increment the counter
     ...:     # return something, or else 
     ...:     # will raise a 'NoneType' has no len() exception
     ...:     return counter

In [260]: np.apply_along_axis(funct, 1, np.array([[0],[0],[0]]))
Out[260]: array([1, 2, 3])

# revert counter to initial state or the counter will keep raising afterward
In [261]: counter = 0 

# or you could just delete it if you had no use for it anymore
In [262]: del counter

Hoped, could be some help for you :)

Upvotes: 0

hpaulj
hpaulj

Reputation: 231395

For a 2d array with axis=1, apply_along_axis is the same as iteration of the rows of the array

In [149]: np.apply_along_axis(some_func, 1, x)
Out[149]: 
array([[2, 0],
       [3, 1],
       [4, 2],
       [5, 3]])
In [151]: np.array([some_func(i) for i in x])
Out[151]: 
array([[2, 0],
       [3, 1],
       [4, 2],
       [5, 3]])

For axis=0 we could iterate on x.T. apply_along_axis is more useful when the array is 3d, and we want to iterate on all dimensions except one. Then it saves on some tedium. But it isn't a speed solution.

With your revised function, we can use standard enumerate to get both row and index:

In [153]: np.array([some_func(v,i) for i,v in enumerate(x)])
Out[153]: 
array([[0, 2, 0],
       [1, 3, 1],
       [2, 4, 2],
       [3, 5, 3]])

or with a simple range iteration:

In [157]: np.array([some_func(x[i],i) for i in range(x.shape[0])])
Out[157]: 
array([[0, 2, 0],
       [1, 3, 1],
       [2, 4, 2],
       [3, 5, 3]])

There are various tools for getting the indexing for higher dimenions, things like ndenumerate and ndindex.

The fast solution - work on all rows at once:

In [158]: np.column_stack((np.arange(4), x[:,2], x[:,0]))
Out[158]: 
array([[0, 2, 0],
       [1, 3, 1],
       [2, 4, 2],
       [3, 5, 3]])

Upvotes: 1

Related Questions