优化函数 python dataframe

Question

我有 supertrend 实现的 python 代码。我正在使用 pandas 数据框。代码工作正常，但是随着数据帧长度的增加，超级趋势函数运行得越来越慢。我想知道是否可以在代码中更改任何内容来优化它并使其运行得更快，即使数据帧长度很大。

def trueRange(df):
    df['prevClose'] = df['close'].shift(1)
    df['high-low'] = df['high'] - df['low']
    df['high-pClose'] = abs(df['high'] - df['prevClose'])
    df['low-pClose'] = abs(df['low'] - df['prevClose'])
    tr = df[['high-low','high-pClose','low-pClose']].max(axis=1)
    
    return tr

def averageTrueRange(df, peroid=12):
    df['trueRange'] = trueRange(df)
    the_atr = df['trueRange'].rolling(peroid).mean()
    
    return the_atr
    

def superTrend(df, peroid=5, multipler=1.5):
    df['averageTrueRange'] = averageTrueRange(df, peroid=peroid)
    h2 = ((df['high'] + df['low']) / 2)
    df['Upperband'] = h2 + (multipler * df['averageTrueRange'])
    df['Lowerband'] = h2 - (multipler * df['averageTrueRange'])
    df['inUptrend'] = None

    for current in range(1,len(df.index)):
        prev = current- 1
        
        if df['close'][current] > df['Upperband'][prev]:
            df['inUptrend'][current] = True
            
        elif df['close'][current] < df['Lowerband'][prev]:
            df['inUptrend'][current] = False
        else:
            df['inUptrend'][current] = df['inUptrend'][prev]
            
            if df['inUptrend'][current] and df['Lowerband'][current] < df['Lowerband'][prev]:
                df['Lowerband'][current] = df['Lowerband'][prev]
                
            if not df['inUptrend'][current] and df['Upperband'][current] > df['Upperband'][prev]:
                df['Upperband'][current] = df['Upperband'][prev]

矢量版本

def superTrend(df, peroid=5, multipler=1.5):
    df['averageTrueRange'] = averageTrueRange(df, peroid=peroid)
    h2 = ((df['high'] + df['low']) / 2)
    df['Upperband'] = h2 + (multipler * df['averageTrueRange'])
    df['Lowerband'] = h2 - (multipler * df['averageTrueRange'])
    df['inUptrend'] = None


    cond1 = df['close'].values[1:] > df['Upperband'].values[:-1]
    cond2 = df['close'].values[1:] < df['Lowerband'].values[:-1]

    df.loc[cond1, 'inUptrend'] = True
    df.loc[cond2, 'inUptrend'] = False

    df.loc[(~cond1) & (cond2), 'inUptrend'] = df['inUptrend'][:-1]
    df.loc[(~cond1) & (cond2) & (df['inUptrend'].values[1:] == True) & (df['Lowerband'].values[1:] < df['Lowerband'].values[:-1]), 'Lowerband'] = df['Lowerband'][:-1]
    df.loc[(~cond1) & (cond2) & (df['inUptrend'].values[1:] == False) & (df['Upperband'].values[1:] > df['Upperband'].values[:-1]), 'Upperband'] = df['Upperband'][:-1]
   

Answer 1

import pandas as pd尝试使用Modin代替。Modin 会自动让 pandas 变得更快。做就是了import modin.pandas as pd。除了导入之外，您不需要更改任何代码。

如果您需要使用该df.apply()方法，有一个名为Swifter的包。在你之后pip install swifter，你需要做的就是import swifter，然后不是做df.apply()，而是做df.swifter.apply()。方便的是 Swifter 还可以与 Modin 配合使用。