numpy shorthand for taking jagged slice

Question

I have an operation that I\'m doing commonly which I\'m calling a \"jagged-slice\" because I don\'t know the real name for it. It\'s best explained by example:

Answer 1

I think that your current method is probably the best way.

You can also use choose for this kind of selection. This is syntactically clearer, but is trickier to get right and potentially more limited. The equivalent with this method would be:

entries_of_interest.choose(a.T)

Answer 2

The elements in jagged_slice_of_a are the diagonal elements of a[:,entries_of_interest]

A slightly less cumbersome way of doing this would therefore be to use np.diagonal to extract them.

jagged_slice_of_a = a[:, entries_of_interest].diagonal()

Answer 3

This is combersome only in the sense that it requires more typing for a task that seems so simple to you.

a[np.arange(a.shape[0]), entries_of_interest]

But as you note, the syntactically simpler a[:, entries_of_interest] has another interpretation in numpy. Choosing a subset of the columns of an array is a more common task that choosing one (random) item from each row.

Your case is just a specialized instance of

a[I, J]

where I and J are 2 arrays of the same shape. In the general case entries_of_interest could be smaller than a.shape[0] (not all the rows), or larger (several items from some rows), or even be 2d. It could even select certain elements repeatedly.

I have found in other SO questions that performing this kind of element selection is faster when applied to a.flat. But that requires some math to construct the I*n+J kind of flat index.

With your special knowledge of J, constructing I seems extra work, but numpy can't make that kind of assumption. If this selection was more common someone could write a function that wraps your expression

def  peter_selection(a,I):
   # check the a.shape[0]==I.shape[0]
   return a[np.arange(a.shape[0]), I]