1

Let us have a numpy array (float) with shape equal to (36, 2, 400, 400). Let us say the 400 by 400 represents an image. Then for each pixel I would like to find the two values (second dimension) which are when taking the norm over the second dimension, the lowest with respect to the first dimension. So I end up with an array of shape (2, 400, 400).

With np.argmin(np.linalg.norm(array, axis=1), axis=0) I am able to get the index for each of those 2 by 400 by 400 pixels which is almost what I want. But now I want to use this number to slice the original array in the first dimension so I am left with an array of shape (2, 400, 400).

What I can do is loop over all indices and construct the result pixel by pixel, but I am convinced there is a smarter way. Can anyone help me with a smarter way?

A minimal reproducible example as requested where distances is the array:

shape = (400, 400)
centers = np.random.randint(400, size=(36, 2))

distances = np.array([np.indices(shape) - np.array(center)[:, None, None] for center in centers])
nearest_center_index = np.argmin(np.linalg.norm(distances, axis=1), axis=0)

print(distances.shape)
print(nearest_center_index.shape)
plt.imshow(nearest_center_index)

out:

(36, 2, 400, 400)
(400, 400)

Shows the index which I would like to use to slice

I was able, with help from the comments, to produce a somewhat ugly answer, which helped me futher to understand the problem. Let me elaborate. What is possible to do is to flatten the image and argmin results and then use advanced indexing with argmin and indices over the image to produce the results.

flatten_indices = nearest_center_index.reshape(400**2)
image_indices = range(400**2)
results = distances.reshape(36, 2, 400**2)[flatten_indices, :, image_indices].reshape(400, 400, 2).swapaxes(0, 2)

However, I think it happens a lot that you have indices that are shaped as a subset of the dimensions and have values containing indices of another dimension. I would expect a generic method to slice this.

Thus let us have and array with n dimensions with shape = (x1, x2, ..., xn) and let us say we have a array representing indices for a dimension, e.g., xi, that has shape which is a subset of the shape of the original array and not containing xi. Then I would expect a method to slice this array.

F.Wessels
  • 179
  • 11
  • 1
    Something in here may help you: https://numpy.org/doc/stable/user/basics.indexing.html – jjramsey Jan 18 '22 at 14:07
  • @jjramsey - I doubt anyone writing a two-line Voronoi map algorithm is unfamiliar with basic indexing. – Michael Szczesny Jan 18 '22 at 14:10
  • 2
    @MichaelSzczesny True, but despite the name in the URL, the content actually covers various kinds of advanced indexing as well. – jjramsey Jan 18 '22 at 14:13
  • 1
    `np.take_along_axis(distances, nearest_center_index[None,None,:,:], 0).swapaxes(3,2)` This keeps the dimension of *arr* like slices usually do. And is a little bit faster. – Michael Szczesny Jan 18 '22 at 14:47
  • @MichaelSzczesny I think this is it, what I was looking for. At least it covers what I can imagine as the generic problem, at least for all usecases that I can think of, thx! – F.Wessels Jan 18 '22 at 14:54

1 Answers1

0

The function I was looking for is numpy.take_along_axis().

For the specific example the only thing needed to be done is making sure the nearest_center_index (output of argmin) has equal amount of dimensions as the to be sliced array. In the example this can be achieved by passing keepdims=True to both norm and argmin which then can be directly used as the second argument of the numpy function. The third argument should be the xi axis (in the example axis 0).

Without passing the keepdims=True, following the exact example, the stated objective can be achieved by:

result = np.take_along_axis(distances, nearest_center_index[None,None,:,:], 0)[0]
F.Wessels
  • 179
  • 11
  • 2
    If you pass `keepdims=True` to `norm` and `argmin`, you will get `nearest_center_index` of the correct shape automatically. – Mad Physicist Jan 18 '22 at 15:47