Início / coding / Perguntas / 78748429
Accepted
sams-studio
Asked: 2024-07-15 14:58:46 +0800 CST

Aplicar matriz de permutação em vários eixos em numpy

  • 772

Digamos que eu tenha uma série de permutações permque poderiam ser assim:

perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])

Se eu quiser aplicá-lo a um eixo, posso escrever algo como:

v = np.arange(9).reshape(3, 3)
print(v[perm])

Saída:

array([[[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]],

       [[3, 4, 5],
        [6, 7, 8],
        [0, 1, 2]],

       [[0, 1, 2],
        [6, 7, 8],
        [3, 4, 5]],

       [[6, 7, 8],
        [3, 4, 5],
        [0, 1, 2]]])

Agora gostaria de aplicá-lo a dois eixos ao mesmo tempo. Eu descobri que posso fazer isso através de:

np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

Mas acho isso bastante ineficiente, porque precisa criar uma grade de malha e também requer um loop for. Criei um array pequeno neste exemplo, mas na realidade tenho arrays muito maiores com muitas permutações, então adoraria ter algo que fosse tão rápido e simples quanto a versão de um eixo.

python

3 respostas

  1. Best Answer

    Que tal:

    p1 = perm[:, :, np.newaxis]
p2 = perm[:, np.newaxis, :]
v[p1, p2]

    O eixo zero de p1e p2é apenas a dimensão de "lote" de perm, o que permite fazer muitas permutações em uma operação.

    A outra dimensão de perm, que corresponde aos índices, está alinhada ao longo do primeiro eixo p1e do segundo em p2. Como os eixos são ortogonais, os arrays são transmitidos , basicamente como os arrays que você usou meshgrid- mas ainda têm a dimensão do lote.

    É o melhor que posso fazer no meu celular :) Posso tentar esclarecer mais tarde, se necessário, mas a ideia principal é a transmissão.

    Comparação:

    import numpy as np
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])
v = np.arange(9).reshape(3, 3)

ref = np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

p1 = perm[:, :, np.newaxis]
p2 = perm[:, np.newaxis, :]
res = v[p1, p2]

np.testing.assert_equal(res, ref)
# passes

%timeit np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])
# 107 µs ± 20.6 µs per loop

%timeit v[perm[:, :, np.newaxis], perm[:, np.newaxis, :]]
# 3.73 µs ± 1.07 µs per loop

    Um exemplo mais simples (sem dimensão de lote) de índices de transmissão:

    import numpy as np
i = np.arange(3)
ref = np.meshgrid(i, i, indexing="ij")
res = np.broadcast_arrays(i[:, np.newaxis], i[np.newaxis, :])
np.testing.assert_equal(res, ref)
# passes

    No código da solução na parte superior, a transmissão está implícita. Não precisamos ligar broadcast_arraysporque isso acontece automaticamente durante a indexação.

    • 4

  2. Você pode se livrar do meshgridcom

    a = np.array([v[p][:,p] for p in perm])
b = np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])
print(np.all(b == a)) # True

    Isso é 5x mais rápido no seu array de exemplo:

    import timeit
%timeit np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm]) # 42.7 µs
%timeit np.array([v[p][:,p] for p in perm]) # 8.18 µs

    Eu diria que o loop for é praticamente irrelevante. Se você estiver preocupado com otimização adicional, especifique as formas com as quais está trabalhando...

    • 2

  3. Para aplicar permutações a dois eixos simultaneamente sem usar loops ou meshgrid, você pode usar indexação e transmissão avançadas. Aqui está uma abordagem mais eficiente:

    import numpy as np

# Define permutations and initial array
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])
v = np.arange(9).reshape(3, 3)

# Apply permutations using broadcasting
permuted_indices = perm[:, :, None] * 3 + perm[:, None, :]
result = v.ravel()[permuted_indices]

print(result)
    1. perm[:, :, None] * 3 + perm[:, None, :] calcula os índices 1D dos valores permutados na matriz achatada usando transmissão.
    2. v.ravel() nivela o array v .
    3. O resultado é a matriz de saída desejada com permutações aplicadas a ambos os eixos.

    Essa abordagem é mais eficiente e evita loops explícitos e meshgrid .

    • 0

relate perguntas