Advanced indexing and grids

Get the best from each row

The argmax(axis=1) gives the position of the largest value in each row. To get the values at these positions, the take_along_axis uses the positions along an axis, and the [:, None] prepares them to combine by broadcasting:

import numpy as np

notas = np.array([[7, 9, 5], [4, 8, 6]])
melhor = notas.argmax(axis=1)
print(melhor, np.take_along_axis(notas, melhor[:, None], axis=1).ravel())
print(np.take_along_axis(notas, np.argsort(notas, axis=1), axis=1))

[1 1] [9 8]
[[5 7 9]
 [4 6 8]]

The second line sorts each row by its own order, something that the np.sort(axis=1) would also do. The advantage of the take_along_axis appears when you want to reorder another array with the positions of a first.

Select rows and columns at the same time

With two lists, fancy indexing combines (chapter 9). For the Cartesian product (all combinations of selected rows and columns), the pairwise builds the right indices: np.ix_ The result is the block formed by rows 0 and 2 and columns 1 and 3.

m = np.arange(16).reshape(4, 4)
print(m[np.ix_([0, 2], [1, 3])])

[[ 1  3]
 [ 9 11]]

generates, from two vectors, two grids with

Coordinate grids

The np.meshgrid all the coordinates of a plane. It is the basis for calculating a function over an entire region and then plotting it: calculates the distance from each point of the grid to the origin. The result is the distance over a grid, without any loop.

x = np.linspace(-1, 1, 3)
y = np.linspace(-1, 1, 3)
X, Y = np.meshgrid(x, y)
print(X)
print(np.hypot(X, Y).round(2))

[[-1.  0.  1.]
 [-1.  0.  1.]
 [-1.  0.  1.]]
[[1.41 1.   1.41]
 [1.   0.   1.  ]
 [1.41 1.   1.41]]

The np.hypot(X, Y) of a 2D array returns the position

Where is the maximum?

The argmax in the flattened array converts it into coordinates. The np.unravel_index list (linha, coluna). And the np.argwhere the positions that satisfy a condition: all the positions that satisfy a condition:

plano = np.array([[3, 8, 1], [9, 2, 7]])
pos = np.unravel_index(plano.argmax(), plano.shape)
print(pos, plano[pos])
print(np.argwhere(plano > 6))

(np.int64(1), np.int64(0)) 9
[[0 1]
 [1 0]
 [1 2]]

Notice how NumPy 2 prints the tuple: each coordinate appears as np.int64(...), because they are NumPy scalars, not Python integers. It's just the representation, and the calculations work the same. In NumPy 1.x, you would only see the numbers.

One-hot by indexing

Transforming numeric labels into "one-hot" vectors (a 1 in the class position, zeros elsewhere) is a common operation in machine learning, and indexing combines an identity matrix does it in one line:

rotulos = np.array([0, 2, 1])
print(np.eye(4)[rotulos])

[[1. 0. 0. 0.]
 [0. 0. 1. 0.]
 [0. 1. 0. 0.]]