Numba @jit не может оптимизировать простую функцию

Вот переработка вашего кода, который выполняется примерно за 0,4 секунды на моей машине:

def ising_model_1d(N=200,J=1,T=1e-2,H=0,n_iter=1e6):
    n_iter = int(n_iter)
    beta = 1/T
    s = randn(N) > 10
    s[N-1] = s[0]

    mag = zeros(n_iter)
    aux_idx =  randint(low=0,high=N,size=n_iter)

    pre_rand = rand(n_iter)

    _ising_jitted(n_iter, aux_idx, s, J, N, H, beta, pre_rand, mag)

    return mag


@jit(nopython=True)
def _ising_jitted(n_iter, aux_idx, s, J, N, H, beta, pre_rand, mag):
    for i1 in range(n_iter):
        rnd_idx = aux_idx[i1]
        s_1 = s[rnd_idx*2] - 1
        s_2 = s[(rnd_idx+1)%(N)]*2 - 1
        s_3 = s[(rnd_idx-1)%(N)]*2 - 1
        delta_E = 2.0*J*(s_2+s_3)*s_1 + 2.0*H*s_1
        t = rand()
        if delta_E < 0:
            s[rnd_idx] = not s[rnd_idx]
        elif np.exp(-1*beta*delta_E) >= pre_rand[i1]:
            s[rnd_idx] = not s[rnd_idx]

        s[N-1] = s[0]
        mag[i1] = (s*2-1).sum()*1.0/N

Пожалуйста, убедитесь, что результаты соответствуют ожиданиям! Я изменил многое из того, что у вас было, и не могу гарантировать правильность расчетов!

Работа с numba требует осторожности. Функции Python, как и большинство функций numpy, не могут быть оптимизированы компилятором. Одна вещь, которую я считаю полезной, это использовать опцию nopython для @jit. Это означает, что компилятор будет жаловаться всякий раз, когда вы даете ему код, который он не может оптимизировать. Затем вы можете просмотреть сообщение об ошибке и найти строку, которая, вероятно, замедлит ваш код.

Я считаю, что хитрость заключается в том, чтобы написать функцию «шлюза» на Python, которая выполняет как можно больше работы, используя numpy и его векторизованные функции. Он должен создать пустые массивы, в которых вам нужно будет хранить результаты. Он должен упаковать все данные, которые вам понадобятся во время вычислений. Затем он должен передать все это вашей jit-функции в одном большом длинном списке аргументов.

Показательный пример: обратите внимание, как я обрабатываю генерацию случайных чисел в джиттинг-коде. В исходном коде вы вызвали rand():

elif(np.exp(-1*beta*delta_E) >= rand()):

Но rand() нельзя оптимизировать с помощью numba (по крайней мере, в старых версиях numba. В более новых версиях можно, при условии, что rand вызывается без аргументов). Наблюдение состоит в том, что вам нужно одно случайное число для каждой из n_iter итераций. Итак, мы просто создаем случайный массив, используя numpy в нашей функции-оболочке, а затем передаем этот случайный массив функции jitted. Получить случайное число так же просто, как индексировать этот массив.

Наконец, список numpy функций, которые можно оптимизировать с помощью последней версии компилятора, см. на странице здесь. В моей переработке вашего кода я был агрессивен в удалении вызовов функций numpy, чтобы код работал на большем количестве версий numba.