Numpy apply along axis based on row index

Question

Trying to apply numpy inbuilt function apply_along_axis based on row index position

import numpy as np
sa = np.array(np.arange(4))
sa_changed = (np.repeat(sa.reshape(1,len(sa)),repeats=2,axis=0))
print (sa_changed)

OP:

[[0 1 2 3]
 [0 1 2 3]]

The function:

np.apply_along_axis(lambda x: x+10,0,sa_changed)

Op:

array([[10, 11, 12, 13],
       [10, 11, 12, 13]])

But is there a way to use this function based on row index position for example, if its a even row index then add 10 and if its a odd row index then add 50

Sample:

def func(x):
   if x.index//2==0:
      x = x+10
   else:
      x = x+50
   return x

`apply_along_axis` is not a speed tool. It can make code look prettier, especially if the array is 3d or larger, but it doesn't do things faster. — hpaulj, May 07 '21 at 03:09
@hpaulj Thanks for the solution. Do you have any other suggestions for faster approach? Not exact code, just suggestions would be fine. — data_person, May 07 '21 at 03:10

hpaulj · Accepted Answer · 2021-05-07T03:13:03.947

When iterating on array, directly or with apply_along_axis, the subarray does not have a .index attribute. So we have to pass an explicit index value to your function:

In [248]: def func(i,x):
     ...:    if i//2==0:
     ...:       x = x+10
     ...:    else:
     ...:       x = x+50
     ...:    return x
     ...: 
In [249]: arr = np.arange(10).reshape(5,2)

apply doesn't have a way to add this index, so instead we have to use an explicit iteration.

In [250]: np.array([func(i,v) for i,v in enumerate(arr)])
Out[250]: 
array([[10, 11],
       [12, 13],
       [54, 55],
       [56, 57],
       [58, 59]])

replacing // with %

In [251]: def func(i,x):
     ...:    if i%2==0:
     ...:       x = x+10
     ...:    else:
     ...:       x = x+50
     ...:    return x
     ...: 
In [252]: np.array([func(i,v) for i,v in enumerate(arr)])
Out[252]: 
array([[10, 11],
       [52, 53],
       [14, 15],
       [56, 57],
       [18, 19]])

But a better way is to skip the iteration entirely:

Make an array of the row additions:

In [253]: np.where(np.arange(5)%2,10,50)
Out[253]: array([50, 10, 50, 10, 50])

apply it via broadcasting:

In [256]: x+np.where(np.arange(5)%2,50,10)[:,None]
Out[256]: 
array([[10, 11],
       [52, 53],
       [14, 15],
       [56, 57],
       [18, 19]])

score 1 · Answer 2 · answered May 07 '21 at 03:14

here's one way to do this

import numpy as np

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

y = x.copy() # if you dont wish to modify x

for even row index

y[::2] = y[::2] + 10

and for odd row index

y[1::2] = y[1::2] + 50

output :

array([[10, 11, 12, 13],
       [50, 51, 52, 53]])

Numpy apply along axis based on row index

2 Answers2