Indexing, masking and fancy selection

An index carries one entry per axis, in order, separated by commas. Axes you do not mention are taken whole:

a.shape == (4, 4)

a[1, 2]      row 1, column 2          one number
a[1]         row 1, every column      the same as a[1, :]
a[1:3, 2]    rows 1 and 2, column 2
a[:, 1:3]    every row, columns 1 and 2

What decides the rank of the result is short enough to say in one line, and it is worth saying in exactly these words:

An integer removes that axis. A slice keeps it.

So on a (4, 4) array, a[1, 2] removes both axes and leaves a scalar, a[1:3, 2] removes one and keeps one, and a[1:3, 2:3] keeps both. Three expressions, overlapping numbers, three different ranks.

The pair that catches everybody is the smallest one:

a[0]      -> shape (4,)     an integer, so axis 0 is gone
a[0:1]    -> shape (1, 4)   a slice, so axis 0 stays, at length one

Same four numbers, different object. a[0:1] is still a stack of rows that happens to hold one row, and it broadcasts, prints and combines like one.

Slices take a step as well, and negative positions count from the end: a[::2] is every other row, a[:, -1] is the last column, and a[::-1] is the rows in reverse.

Instead of a slice, give an axis a list of positions. NumPy fetches exactly those, in exactly that order:

a.shape == (4, 4)

a[[2, 0, 0]]     rows 2, 0 and 0 -- in that order, and 0 twice
a[:, [3, 1]]     every row, columns 3 and 1, in that order

This is fancy indexing, also called integer indexing, and it does two things no slice can do: it takes positions in any order, and it takes the same position more than once. Neither can be described by a start and a step, which is why the previous lesson put it firmly on the copy side of the line.

Now the part that is got wrong more than anything else in this lesson. The shape of the result is not the shape of the array:

The result takes the shape of the index array.

a.shape == (4, 4)

a[[2, 0, 0]]        -> (3, 4)   three positions, so three rows
a[[1]]              -> (1, 4)   one position, so one row -- still 2-D
a[1]                -> (4,)     an integer, so the axis is gone

And the index does not have to be flat. Whatever shape you index with is the shape you get back, with the indexed axis replaced by the whole of it:

v.shape == (6,)

v[[0, 3]]             -> (2,)
v[[[0, 3], [3, 0]]]   -> (2, 2)   a 2x2 index, so a 2x2 result

That reads like a curiosity and it is the mechanism under embedding lookup, gather, one-hot decoding and every table join written in NumPy: an index array of any shape, against a table, gives a result of that shape with the table's rows dropped in.

Compare an array with a number and you do not get a yes or a no. You get a bool array of the same shape, one answer per element:

a       ==  [[ 3, -1,  4],
             [-5,  9, -2]]

a > 0   ==  [[ True, False,  True],
             [False,  True, False]]

That array is a mask, and indexing with it keeps every element sitting under a True:

a[a > 0]   ->   [3, 4, 9]

Two things are always true of that result, and both surprise people.

• It is 1-D. Always, whatever the rank of the array. The elements that passed do not form a rectangle, so there is no shape to keep, and they come out in row-major order -- across the first row, then the second.
• Its length is not known until it runs. The shape depends on the data rather than on the code, which is why a masked result is the one thing in array programming you cannot check by reading.

Masks combine with & and |, elementwise, and the parentheses are not optional:

a[(a > 0) & (a < 7)]     both conditions
a[(a < 0) | (a > 7)]     either one

& binds tighter than >, so without the parentheses Python groups it as a > (0 & a) < 7 and you get an error rather than an answer. Plain and and or do not work at all: they ask the whole array for a single truth value, and it refuses.

Because True counts as 1, the mask is also how you count without selecting anything:

(a > 0).sum()     how many passed
(a > 0).mean()    what fraction passed
(a > 0).any()     did anything
(a > 0).all()     did everything

Clipping every negative to zero is one line, and it is the line this idea exists for:

a[a < 0] = 0

Which looks like it contradicts the previous lesson. A boolean mask copies, so how can writing through one reach the original?

It does not. Nothing is written through anything, because no intermediate array is ever built. a[mask] on the right of an assignment is a read, and reads gather into new memory. a[mask] = 0 is not a read at all -- it is a request to store into a, and the mask only says which elements.

So these three lines are three different programs:

a[a < 0] = 0          a is clipped

a[a < 0] + 0          a new 1-D array of results, immediately discarded

