Как заменить циклы for матричными операциями?

Question

for x in range(0, width):
    for y in range(0, width):
        D[x,y] = dist(A[x,y], B[x,y])

def dist(a, b):
    return ((a-b)**2).sum()

Значения в A и B являются матрицами формы (3,3)

Как можно добиться того же результата, используя матричные операции вместо циклов for? Следующий код не работает.

D[[range(0,height)],[range(0,width)]] = dist(A[[range(0,height)],[range(0,width)], 
                                             B[[range(0,height)],[range(0,width)])]

Answer 1

Моя интерпретация заключается в том, что A и B являются массивами NumPy с размерами h * w * 3 * 3.

Простейшим способом использования операций с массивами для определения квадрата расстояния между матрицами 3 на 3 A[i,j,:,:] и B[i,j,:,:] для каждого i,j является запись:

distances = np.sum((A-B)**2, axis=(2,3))

т.е.. просто возводя в квадрат все и суммируя по двум последним осям.

Это включает в себя создание некоторых ненужных временных массивов (A-B, затем его квадрат), поэтому, если вам нужна эффективность extreme , вы можете использовать другую опцию, такую ​​как цикл Numba joted for, который выполняется np.vdot для каждой пары A[i,j,:,:] и B[i,j,:,:]. Но вряд ли это ваше узкое место.