Можно ли ускорить этот кусок кода Python? ⇐ Python

[Anonymous]

У меня есть следующий код, который я хочу ускорить:

Код: Выделить всё

import numpy as np

def do_calc(a, b, d, sum_)->None:
idx = np.ndindex(a.shape[:1])
sum_[:] = 0
for i in idx:
sum_[0:d[i[0]]] -= a[i, 0:d[i[0]]][0]
sum_[d[i[0]]] += b[i, d[i[0]]][0]
return sum_

a=np.ones((5,10), order='F')
b=np.ones((5,10), order='F')*10

sum_ = np.empty(a.shape[1])

d = [8,5,2,7,1]
do_calc(a, b, d, sum_)

Итак, a и b — это статические массивы размером 5x10, и я хочу суммировать их в sum_, 1d-массив длиной 10, следующим образом:

• для каждой строки i из b я хочу включить только индекс d, то есть b[i , d], который войдет в положение sum_[d]
• для каждой строки i из, я хочу включить все значения до индекса d, то есть a[i, 0:d], который перейдет в позиции sum_[0:d]

в этом примере у нас есть:

Код: Выделить всё

a= array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

b= array([[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.],
[10., 10., 10., 10., 10., 10., 10., 10., 10., 10.]])

d= [8, 5, 2, 7, 1]

sum_= array([-5.,  6.,  7., -3., -3.,  8., -2.,  9., 10.,  0.])

Я пробовал следующее:

• Используя порядок Fortan, поскольку в основном мы копируем последний индекс.

  Код: Выделить всё

  sum_

  передается в функцию do_calc, поэтому ее не нужно запускать каждый раз. Вокруг последних двух строк происходит большой цикл, в котором мы генерируем новые значения для d и каждый раз запускаем do_calc.
• Использование ndindex для перебора строк.
  Интересно, можно ли это сделать более эффективно, поскольку именно на эту часть код тратит все время. На самом деле a и b не являются постоянными, поэтому мы не можем использовать их в этом примере.
  
  Подробнее здесь: https://stackoverflow.com/questions/791 ... e-speed-up