1

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.

Improve this question
2
  • Basically, given A, B and C, you are looking for a vectorized (no explicit looping) approach to compute D. Is that correct? Commented Feb 19, 2020 at 14:35
  • Yes, that's exactly what I'm looking for! Commented Feb 19, 2020 at 15:18

2 Answers 2

Reset to default
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.

Improve this answer
Sign up to request clarification or add additional context in comments.

3 Comments

I've seen the "take_along_axis"-function, but couldn't try it out. My numpy version seems too old (and I don't dare upgrade it due to incompatibility risks). Thanks for trying it out for me and timing it, though!
Which version of numpy do you have? Is np.apply_along_axis() / np.take() available to you?
My numpy version is 1.14.5 - I can call both np.apply_along_axis() and np.take(). I've come across both functions, but couldn't wrap my head around how to use them here.
0 
y, x = np.ogrid[:C.shape[0],:C.shape[1]]
D = A[C[y, x], y, x]
Improve this answer

2 Comments

This is exactly what I've been looking for, even if it might not be the fastest version. Thanks alot!
You're welcome. Readability matters as much or more than efficiency, in my experience.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.