Ускорение цикла Python
Я создал цикл for, который принимает значение из одного столбца и смотрит вперед, если это значение было превышено один или два раза в последующих данных. Код работает, но поскольку набор данных, с которым он работает, очень велик, код работает очень медленно. Я подозреваю, что это особенно важно потому, что на каждой итерации подсчитывается, сколько раз значение превышено (около 500 тыс. строк). Есть ли способ ускорить это?
import pandas as pd
df1 = pd.DataFrame({'index': [0,1,2,3,4], 'Time': ['2022-01-01','2022-01-02','2022-01-03','2022-01-04','2022-01-05'], 'A':[234,456,323,576,234], 'B': [0,1,0,1,0], 'B.v': [0,234,0,323,0], 'in' : [0,0,0,0,0], 'out':[0,0,0,0,0]})
def calc(df1):
df2 = pd.DataFrame(df1[df1['B'] == 1])
for x in range(len(df2)):
index = df2.iloc[x, df2.columns.get_loc('index')]
tvalue = df2.iloc[x, df2.columns.get_loc('A')]
pointvalue = df2.iloc[x, df2.columns.get_loc('B.v')]
postrates = df1['A'].values[range(index,len(df1))]
if sum(pointvalue > postrates) == 1:
df1.iloc[index, df1.columns.get_loc('in')] = 1
if sum(pointvalue > postrates) >= 2:
df1.iloc[index, df1.columns.get_loc('in')] = 2
if sum(tvalue < postrates) == 1:
df1.iloc[index, df1.columns.get_loc('out')] = 1
if sum(tvalue < postrates) >= 2:
df1.iloc[index, df1.columns.get_loc('out')] = 2
return df1
if __name__ == "__main__":
print(calc(df1))
Я не уверен, что понимаю ваш комментарий, но я сделал это таким образом, потому что «посттраты» меняются со временем, поскольку они зависят от индекса «B.v». в основном вы хотите сравнивать только с будущими данными.
Обратите внимание, что использование sum в массиве Numpy неэффективно. Использование списков явно неоптимально из-за объектов. Известно, что использование iloc происходит медленно, хотя проблема, безусловно, связана с суммой. Кроме того, когда вы вводите sum(tvalue < postrates) дважды, он вычисляется дважды, что неэффективно. Имейте в виду, что CPython — это интерпретатор, поэтому он почти ничего не оптимизирует, и вам не придется много раз пересчитывать данные.
Ответы 1
Я бы использовал numba, чтобы ускорить вычисления (с добавленным коротким замыканием: когда две суммы уже >= 2, дальше считать нет необходимости):
import numba
@numba.njit(parallel=True)
def calc_in_out(A, B, Bv, out_in, out_out):
for idx in numba.prange(len(B)):
val_b = B[idx]
if val_b != 1:
continue
val_a = A[idx]
val_Bv = Bv[idx]
s1, s2 = 0, 0
for idx2 in range(idx, len(B)):
s1 += val_Bv > A[idx2]
s2 += val_a < A[idx2]
# no need to count further:
if s1 >= 2 and s2 >= 2:
break
if s1 == 1:
out_in[idx] = 1
elif s1 >= 2:
out_in[idx] = 2
if s2 == 1:
out_out[idx] = 1
elif s2 >= 2:
out_out[idx] = 2
Тест:
from timeit import timeit
import numba
import numpy as np
@numba.njit(parallel=True)
def calc_in_out(A, B, Bv, out_in, out_out):
for idx in numba.prange(len(B)):
val_b = B[idx]
if val_b != 1:
continue
val_a = A[idx]
val_Bv = Bv[idx]
s1, s2 = 0, 0
for idx2 in range(idx, len(B)):
s1 += val_Bv > A[idx2]
s2 += val_a < A[idx2]
# no need to count further:
if s1 >= 2 and s2 >= 2:
break
if s1 == 1:
out_in[idx] = 1
elif s1 >= 2:
out_in[idx] = 2
if s2 == 1:
out_out[idx] = 1
elif s2 >= 2:
out_out[idx] = 2
def setup_df(N=500_000):
return pd.DataFrame(
{
"Time": ["2022-01-01"] * N,
"A": np.random.randint(10, 1000, size=N),
"B": np.random.randint(0, 2, size=N),
"B.v": np.random.randint(10, 1000, size=N),
"in": [0] * N,
"out": [0] * N,
}
)
def main():
df1 = pd.DataFrame(
{
"index": [0, 1, 2, 3, 4],
"Time": [
"2022-01-01",
"2022-01-02",
"2022-01-03",
"2022-01-04",
"2022-01-05",
],
"A": [234, 456, 323, 576, 234],
"B": [0, 1, 0, 1, 0],
"B.v": [0, 234, 0, 323, 0],
"in": [0, 0, 0, 0, 0],
"out": [0, 0, 0, 0, 0],
}
)
# this will compile calc_in_out
calc_in_out(
df1["A"].values,
df1["B"].values,
df1["B.v"].values,
df1["in"].values,
df1["out"].values,
)
print(df1)
to_run = """calc_in_out(
df1["A"].values,
df1["B"].values,
df1["B.v"].values,
df1["in"].values,
df1["out"].values,
)"""
t = timeit(to_run, setup = "df1=setup_df()", globals=globals(), number=1)
print(t)
if __name__ == "__main__":
main()
Это печатается на моей машине (AMD 5700x):
0.023107149987481534
Вау... отличный ответ, спасибо! разница в скорости невероятная. Мне придется применить эти принципы ко всему моему коду мясника.
Почему бы вам просто не выполнить один раз цикл по данным и группировку вхождений, а затем выполнить цикл по сегментам, чтобы найти те, которые вам нужны?