Numpy: Quick setting of array elements

I don't think it is the indexing with numbers that's slowing it down. Usually indexing a single dimension with a boolean vector is slower than indexing with the corresponding np.where.

Something else is going on here. Look at these shapes:

In [14]: a.shape
Out[14]: (10, 100, 100)
In [15]: np.argmax(a,axis=0).shape
Out[15]: (100, 100)
In [16]: a[np.argmax(a,axis=0)].shape
Out[16]: (100, 100, 100, 100)

The indexed a is much larger than the original, 1000x.

@MSeifert's solution is faster, but I can't help feeling it is more complex than needed.

In [35]: %%timeit 
   ....: a=np.arange(x*y*z).reshape((z,y,x))
   ....: pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
   ....: pos_max_indices = np.arange(a.shape[0]).reshape(10,1,1) == pos_max
   ....: a[pos_max_indices]=1
   ....: 
1000 loops, best of 3: 1.28 ms per loop

I'm still working on an improvement.

The sample array isn't a good one - it's too big to display, and all the max values on the last z plane:

In [46]: x,y,z=4,2,3
In [47]: a=np.arange(x*y*z).reshape((z,y,x))
In [48]: a
Out[48]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])

In [49]: a[np.argmax(a,axis=0)]=1
In [50]: a
Out[50]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[ 1,  1,  1,  1],
        [ 1,  1,  1,  1]]])

I could access those same argmax values with:

In [51]: a[-1,...]
Out[51]: 
array([[1, 1, 1, 1],
       [1, 1, 1, 1]])

Let's try a random array, where the argmax can be in any plane:

In [57]: a=np.random.randint(2,10,(z,y,x))
In [58]: a
Out[58]: 
array([[[9, 7, 6, 5],
        [6, 3, 5, 2]],

       [[5, 6, 2, 3],
        [7, 9, 6, 9]],

       [[7, 7, 8, 9],
        [2, 4, 9, 7]]])
In [59]: a[np.argmax(a,axis=0)]=0
In [60]: a
Out[60]: 
array([[[0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0]]])

Oops - I turned everything to 0. Is that what you want?

Let's try the pos_max method:

In [61]: a=np.random.randint(0,10,(z,y,x))
In [62]: a
Out[62]: 
array([[[9, 3, 9, 0],
        [6, 6, 2, 4]],

       [[9, 9, 4, 9],
        [5, 9, 7, 9]],

       [[1, 8, 1, 7],
        [1, 0, 2, 3]]])
In [63]: pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
In [64]: pos_max
Out[64]: 
array([[[0, 1, 0, 1],
        [0, 1, 1, 1]]], dtype=int32)
In [66]: pos_max_indices = np.arange(a.shape[0]).reshape(z,1,1) == pos_max
In [67]: pos_max_indices
Out[67]: 
array([[[ True, False,  True, False],
        [ True, False, False, False]],

       [[False,  True, False,  True],
        [False,  True,  True,  True]],

       [[False, False, False, False],
        [False, False, False, False]]], dtype=bool)
In [68]: a[pos_max_indices]=0
In [69]: a
Out[69]: 
array([[[0, 3, 0, 0],
        [0, 6, 2, 4]],

       [[9, 0, 4, 0],
        [5, 0, 0, 0]],

       [[1, 8, 1, 7],
        [1, 0, 2, 3]]])

That looks more reasonable. There still is a 9 in the 2nd plane, but that's because there was also a 9 in the 1st.

This still needs to be cleaned up, but here's a non-boolean mask solution:

In [98]: a=np.random.randint(0,10,(z,y,x))
In [99]: a1=a.reshape(z,-1)   # it's easier to work with a 2d view
In [100]: ind=np.argmax(a1,axis=0)
In [101]: ind
Out[101]: array([2, 2, 1, 0, 2, 0, 1, 2], dtype=int32)
In [102]: a1[ind,np.arange(a1.shape[1])]  # the largest values
Out[102]: array([9, 8, 7, 4, 9, 7, 9, 6])
In [104]: a1
Out[104]: 
array([[3, 1, 5, 4, 2, 7, 4, 5],
       [4, 4, 7, 1, 3, 7, 9, 4],
       [9, 8, 3, 3, 9, 1, 2, 6]])

In [105]: a1[ind,np.arange(a1.shape[1])]=0
In [106]: a
Out[106]: 
array([[[3, 1, 5, 0],
        [2, 0, 4, 5]],

       [[4, 4, 0, 1],
        [3, 7, 0, 4]],

       [[0, 0, 3, 3],
        [0, 1, 2, 0]]])

Working with a1 the 2d view is easier; the exact shape of the x,y dimensions is not important to this problem. We are changing individual values, not columns or planes. Still I'd like to do get it working without `a1.

Here are two functions that replace the maximum value (in the 1st plane). I use copy since it makes repeated time testing easier.

def setmax0(a, value=-1):
    # @MSeifert's
    a = a.copy()
    z = a.shape[0]
    # a=np.arange(x*y*z).reshape((z,y,x))
    pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
    pos_max_indices = np.arange(z).reshape(z,1,1) == pos_max
    a[pos_max_indices]=value
    return a

def setmax1(a, value=-2):
    a = a.copy()
    z = a.shape[0]
    a1 = a.reshape(z, -1)
    ind = np.argmax(a1, axis=0)
    a1[ind, np.arange(a1.shape[1])] = value
    return a

They produce the same result in a test like:

ab = np.random.randint(0,100,(20,1000,1000))
test = np.allclose(setmax1(ab,-1),setmax0(ab,-1))

Timings (using ipython timeit) are basically the same.

They do assign values in different orders, so setmax0(ab,-np.arange(...)) will be different.