Pythonic algorithm for applying a function to a 3D array along the first axis indices

Question

I have a 3D array like A

A = np.random.randint(20,size=(4,2,2))
array([[[18,  8],
    [ 2, 11]],

   [[ 9,  8],
    [ 9, 10]],

   [[ 0,  1],
    [10,  6]],

   [[ 1,  8],
    [ 4,  2]]])

What I want to do is to apply a function to some indices along the axis=0. For example, I want to multiply the A[1] and A[3] by 2 and add 10 to them. I know one option is this:

for index in [1,3]:
    A[index] = A[index]*2+10

Which gives:

array([[[18,  8],
        [ 2, 11]],

       [[28, 26],
        [28, 30]],

       [[ 0,  1],
        [10,  6]],

       [[12, 26],
        [18, 14]]])

But my original array is of the size of (2500, 300, 300) and I need to apply the function to 500 non-consecutive indices along the axis=0 every time. Is there a faster and more pythonic way to do it?

No, they are no regularly spaced. They are like 5 sets of 500 random integers which will be used for applying 5 different functions to 2500 indices in the first axis. @unutbu — Monobakht
– Monobakht, Commented Jun 1, 2017 at 22:26

piRSquared · Accepted Answer · 2017-06-01 23:06:14Z

2

You could use stepped slicing

A[1::2] = A[1::2] * 2 + 10
A

array([[[18,  8],
        [ 2, 11]],

       [[28, 26],
        [28, 30]],

       [[ 0,  1],
        [10,  6]],

       [[12, 26],
        [18, 14]]])

Or assuming your slice is named slc

A[slc] = A[slc] * 2 + 10

edited Jun 1, 2017 at 23:06

answered Jun 1, 2017 at 22:24

piRSquared

296k68 gold badges509 silver badges654 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Monobakht Over a year ago

As I mentioned in a comment above, my indices are not regularly spaced. So I can not use stepped slicing. But thanks ;)

piRSquared Over a year ago

@Monobakht so do you know what the indices will be when you are applying the function?

Monobakht Over a year ago

Yes, I do have a list of 500 indices in my case.

Monobakht Over a year ago

That is what I was looking for. Thanks a lot.

Collectives™ on Stack Overflow

Pythonic algorithm for applying a function to a 3D array along the first axis indices

1 Answer 1

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related