Как «переместить» значения двумерного массива NumPy в соответствии с матрицей замещения

Question

DD - это двумерный массив размером (Nx, Ny). Ux и Uy - это два двумерных массива целых чисел одинакового размера, которые определяют, как перемещать значения DD в новую матрицу DD_new.

Например, если Ux[0,0]=2 и Uy[0,0]=3, то значение DD[0,0] будет сохранено в DD_new[2,3]. Если несколько значений DD необходимо сохранить в одном и том же расположении DD_new, они будут суммированы.

Я ищу оптимизированный сценарий, чтобы сделать это "удаление" как можно быстрее.

У меня уже есть короткий скрипт, но если кто-то знает, как его улучшить, используя манипуляции с матрицами и массивами, он мне поможет:

XX, YY = np.meshgrid(range(Nx), range(Ny))

XX_new = (XX + Ux).astype(int)
YY_new = (YY + Uy).astype(int)

for ix in range(Nx):
    for iy in range(Ny):
        x_new = XX_new[ix, iy]
        y_new = YY_new[ix, iy]
        if (x_new >= 0) and (x_new < Nx) and (y_new >= 0) and (y_new < Ny):
            DD_new[x_new, y_new] += DD[XX[ix, iy], YY[ix, iy]]

Кроме того, если кто-то знает, как обобщить эту проблему с помощью нецелочисленных матриц Ux и Uy и путем отправки значений DD для нескольких соседних расположений матрицы DD_new, это может помочь мне !

Answer 1

Вы можете сделать это с помощью np.add.at:

import numpy as np

DD = [[ 1,  2,  3, 4],
      [ 5,  6,  7, 8],
      [ 9, 10, 11, 12]]
Ux = [[ 0,  1,  2,  0],
      [ 1,  2,  0,  1],
      [ 2,  0,  2,  1]]
Vx = [[ 0,  1,  2,  3],
      [ 0,  1,  2,  3],
      [ 0,  1,  2,  3]]
DDnew = np.zeros_like(DD)
np.add.at(DDnew, (Ux, Vx), DD)
print(DDnew)
# [[ 1 10  7  4]
#  [ 5  2  0 20]
#  [ 9  6 14  0]]