Select and apply condition on numpy n-dimensional inner array

Question

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)

`and` is like `if`. It only works with scalar values. `arr[:,:,0] > 100` is a boolean array. If cannot be used with `if` or `and` (or `or`). — hpaulj, Oct 01 '21 at 19:34
Another issue, `arr[:,:,0] = 255` is an assignment, not a equality test. — hpaulj, Oct 01 '21 at 19:36
@hpaulj I know the code above doesn't work, it is just as an illustration of how I would normally do it using `np.select()'. I need a way to get to the result that is efficient. — Jose Vega, Oct 01 '21 at 19:38
Before throwing all those expressions into `select`, make sure they work individually. `select` may still work, but it is a python function. So all arguments are evaluated before hand. If necessary, read the `select` docs carefully, and test small cases. Often when working with a new function it's a good idea to start with the documented example, and incrementally make things more complicated. In other words, start small, and work you way to a fuller implementation. — hpaulj, Oct 01 '21 at 19:40
Does this answer your question? [ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()](https://stackoverflow.com/questions/10062954/valueerror-the-truth-value-of-an-array-with-more-than-one-element-is-ambiguous) — jjramsey, Oct 01 '21 at 19:41

hpaulj · Accepted Answer · 2021-10-01T20:22:45.180

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]])

Select and apply condition on numpy n-dimensional inner array

1 Answers1