3D numpy array threshold comparison

I’m trying to solve whether or not each element in a 3d numpy array fits within a certain threshold. Currently I have nested for loops that get the job done, but it’s really slow and I don’t have all day to wait for my code to run LOL. The input spec is a little bit weird so I haven’t been able to find any more efficient solutions. Here’s an example:

input matrix:

[
    [
        [4, 5, 2],
        [5, 6, 4]
    ],
    [
        [4, 4, 2],
        [5, 8, 0]
    ],
    [
        [8, 4, 6],
        [0, 6, 8]
    ]
]

upper threshold: 5

lower threshold:

[
    [1, 0, 4],
    [1, 2, 2]
]

and the output should be:

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

Answer

Broadcasting saves the day!

x = np.array([
    [
        [4, 5, 2],
        [5, 6, 4],
    ],
    [
        [4, 4, 2],
        [5, 8, 0],
    ],
    [
        [8, 4, 6],
        [0, 6, 8],
    ],
])

upper_thresh = 5

lower_thresh = np.array([
    [1, 0, 4],
    [1, 2, 2],
])

within_lower = x > lower_thresh
within_upper = x < upper_thresh
within_bound = (within_lower & within_upper).astype(dtype=np.int64)

And now:

>>> print(within_bound)
[[[1 0 0]
  [0 0 1]]

 [[1 1 0]
  [0 0 0]]

 [[0 1 0]
  [0 0 0]]]

(Not sure if you wanted inclusive thresholds. In that case, use <= and >= for comparisons instead.)