Summarize ndarray by 2d array in Python

Question

I want to summarize a 3d array dat using indices contained in a 2d array idx.

Consider the example below. For each margin along dat[:, :, i], I want to compute the median according to some index idx. The desired output (out) is a 2d array, whose rows record the index and columns record the margin. The following code works but is not very efficient. Any suggestions?

import numpy as np
dat = np.arange(12).reshape(2, 2, 3)
idx = np.array([[0, 0], [1, 2]])

out = np.empty((3, 3))
for i in np.unique(idx):
    out[i,] = np.median(dat[idx==i], axis = 0)
print(out)

Output:

[[ 1.5  2.5  3.5]
 [ 6.   7.   8. ]
 [ 9.  10.  11. ]]

Answer 1

To visualize the problem better, I will refer to the 2x2 dimensions of the array as the rows and columns, and the 3 dimension as depth. I will refer to vectors along the 3rd dimension as "pixels" (pixels have length 3), and planes along the first two dimensions as "channels".

Your loop is accumulating a set of pixels selected by the mask idx == i, and taking the median of each channel within that set. The result is an Nx3 array, where N is the number of distinct incides that you have.

One day, generalized ufuncs will be ubiquitous in numpy, and np.median will be such a function. On that day, you will be able to use reduceat magic¹ to do something like

unq, ind = np.unique(idx, return_inverse=True)
np.median.reduceat(dat.reshape(-1, dat.shape[-1]), np.r_[0, np.where(np.diff(unq[ind]))[0]+1])

_{¹ See Applying operation to unevenly split portions of numpy array for more info on the specific type of magic.}

Since this is not currently possible, you can use scipy.ndimage.median instead. This version allows you to compute medians over a set of labeled areas in an array, which is exactly what you have with idx. This method assumes that your index array contains N densely packed values, all of which are in range(N). Otherwise the reshaping operations will not work properly.

If that is not the case, start by transforming idx:

_, ind = np.unique(idx, return_inverse=True)
idx = ind.reshape(idx.shape)

OR

idx = np.unique(idx, return_inverse=True)[1].reshape(idx.shape)

Since you are actually computing a separate median for each region and channel, you will need to have a set of labels for each channel. Flesh out idx to have a distinct set of indices for each channel:

chan = dat.shape[-1]
offset = idx.max() + 1
index = np.stack([idx + i * offset for i in range(chan)], axis=-1)

Now index has an identical set of regions defined in each channel, which you can use in scipy.ndimage.median:

out = scipy.ndimage.median(dat, index, index=range(offset * chan)).reshape(chan, offset).T

The input labels must be densely packed from zero to offset * chan for index=range(offset * chan) to work properly, and the reshape operation to have the right number of elements. The final transpose is just an artifact of how the labels are arranged.

Here is the complete product, along with an IDEOne demo of the result:

import numpy as np
from scipy.ndimage import median

dat = np.arange(12).reshape(2, 2, 3)
idx = np.array([[0, 0], [1, 2]])

def summarize(dat, idx):
    idx = np.unique(idx, return_inverse=True)[1].reshape(idx.shape)
    chan = dat.shape[-1]
    offset = idx.max() + 1
    index = np.stack([idx + i * offset for i in range(chan)], axis=-1)
    return median(dat, index, index=range(offset * chan)).reshape(chan, offset).T

print(summarize(dat, idx))