首页 > 解决方案 > 将自定义函数应用于数据框列

问题描述

有没有办法优化以下代码片段?我正在尝试使用前一行列值和自定义函数中指定的周期以及当前行列中的价格来计算当前行列的值。

import pandas as pd

class EMA_Period:
   fast = 8
   slow = 17

def calculate_ema(prev_ema, price, period):
    return prev_ema + (2.0 / (1.0 + period)) * (price - prev_ema)

times = [1578614400, 1578614700, 1578615000, 1578615300, 1578615600]
closes = [10278.6, 10276.0, 10275.6, 10274.8, 10277.0]
fast_ema = [10278.6, 0, 0, 0, 0]

df = pd.DataFrame(data={'time': times, 'close': closes, 'fast_ema': fast_ema})

df.set_index('time', inplace=True)

for i in range(1, df.shape[0]):
    df.iloc[i]['fast_ema'] = calculate_ema(df.iloc[i-1]['fast_ema'], df.iloc[i]['close'], EMA_Period.fast)

标签: pythonpandasdataframe

解决方案


谢谢@火星

def calc_ema(df, period=8, col_name='fast'):
    prev_value = df.iloc[0][col_name]
    def func2(row):
        # non local variable ==> will use pre_value from the new_fun function
        nonlocal prev_value
        prev_value = prev_value + (2.0 / (1.0 + period)) * (row['close'] - prev_value)
        return prev_value
    # This line might throw a SettingWithCopyWarning warning
    df.iloc[1:][col_name] = df.iloc[1:].apply(func2, axis=1)
    return df
df = calc_ema(df)

推荐阅读