Как увеличить производительность цикла for в python

Question

Я много симулирую с python, симулирую системные ответы.

В настоящее время я использую схему Рунге-Кутты, но наткнулся на другую схему, которую я тестировал.

При тестировании в Matlab я достигаю исключительной производительности по сравнению с моей Runge-Kutta. Однако когда я перенес это в Python, это было значительно медленнее.

Я не уверен, так ли это, или я мог бы улучшить свой способ кодирования, поэтому я хотел бы услышать некоторые ваши отзывы, если это возможно.

Код в Matlab, пример:

dt = 0.0001;
f = randn(1, (60 / dt));
ns = length(f);
yo = zeros(3,1);
P1 = [0; 0.001; 0];
F = [1 0.0001 0; 0.001 1 0.0001; 0.001 0 1];
y1 = zeros(3, ns);
tic
for i = 1:ns
    y1(:, i) = P1*f(:, i) + F*yo;
    yo = y1(:, i);
end
toc

В котором цикл выполняется за 0,55-0,61 сек.

Код на Python, пример:

dt = 0.0001
f = np.random.randn(1, int(60 / dt))
ns = np.size(f)
yo = np.zeros((3))
F = np.array([[1, 0.0001, 0], [0.001, 1, 0.0001], [0.001, 0, 1]])
P1 = np.transpose(np.array([[0, 0.0001, 0]]))
y1 = np.zeros((3, ns), order='F')
start_time = time.time()
for i in range(ns-1):
    y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
    yo = y1[: , i]
print("--- %s seconds ---" % (time.time() - start_time))

В котором цикл выполняется за 2,8 -3,1 с.

Могу ли я сделать что-нибудь, чтобы улучшить это?

Спасибо за рассмотрение моего вопроса.

Answer 1

Я немного оптимизировал ваш код, время выполнения для меня сократилось с 2,8 с до 1,2 с. Прежде чем искать более быстрых интерпретаторов, я рекомендую вам выполнить профилирование (см. Line_profiler) и попытаться удалить все, что вы можете, из внутреннего цикла. Лучше избегать любых явных циклов for и полагаться на такие простые функции, как точка, einsum и т. Д.

Я думаю, что еще есть место для оптимизации. Я не думаю, что я изменил ваши ценности, но лучше проверить. С другими инструментами, такими как numba или cython ( cython.org ) или pypy ( pypy.org ), я думаю, ваше время выполнения значительно улучшится.

#!/usr/bin/env python3

import numpy as np
import time

np.random.seed(0)

#@profile
def run():
    dt = 0.0001
    f = np.random.randn(1, int(60 / dt))
    ns = np.size(f)
    yo = np.zeros((3))
    F = np.array([[1, 0.0001, 0], [0.001, 1, 0.0001], [0.001, 0, 1]])
    P1 = np.transpose(np.array([[0, 0.0001, 0]]))
    start_time = time.time()
    y1 = np.outer(f, P1)
    for i in range(ns-1):
        y1[i] += F@yo
        yo = y1[i]
    print("--- %s seconds ---" % (time.time() - start_time))
    y1 = y1.T
    print(yo)

run()

Answer 2

Я предложил использовать numba в комментариях. Вот пример:

import numba
import numpy as np

def py_func(dt, F, P1):
    f = np.random.randn(1, int(60 / dt))
    ns = f.size
    yo = np.zeros((3))
    y1 = np.zeros((3, ns), order='F')
    for i in range(ns-1):
        y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
        yo = y1[: , i]
    return yo

@numba.jit(nopython=True)
def numba_func(dt, F, P1):
    f = np.random.randn(1, int(60 / dt))
    ns = f.size
    yo = np.zeros((3))
    y1 = np.zeros((3, ns))
    for i in range(ns-1):
        y1[:, i] = np.dot(P1, f[:, i]) + np.reshape(np.dot(F, yo), (3))
        yo = y1[: , i]
    return yo

Вы не можете использовать порядок 'F' с numba, поскольку он использует массивы типа C, а не массивы FORTRAN.

Разница во времени показана ниже:

Чистый цикл Python:

%%timeit
py_func(dt, F, P1)

Результаты:

2.88 s ± 100 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Numba:

%%timeit
numba_func(dt, F, P1)

Результаты:

588 ms ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)