Indexação avançada e malhas

Pegar o melhor de cada linha

O argmax(axis=1) dá a posição do maior valor em cada linha. Para pegar os valores nessas posições, o take_along_axis usa as posições ao longo de um eixo, e o [:, None] as prepara para combinar por 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]]

A segunda linha ordena cada linha pela sua própria ordem, algo que o np.sort(axis=1) também faria. A vantagem do take_along_axis aparece quando você quer reordenar outro array com as posições de um primeiro.

Selecionar linhas e colunas ao mesmo tempo

Com duas listas, a indexação fancy combina par a par (capítulo 9). Para o produto cartesiano (todas as combinações de linhas e colunas escolhidas), o np.ix_ monta os índices certos:

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

[[ 1  3]
 [ 9 11]]

O resultado é o bloco formado pelas linhas 0 e 2 e pelas colunas 1 e 3.

Malhas de coordenadas

O np.meshgrid gera, a partir de dois vetores, duas grades com todas as coordenadas de um plano. É a base para calcular uma função em uma região inteira e depois desenhá-la:

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]]

O np.hypot(X, Y) calcula a distância de cada ponto da malha à origem. O resultado é a distância sobre uma grade, sem nenhum laço.

Onde está o máximo?

O argmax de um array 2D devolve a posição no array achatado. O np.unravel_index a converte em coordenadas (linha, coluna). E o np.argwhere lista todas as posições que satisfazem uma condição:

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]]

Repare na forma como o NumPy 2 imprime a tupla: cada coordenada aparece como np.int64(...), porque elas são escalares NumPy, e não inteiros do Python. É só a representação, e as contas funcionam igual. Em NumPy 1.x, você veria apenas os números.

One-hot por indexação

Transformar rótulos numéricos em vetores "one-hot" (um 1 na posição da classe, zeros no resto) é uma operação comum em aprendizado de máquina, e a indexação fancy de uma matriz identidade a faz em uma linha:

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

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