Можно ли векторизовать это вычисление for-l oop np.meshgrid?

Question

У меня есть раздел внутри моего кода, где я через for l oop распределяю результаты из np.meshgrid, примененные к фрагменту my_array, в заполнитель res_array следующим образом:

p = 360; q = 50; r = 50
my_array   =  np.random.rand(p, q, 2)
res_array  = np.zeros((p,q,r), dtype=np.float32)
for i in range(p):
   x, x_ = np.meshgrid(my_array[i,:,0], my_array[i,:,0])
   res_array[i] = (x_-x)

В моем коде есть еще пара подобных операций, и мне приходится запускать свой код несколько тысяч раз. Поэтому скорость становится проблемой. Я думал, что смогу сэкономить здесь время, если смогу обойтись без for l oop, однако я не могу сделать это сам.

Как это сделать без for l oop?

ИЗМЕНИТЬ ИСТОРИЮ:

1. Обратите внимание, что my_array[i,:,1] не используется в показанном выше фрагменте. Я аналогичным образом использую my_array[i,:,1] в другом блоке кода (здесь не показан, но похож на l oop). Это то, что я имею в виду под словами «В моем коде есть еще пара похожих операций ...»

Answer 1

Вы можете переписать np.meshgrid как более быстрые операции нижнего уровня numpy:

p = 360; q = 50; r = 50
my_array   =  np.random.rand(p, q, 2)
res_array  = np.zeros((p,q,r), dtype=np.float32)
for i in range(p):
    x = my_array[None,i,:,0].repeat(q, axis=0).reshape(q, q)
    y = my_array[None,i,:,0].repeat(q, axis=1).reshape(q, q)
    res_array[i] = y - x

Этот код в 2 раза быстрее на моей машине. Numba @njit может использоваться для ускорения вышеуказанного кода (снова в 3 раза быстрее), но с его помощью возможна гораздо более эффективная реализация. Вот он:

@njit(parallel=True)
def fasterImpl(my_array, p, q, r):
    res_array = np.zeros((p, q, r))
    for i in prange(p):
        for j in range(q):
            for k in range(r):
                res_array[i,j,k] = my_array[i,j,0] - my_array[i,k,0]
    return res_array

p = 360; q = 50; r = 50
my_array   = np.random.rand(p, q, 2)
res_array  = fasterImpl(my_array, p, q, r)

Эта финальная реализация в 29 раз быстрее, чем исходная на моей машине!

Answer 2

вы можете сделать так:

my_array   =  np.random.rand(p, q, 2)
res_array = np.array([x[1]-x[0] for x in [np.meshgrid(my_array[i,:,0], my_array[i,:,0]) for i in range(p)]], np.float32).reshape(p,q,r)

Но требует больше времени, чем вы.