Оптимизация кода Python с использованием Numpy и Pandas

Question

У меня работает следующий код:

import numpy as np
import pandas as pd
colum1 = [0.05,0.05,0.05,0.05,0.05,0.05,0.05,0.05,0.05,0.05,0.05,0.05]
colum2 = [1,2,3,4,5,6,7,8,9,10,11,12]
colum3 = [0.85,0.80,0.80,0.80,0.85,0.0,0.0,0.0,0.0,0.0,0.0,0.0]
colum4 = [1743.85, 1485.58, 1250.07, 1021.83, 818.96, 628.05, 455.40, 319.03, 190.86 , 97.07, 26.96 , 0.00]
df = pd.DataFrame({
    'colum1' : colum1,
    'colum2' : colum2,
    'colum3' : colum3,
    'colum4' : colum4,
});

df['result'] = 0
for i in range(len(colum2)):
    df['result'] = np.where(
        df['colum2'] <= 5,
        np.where(
            df['colum2'] == 1,
            df['colum4'],
            np.where(
                ( df['colum4'] - (df['result'].shift(1) * (df['colum1'] * df['colum3'])) )>0,
                ( df['colum4'] - (df['result'].shift(1) * (df['colum1'] * df['colum3'])) ),
                0
            )
        ),
        np.where(
            ( df['colum4'] - (df['result'].shift(1) * df['colum1']) )>0,
            ( df['colum4'] - (df['result'].shift(1) * df['colum1']) ),
            0
        )
    )

, и мне нужно выполнить ту же операцию, не прибегая к циклу for.Это было бы очень полезно, так как я работаю с тысячами записей, что очень медленно.

Мой ожидаемый результат следующий:

    colum1  colum2  colum3   colum4       result
0     0.05       1    0.85  1743.85  1743.850000
1     0.05       2    0.80  1485.58  1415.826000
2     0.05       3    0.80  1250.07  1193.436960
3     0.05       4    0.80  1021.83   974.092522
4     0.05       5    0.85   818.96   777.561068
5     0.05       6    0.00   628.05   589.171947
6     0.05       7    0.00   455.40   425.941403
7     0.05       8    0.00   319.03   297.732930
8     0.05       9    0.00   190.86   175.973354
9     0.05      10    0.00    97.07    88.271332
10    0.05      11    0.00    26.96    22.546433
11    0.05      12    0.00     0.00     0.000000

Answer 1

Первым шагом является удаление цикла над индексом и замена этих тестов для чисел больше 0 на np.maximum.Это работает, потому что np.where(a > 0, a, 0) для наших целей эквивалентно np.maximum(0, a).

. В то же время задайте более длинные выражения отдельно, чтобы сделать ваш код читабельным:

s1 = df['colum4'] - (df['result'].shift(1) * (df['colum1'] * df['colum3']))
s2 = df['colum4'] - (df['result'].shift(1) * df['colum1'])

df['result'] = np.where(df['colum2'] <= 5,
                        np.where(df['colum2'] == 1, df['colum4'],
                                 np.maximum(0, s1)),
                        np.maximum(0, s2))

Следующий шагиспользовать np.select для удаления вложенных операторов np.where:

m1 = df['colum2'] <= 5
m2 = df['colum2'] == 1

conds = [m1 & m2, m1 & ~m2]
choices = [df['colum4'], np.maximum(0, s1)]

df['result'] = np.select(conds, choices, np.maximum(0, s2))

Эта версия будет более управляемой.