make_features.py 文件源码

def calculate_indicator(stock_df):
    periods = [3, 5, 10, 20, 30, 60]
    # MA
    for period in periods:
        stock_df['MA' + str(period)] = talib.MA(stock_df['close'].values, timeperiod=period)

    # EMA
    periods = [3, 5, 10, 20, 30, 60]
    for period in periods:
        stock_df['EMA' + str(period)] = talib.EMA(stock_df['close'].values, timeperiod=period)

    # AMTMA
    periods = [5, 10, 20]
    for period in periods:
        stock_df['AMTMA' + str(period)] = talib.MA(stock_df['amount'].values, timeperiod=period)

    # ATR
    periods = [5, 10, 20]
    for period in periods:
        stock_df['ATR' + str(period)] = talib.ATR(stock_df['high'].values, stock_df['low'].values,
                                                  stock_df['close'].values, timeperiod=period)

    # ADX
    period = 14
    stock_df['ADX' + str(period)] = talib.ADX(stock_df['high'].values, stock_df['low'].values,
                                              stock_df['close'].values, timeperiod=period)

    # MACD
    stock_df['MACD_DIFF'], stock_df['MACD_DEA'], stock_df['MACD_HIST'] = talib.MACD(
        stock_df['close'].values, fastperiod=12, slowperiod=26, signalperiod=9)

    # CCI
    period = 14
    stock_df['CCI' + str(period)] = talib.CCI(stock_df['high'].values, stock_df['low'].values,
                                              stock_df['close'].values, timeperiod=period)

    # MFI
    period = 14
    stock_df['MFI' + str(period)] = talib.MFI(stock_df['high'].values, stock_df['low'].values,
                                              stock_df['close'].values, stock_df['volume'].values,
                                              timeperiod=period)

    # ROCP
    periods = [5, 10, 20]
    for period in periods:
        stock_df['ROCP' + str(period)] = talib.ROCP(stock_df['close'].values, timeperiod=period)