Векторизация матричной операции на основе индекса в numpy

Question

Как я могу векторизовать следующий цикл?

def my_fnc():
   m = np.arange(27.).reshape((3,3,3))
   ret = np.empty_like(m)

   it = np.nditer(m, flags=['multi_index'])
   for x in it:
      i,j,k = it.multi_index
      ret[i,j,k] = x / m[i,j,i]
   return ret

По сути, я делю каждое значение в м на что-то похожее на диагональ. Не все значения в m будут разными, arange - только пример.

Заранее спасибо! ~

P.S .: Вот результат функции выше, не обращайте внимания на nans:)

array([[[        nan,         inf,         inf],
        [ 1.        ,  1.33333333,  1.66666667],
        [ 1.        ,  1.16666667,  1.33333333]],

       [[ 0.9       ,  1.        ,  1.1       ],
        [ 0.92307692,  1.        ,  1.07692308],
        [ 0.9375    ,  1.        ,  1.0625    ]],

       [[ 0.9       ,  0.95      ,  1.        ],
        [ 0.91304348,  0.95652174,  1.        ],
        [ 0.92307692,  0.96153846,  1.        ]]])

Answer 1

Используйте advanced-indexing, чтобы получить эквивалент m[i,j,i] за один раз, а затем просто разделите входной массив на него -

r = np.arange(len(m))
ret = m/m[r,:,r,None] # Add new axis with None to allow for broadcasting