b = a[a < 0]
b[:] = 0              b is a copy; writing into it never reaches a

The rule generalizes past masks. Assignment lands for every kind of index, fancy indexing included -- a[[0, 2]] = 0 sets those two rows, even though a[[0, 2]] on its own hands back a copy.

The right-hand side broadcasts against however many elements the index selected, so it is either a scalar or something whose shape matches -- and with a mask that count is only known at run time:

a[a < 0] = 0            a scalar, against any number of elements
a[a < 0] *= -1          the same elements, read and written
a[a < 0] = [1, 2]       fine only if exactly two elements passed

Two index arrays in one subscript are zipped, not crossed. They are walked together, and you get one element per pair:

a[[0, 1], [0, 1]]      ->  a[0, 0] and a[1, 1]        2 elements
a[[0, 1]][:, [0, 1]]   ->  rows 0 and 1, then columns 0 and 1   a 2x2 block

Two subscripts that look like the same request and are not. The first names two positions; the second names a set of rows, and then a set of columns of what came back. If you wanted the block and wrote the pair, you get the diagonal of the block and no error at all.

The shape rule from the fancy-indexing idea still holds, applied to both arrays at once:

The result takes the shape the index arrays broadcast to.

Two lists of length 3 broadcast to (3,), so you get three elements. Which also says how to ask for the rectangle in a single subscript -- give the two indexes shapes that broadcast into a grid:

rows = [0, 1]
cols = [0, 1]

a[rows, cols]                  (2,) with (2,)     -> (2,)     the pairs
a[[[0], [1]], [0, 1]]          (2, 1) with (2,)   -> (2, 2)   the block
a[np.array(rows)[:, None], cols]     the same thing, as you will see it written

The zipped form is the more useful of the two, and it is worth recognizing on sight, because it is how you pick one element per row: one index array walking the rows and one saying which column to take from each.

np.where(condition, then, otherwise) is an if/else evaluated for every element at once. It keeps the shape, and both branches broadcast:

np.where(a > 0, a, 0)      the clip again, as a new array
np.where(a > 0, 1, -1)     a sign, with no loop

Hold it against a[a > 0], which is the other thing a mask can do. The mask selects and gives you a shorter 1-D array; where replaces and gives you the same shape back. Wanting one and writing the other is a shape error at best.

The other half of this idea is a family of functions that return positions rather than values, and they all begin with arg:

v = [30, 10, 40, 20]

np.argmax(v)        -> 2               where the largest is, not what it is
np.argsort(v)       -> [1, 3, 0, 2]    the order that would sort it
v[np.argsort(v)]    -> [10, 20, 30, 40]

That last line is the pattern. An arg function computes an index array, and the previous ideas say what happens when you index with one. Sorting a second array into the same order is the same expression with a different array in front of the subscript, which is how you sort a table by one of its columns.

On a 2-D array, argmax(axis=1) gives one position per row, and turning those positions back into values is exactly the zipped pair from the previous idea:

best = a.argmax(axis=1)              one column number per row
a[np.arange(len(a)), best]           the winning value in each row

Compute positions, then index with them. Top-k with argsort, ranking, taking the predicted class's probability out of a batch of scores -- all of it is that one habit.

Everything above selects. Whether what you are handed is a window onto a or a fresh array of its own is decided by the rule from the previous lesson, and now there is enough vocabulary to run it on every form:

a[1:3, 2]           view    basic slicing, evenly spaced
a[:, ::2]           view    a start and a step is all it is
a[0]                view    an integer is basic indexing too

a[[0, 2]]           copy    arbitrary positions
a[a > 0]            copy    arbitrary positions
a[[0, 1], [0, 1]]   copy    arbitrary positions
np.where(a > 0, a, 0)  copy  new values, so new memory

Indexing twice adds one clause: a chain reaches the original only if every step is a view. Once any step gathers, everything after it is working on the gathered array:

a[1:3][:, [0]]     the slice is a view; the gather is not
                   what you hold is a copy, and writing into it
                   changes nothing you will ever look at again

And the asymmetry from the assignment idea is worth restating next to the table, because the two rules pull in opposite directions and both are true:

• Reading through fancy indexing or a mask gives you a copy.
• Assigning through fancy indexing or a mask lands on the original.

Which is why a[mask] = 0 works and b = a[mask]; b[:] = 0 does not, and why the second one raises nothing while it fails to do anything.