title: Slice based on numpy.argmin results
body: <p>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).</p>
  <p>With <code>np.argmin(np.linalg.norm(array, axis=1), axis=0)</code> 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).</p>
  <p>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?</p>
  <p>A minimal reproducible example as requested where distances is the array:</p>
  <pre class="lang-py prettyprint-override"><code>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)
  </code></pre>
  <p>out:</p>
  <pre><code>(36, 2, 400, 400)
  (400, 400)
  </code></pre>
  <p><a href="https://i.stack.imgur.com/NQwPA.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/NQwPA.png" alt="Shows the index which I would like to use to slice" /></a></p>
  <p>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.</p>
  <pre><code>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)
  </code></pre>
  <p>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.</p>
  <p>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.</p>