How to slice a 3D array from a 2D array without nested loops?

Question

I have four matrices, partially created from each other:

• A is a 3D matrix that represents a stack of grayscale images and is of shape (n, h, w)
• B is a 3D matrix that also represents a stack of images, where each slice is individually calculated from the corresponding slice in A and is also of shape (n, h, w)
• C is a 2D matrix, containing the index with the maximum value of B in z direction and is of shape (h, w)
• D is a 2D matrix, where a value from A is copied from a certain slice, which is indicated by the value in C at position (x, y).

A minimum example implemented with loops would look as follows:

import numpy as np

A = np.random.randint(0, 255, size=(3, 4, 5))
B = np.random.randint(0, 255, size=(3, 4, 5))
C = np.argmax(B, axis=0)

D = np.zeros(C.shape, dtype=int)
for y in range(C.shape[0]):
    for x in range(C.shape[1]):
        D[y, x] = A[C[y, x], y, x]


> A
array([[[ 24,  84, 155,   8, 147],
        [ 25,   4,  49, 195,  57],
        [ 93,  76, 233,  83, 177],
        [ 70, 211, 201, 132, 239]],

       [[177, 144, 247, 251, 207],
        [242, 148,  28,  40, 105],
        [181,  28, 132,  94, 196],
        [166, 121,  72,  14,  14]],

       [[ 55, 254, 140, 142,  14],
        [112,  28,  85, 112, 145],
        [ 16,  72,  16, 248, 179],
        [160, 235, 225,  14, 211]]])

> B
array([[[246,  14,  55, 163, 161],
        [  3, 152, 128, 104, 203],
        [ 43, 145,  59, 169, 242],
        [106, 169,  31, 222, 240]],

       [[ 41,  26, 239,  25,  65],
        [ 47, 252, 205, 210, 138],
        [194,  64, 135, 127, 101],
        [ 63, 208, 179, 137,  59]],

       [[112, 156, 183,  23, 253],
        [ 35,   6, 233,  42, 100],
        [ 66, 119, 102, 217,  64],
        [ 82,  67, 135,   6,   8]]])

> C
array([[0, 2, 1, 0, 2],
       [1, 1, 2, 1, 0],
       [1, 0, 1, 2, 0],
       [0, 1, 1, 0, 0]])

> D
array([[ 24, 254, 247,   8,  14],
       [242, 148,  85,  40,  57],
       [181,  76, 132, 248, 177],
       [ 70, 121,  72, 132, 239]])

My question is: How to slice A with C efficiently eliminating the nested for-loops? My initial idea was to expand C to a 3D boolean mask, where only the positions [c, y, x] are set to True and then to simply multiply it elementwise with A and take the sum over z-axis. But I can't think of an pythonesque implementation without loops (and I probably wouldn't need to create a boolean mask anymore, if I'd knew how to write that).

The closest already implemented function I found is np.choose(), but it only takes 32 elements for C.

Answer 1

The standard approach here is to use np.take_along_axis() in conjunction with np.expand_dims() (the core idea is presented also in the np.argmax() documentation):

np.take_along_axis(A, np.expand_dims(C, axis=0), axis=0).squeeze()

Comparing the proposed approach with the explicit loops and the np.ogrid()-based approaches one would get:

import numpy as np


def take_by_axis_loop_fix(arr, pos):
    result = np.zeros(pos.shape, dtype=int)
    for i in range(pos.shape[0]):
        for j in range(pos.shape[1]):
            result[i, j] = arr[pos[i, j], i, j]
    return result


def take_by_axis_ogrid_fix(arr, pos):
    i, j = np.ogrid[:pos.shape[0], :pos.shape[1]]
    return arr[pos[i, j], i, j]


def take_by_axis_np(arr, pos, axis=0):
    return np.take_along_axis(arr, np.expand_dims(pos, axis=axis), axis=axis).squeeze()


def take_by_axis_ogrid(arr, pos, axis=0):
    ij = tuple(np.ogrid[tuple(slice(None, d, None) for d in pos.shape)])
    ij = ij[:axis] + (pos[ij],) + ij[axis:]
    return arr[ij]

A_ = np.random.randint(0, 255, size=(300, 400, 500))
B_ = np.random.randint(0, 255, size=(300, 400, 500))
C_ = np.argmax(B_, axis=0)

funcs = take_by_axis_loop_fix, take_by_axis_ogrid_fix, take_by_axis_ogrid, take_by_axis_np
for func in funcs:
    print(func.__name__, np.all(func(A_, C_) == take_by_axis_loop_fix(A_, C_)))
    %timeit func(A_, C_)
    print()

# take_by_axis_loop_fix True
# 10 loops, best of 3: 114 ms per loop

# take_by_axis_ogrid_fix True
# 100 loops, best of 3: 5.94 ms per loop

# take_by_axis_ogrid True
# 100 loops, best of 3: 5.54 ms per loop

# take_by_axis_np True
# 100 loops, best of 3: 3.34 ms per loop

indicating this to be the most efficient approach proposed so far. Note also that the np.take_along_axis()-based and the take_by_axis_ogrid() approaches would work essentially unchanged for inputs with higher dimensionality, contrarily to the _fix approaches.

Particularly, take_by_axis_ogrid() is the axis-agnostic version of take_by_axis_ogrid_fix() which is, essentially, nth's answer.

Answer 2

y, x = np.ogrid[:C.shape[0],:C.shape[1]]
D = A[C[y, x], y, x]