Если вы точно не знаете, что является следствием явных объявлений ввода и вывода, пусть numba решит это. С вашим вводом вы можете использовать 'float64(float64[::1],float64[::1])'. (скалярный вывод, непрерывные входные массивы). Если вы вызовете явно объявленную функцию с пошаговыми вводами, она потерпит неудачу, если Numba сделает работу, она просто перекомпилируется.
Без использования fastmath=True также невозможно использовать SIMD, потому что это изменяет точность результата.
Вычисление не менее 4 частичных сумм (256-битный вектор) и вычисление суммы этих частичных сумм здесь предпочтительнее (Numpy также не вычисляет наивную сумму).
Пример использования превосходной утилиты тестирования MSeiferts
import numpy as np
import numba as nb
from itertools import chain
def python_loop(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
@nb.njit
def numba_loop_zip(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
#Your version with suboptimal input and output (prevent njit compilation) declaration
@nb.jit('float64[:](float64[:],float64[:])')
def numba_your_func(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
@nb.njit(fastmath=True)
def numba_loop_zip_fastmath(a,b):
result = 0.0
for i,j in zip(a,b):
result += (i+j)
return result
@nb.njit(fastmath=True)
def numba_loop_fastmath_single(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
for i in range(size):
result += a[i]+b[i]
return result
@nb.njit(fastmath=True,parallel=True)
def numba_loop_fastmath_multi(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
for i in nb.prange(size):
result += a[i]+b[i]
return result
#just for fun... single-threaded for small arrays,
#multithreaded for larger arrays
@nb.njit(fastmath=True,parallel=True)
def numba_loop_fastmath_combined(a,b):
result = 0.0
size=min(a.shape[0],b.shape[0])
if size>2*10**4:
result=numba_loop_fastmath_multi(a,b)
else:
result=numba_loop_fastmath_single(a,b)
return result
def numpy_methods(a, b):
return a.sum() + b.sum()
def python_sum(a, b):
return sum(chain(a.tolist(), b.tolist()))
from simple_benchmark import benchmark, MultiArgument
arguments = {
2**i: MultiArgument([np.zeros(2**i), np.zeros(2**i)])
for i in range(2, 19)
}
b = benchmark([python_loop, numba_loop_zip, numpy_methods,numba_your_func, python_sum,numba_loop_zip_fastmath,numba_loop_fastmath_single,numba_loop_fastmath_multi,numba_loop_fastmath_combined], arguments, warmups=[numba_loop_zip,numba_loop_zip_fastmath,numba_your_func,numba_loop_fastmath_single,numba_loop_fastmath_multi,numba_loop_fastmath_combined])
%matplotlib notebook
b.plot()
Обратите внимание, что использование numba_loop_fastmath_multi или numba_loop_fastmath_combined(a,b) рекомендуется только в некоторых особых случаях. Чаще всего такая простая функция является частью другой проблемы, которая может быть более эффективно распараллелена (запуск потоков имеет некоторые накладные расходы)