Select and apply condition on numpy n-dimensional inner array-CodePudding

Let's suppose I have a numpy array of shape (2, 4, 3), like this:

import numpy as np

arr = np.array([[[248,  26,   4],
                 [ 99, 126, 156],
                 [ 80, 240, 232],
                 [136,  27, 216]],

                [[221, 130, 119],
                 [253, 188, 232],
                 [159,  21,  98],
                 [ 12,  35,  50]]])

I'd like to check if each item in the smaller array meets a condition, for example:

conditions_list = [(arr[:,:,0] > 100 and arr[:,:,1] > 100 and arr[:,:,2] > 100),
                   (arr[:,:,0] < 25 and arr[:,:,1] < 50 and arr[:,:,2] < 50)]
choices_list = [(arr[:,:,0] = 255, arr[:,:,1] = 255, arr[:,:,2] = 255),
                (arr[:,:,0] = 0, arr[:,:,1] = 0, arr[:,:,2] = 0)]
new_array = np.select(conditions_list, choices, default=(arr[:,:,0], arr[:,:,1], arr[:,:,2]))
#This code doesn't work, but it represents the idea of what I need to get.

The expected result is as follows:

([[[248,  26,   4],
   [ 99, 126, 156],
   [ 80, 240, 232],
   [136,  27, 216]],

  [[255, 255, 255],
   [255, 255, 255],
   [159,  21,  98],
   [  0,   0,   0]]])

When I run the code above, I get an exception saying that I should use arr.any() or arr.all() instead of conditions_list, but those functions do not meet what I need.

How can an inner array be modified based on conditions, as efficiently as possible? There can be thousands of arrays that can be shaped up to (3000, 3000, 3), which is why I thought np.select() was a good option.

EDIT

I know I can use a list comprehension defining a custom function, but that is too slow as it'll iterate over each item. It can be done like this:

new_arr = [[cond_func(x) for x in y] for y in arr]

def cond_func(x):
    if x[0] > 100 and x[1] > 100 and x[2] > 100:
        x[0], x[1], x[2] = 255, 255, 255
    elif x[0] < 25 and x[1] < 50 and x[2] < 50:
        x[0], x[1], x[2] = 0, 0 ,0
    return(x)

CodePudding user response：

In [78]: arr = np.array([[[248,  26,   4],
    ...:                  [ 99, 126, 156],
    ...:                  [ 80, 240, 232],
    ...:                  [136,  27, 216]],
    ...: 
    ...:                 [[221, 130, 119],
    ...:                  [253, 188, 232],
    ...:                  [159,  21,  98],
    ...:                  [ 12,  35,  50]]])
    ...:

Here's a 'masking' way of handling the >100 case:

In [79]: (arr>100)
Out[79]: 
array([[[ True, False, False],
        [False,  True,  True],
        [False,  True,  True],
        [ True, False,  True]],

       [[ True,  True,  True],
        [ True,  True,  True],
        [ True, False, False],
        [False, False, False]]])
In [80]: (arr>100).all(axis=2)
Out[80]: 
array([[False, False, False, False],
       [ True,  True, False, False]])
In [81]: arr[_]
Out[81]: 
array([[221, 130, 119],
       [253, 188, 232]])
In [82]: arr1=arr.copy()
In [83]: arr1[(arr>100).all(axis=2)]=255
In [84]: arr1
Out[84]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [ 12,  35,  50]]])

and for the 2nd condition, using 75 for arr[:,:,2]:

In [94]: arr1[(arr<np.array([25,50,75])).all(axis=2)]=0
In [95]: arr1
Out[95]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [  0,   0,   0]]])

Here's a way of using these masks with select. Note that I have to adjust the dimensions of the masks so they work with the 3d default.

In [104]: mask1 = (arr>100).all(axis=2)
In [105]: mask2 = (arr<np.array([25,50,75])).all(axis=2)
In [106]: mask1.shape
Out[106]: (2, 4)
In [107]: arr.shape
Out[107]: (2, 4, 3)
In [108]: np.select([mask1[:,:,None],mask2[:,:,None]],[255,0], default=arr)
Out[108]: 
array([[[248,  26,   4],
        [ 99, 126, 156],
        [ 80, 240, 232],
        [136,  27, 216]],

       [[255, 255, 255],
        [255, 255, 255],
        [159,  21,  98],
        [  0,   0,   0]]])

Or use:

In [109]: mask1 = (arr>100).all(axis=2, keepdims=True)
In [110]: mask1.shape
Out[110]: (2, 4, 1)

To use your syntax, we need to add () and use & instead of and:

In [111]: (arr[:,:,0] > 100) & (arr[:,:,1] > 100) & (arr[:,:,2] > 100)
Out[111]: 
array([[False, False, False, False],
       [ True,  True, False, False]])