I have two NxMx3 arrays and I want to reduce over the last dimension of the
first array by selecting those elements corresponding to the index of the
maximum value of each 3-vector of the second array to give an NxM result.
Hopefully that makes sense? If not hopefully the example below will shed some
light.
Can anyone think of a more efficient way to do this than my 1st attempt? I
thought that nditer might be the solution; I got it to work (mostly by trial &
error) but found that it's ~50x slower for this task! Is this not a good usecase
for nditer or am I doing something wrong?
In [42]: value_array = np.outer(ones(1000), arange(3)).reshape(20,50,3)
...: index_array = randn(*value_array.shape)
In [43]: indices = index_array.reshape(-1,3).argmax(axis=1)
...: result = value_array.reshape(-1,3)[np.arange(indices.size), indices]
...: result = result.reshape(value_array.shape[0:-1])
In [44]: it = np.nditer([value_array, index_array, None],
...: flags=['reduce_ok', 'external_loop','buffered', 'delay_bufalloc'],
...: op_flags=[['readonly'],['readonly'],['readwrite', 'allocate']],
...: op_axes=[[0,1,2], [0,1,2], [0,1,-1]])
...: it.reset()
...: for values, index_values, out in it:
...: out[...] = values[index_values.argmax()]
...: #
In [45]: allclose(result, it.operands[2])
Out[45]: True
In [46]: %%timeit
...: indices = index_array.reshape(-1,3).argmax(axis=1)
...: result = value_array.reshape(-1,3)[np.arange(indices.size), indices]
...: result = result.reshape(value_array.shape[0:-1])
...:
10000 loops, best of 3: 113 µs per loop
In [47]: %%timeit
...: it = np.nditer([value_array, index_array, None],
...: flags=['reduce_ok', 'external_loop','buffered', 'delay_bufalloc'],
...: op_flags=[['readonly'],['readonly'],['readwrite', 'allocate']],
...: op_axes=[[0,1,2], [0,1,2], [0,1,-1]])
...: it.reset()
...: for values, index_values, out in it:
...: out[...] = values[index_values.argmax()]
...: #
...:
100 loops, best of 3: 5.26 ms per loop
In [48]:
Thanks,
Dave