From 2f53663c68371a3d7ae0fd11ab2801a1505d24cc Mon Sep 17 00:00:00 2001 From: Nikolay Alemasov Date: Thu, 1 Dec 2022 18:09:58 +0700 Subject: [PATCH] Introducing parallel variant of add_all_ta_features. + Add add_all_ta_features_parallel parallelizing all indicators calculation. * Fix __init__.py and README.md. --- README.md | 4 +- ta/__init__.py | 2 + ta/wrapper.py | 475 ++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 478 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 80dac080..149e9fd1 100644 --- a/README.md +++ b/README.md @@ -114,8 +114,8 @@ df = pd.read_csv('ta/tests/data/datas.csv', sep=',') df = dropna(df) # Add all ta features -df = add_all_ta_features( - df, open="Open", high="High", low="Low", close="Close", volume="Volume_BTC") +df = add_all_ta_features_parallel( + df, open="Open", high="High", low="Low", close="Close", volume="Volume_BTC", n_jobs=-1) ``` diff --git a/ta/__init__.py b/ta/__init__.py index ac190cc1..20646d69 100644 --- a/ta/__init__.py +++ b/ta/__init__.py @@ -6,6 +6,7 @@ """ from ta.wrapper import ( + add_all_ta_features_parallel, add_all_ta_features, add_momentum_ta, add_others_ta, @@ -15,6 +16,7 @@ ) __all__ = [ + "add_all_ta_features_parallel", "add_all_ta_features", "add_momentum_ta", "add_others_ta", diff --git a/ta/wrapper.py b/ta/wrapper.py index afb19e63..43f2477c 100644 --- a/ta/wrapper.py +++ b/ta/wrapper.py @@ -4,9 +4,12 @@ .. moduleauthor:: Dario Lopez Padial (Bukosabino) """ - +import numpy as np import pandas as pd +from joblib import Parallel, delayed +from typing import Dict, Tuple, List, Union + from ta.momentum import ( AwesomeOscillatorIndicator, KAMAIndicator, @@ -605,3 +608,473 @@ def add_all_ta_features( df=df, close=close, fillna=fillna, colprefix=colprefix ) return df + + +def _job(target: pd.DataFrame, obj: str, kwargs: Dict[str, Union[pd.Series, bool]], jobs: List[Tuple[str, str]]) -> Dict[str, np.ndarray]: + ohlcv_keys = ["open", "high", "low", "close", "volume"] + instance = globals()[obj](**{k: target[v] if k in ohlcv_keys else v for k, v in kwargs.items()}) + columns = {} + + for column, attr in jobs: + values: pd.Series = getattr(instance, attr)() + columns[column] = values.values + + return columns + + +def add_all_ta_features_parallel( + df: pd.DataFrame, + open: str, # noqa + high: str, + low: str, + close: str, + volume: str, + fillna: bool = False, + colprefix: str = "", + vectorized: bool = False, + n_jobs: int = -1, +) -> pd.DataFrame: + """Add all technical analysis features to dataframe. + + Args: + df (pandas.core.frame.DataFrame): Dataframe base. + open (str): Name of 'open' column. + high (str): Name of 'high' column. + low (str): Name of 'low' column. + close (str): Name of 'close' column. + volume (str): Name of 'volume' column. + fillna(bool): if True, fill nan values. + colprefix(str): Prefix column names inserted. + vectorized(bool): if True, use only vectorized functions indicators. + n_jobs(int): Number of parallel processes to spawn. + + Returns: + pandas.core.frame.DataFrame: Dataframe with new features. + """ + init = [ + # Accumulation Distribution Index + ( + "AccDistIndexIndicator", + dict(high=high, low=low, close=close, volume=volume, fillna=fillna), + [(f"{colprefix}volume_adi", "acc_dist_index")] + ), + # On Balance Volume + ( + "OnBalanceVolumeIndicator", + dict(close=close, volume=volume, fillna=fillna), + [(f"{colprefix}volume_obv", "on_balance_volume")] + ), + # Chaikin Money Flow + ( + "ChaikinMoneyFlowIndicator", + dict(high=high, low=low, close=close, volume=volume, fillna=fillna), + [(f"{colprefix}volume_cmf", "chaikin_money_flow")] + ), + # Force Index + ( + "ForceIndexIndicator", + dict(close=close, volume=volume, window=13, fillna=fillna), + [(f"{colprefix}volume_fi", "force_index")] + ), + # Ease of Movement + ( + "EaseOfMovementIndicator", + dict(high=high, low=low, volume=volume, window=14, fillna=fillna), + [ + (f"{colprefix}volume_em", "ease_of_movement"), + (f"{colprefix}volume_sma_em", "sma_ease_of_movement") + ] + ), + # Volume Price Trend + ( + "VolumePriceTrendIndicator", + dict(close=close, volume=volume, fillna=fillna), + [(f"{colprefix}volume_vpt", "volume_price_trend")] + ), + # Volume Weighted Average Price + ( + "VolumeWeightedAveragePrice", + dict(high=high, low=low, close=close, volume=volume, window=14, fillna=fillna), + [(f"{colprefix}volume_vwap", "volume_weighted_average_price")] + ), + + # Bollinger Bands + ( + "BollingerBands", + dict(close=close, window=20, window_dev=2, fillna=fillna), + [ + (f"{colprefix}volatility_bbm", "bollinger_mavg"), + (f"{colprefix}volatility_bbh", "bollinger_hband"), + (f"{colprefix}volatility_bbl", "bollinger_lband"), + (f"{colprefix}volatility_bbw", "bollinger_wband"), + (f"{colprefix}volatility_bbp", "bollinger_pband"), + (f"{colprefix}volatility_bbhi", "bollinger_hband_indicator"), + (f"{colprefix}volatility_bbli", "bollinger_lband_indicator"), + ] + ), + # Keltner Channel + ( + "KeltnerChannel", + dict(close=close, high=high, low=low, window=10, fillna=fillna), + [ + (f"{colprefix}volatility_kcc", "keltner_channel_mband"), + (f"{colprefix}volatility_kch", "keltner_channel_hband"), + (f"{colprefix}volatility_kcl", "keltner_channel_lband"), + (f"{colprefix}volatility_kcw", "keltner_channel_wband"), + (f"{colprefix}volatility_kcp", "keltner_channel_pband"), + (f"{colprefix}volatility_kchi", "keltner_channel_hband_indicator"), + (f"{colprefix}volatility_kcli", "keltner_channel_lband_indicator"), + ] + ), + # Donchian Channel + ( + "DonchianChannel", + dict(high=high, low=low, close=close, window=20, offset=0, fillna=fillna), + [ + (f"{colprefix}volatility_dcl", "donchian_channel_lband"), + (f"{colprefix}volatility_dch", "donchian_channel_hband"), + (f"{colprefix}volatility_dcm", "donchian_channel_mband"), + (f"{colprefix}volatility_dcw", "donchian_channel_wband"), + (f"{colprefix}volatility_dcp", "donchian_channel_pband"), + ] + ), + # MACD + ( + "MACD", + dict(close=close, window_slow=26, window_fast=12, window_sign=9, fillna=fillna), + [ + (f"{colprefix}trend_macd", "macd"), + (f"{colprefix}trend_macd_signal", "macd_signal"), + (f"{colprefix}trend_macd_diff", "macd_diff"), + ] + ), + # SMAs + ( + "SMAIndicator", + dict(close=close, window=12, fillna=fillna), + [(f"{colprefix}trend_sma_fast", "sma_indicator")], + ), + ( + "SMAIndicator", + dict(close=close, window=26, fillna=fillna), + [(f"{colprefix}trend_sma_slow", "sma_indicator")], + ), + # EMAs + ( + "EMAIndicator", + dict(close=close, window=12, fillna=fillna), + [(f"{colprefix}trend_ema_fast", "ema_indicator")], + ), + ( + "EMAIndicator", + dict(close=close, window=26, fillna=fillna), + [(f"{colprefix}trend_ema_slow", "ema_indicator")], + ), + # Vortex Indicator + ( + "VortexIndicator", + dict(high=high, low=low, close=close, window=14, fillna=fillna), + [ + (f"{colprefix}trend_vortex_ind_pos", "vortex_indicator_pos"), + (f"{colprefix}trend_vortex_ind_neg", "vortex_indicator_neg"), + (f"{colprefix}trend_vortex_ind_diff", "vortex_indicator_diff"), + ] + ), + # TRIX Indicator + ( + "TRIXIndicator", + dict(close=close, window=15, fillna=fillna), + [(f"{colprefix}trend_trix", "trix")], + ), + # Mass Index + ( + "MassIndex", + dict(high=high, low=low, window_fast=9, window_slow=25, fillna=fillna), + [(f"{colprefix}trend_mass_index", "mass_index")], + ), + # DPO Indicator + ( + "DPOIndicator", + dict(close=close, window=20, fillna=fillna), + [(f"{colprefix}trend_dpo", "dpo")], + ), + # KST Indicator + ( + "KSTIndicator", + dict( + close=close, + roc1=10, + roc2=15, + roc3=20, + roc4=30, + window1=10, + window2=10, + window3=10, + window4=15, + nsig=9, + fillna=fillna, + ), + [ + (f"{colprefix}trend_kst", "kst"), + (f"{colprefix}trend_kst_sig", "kst_sig"), + (f"{colprefix}trend_kst_diff", "kst_diff"), + ] + ), + # Ichimoku Indicator + ( + "IchimokuIndicator", + dict( + high=high, + low=low, + window1=9, + window2=26, + window3=52, + visual=False, + fillna=fillna, + ), + [ + (f"{colprefix}trend_ichimoku_conv", "ichimoku_conversion_line"), + (f"{colprefix}trend_ichimoku_base", "ichimoku_base_line"), + (f"{colprefix}trend_ichimoku_a", "ichimoku_a"), + (f"{colprefix}trend_ichimoku_b", "ichimoku_b"), + ] + ), + # Schaff Trend Cycle (STC) + ( + "STCIndicator", + dict( + close=close, + window_slow=50, + window_fast=23, + cycle=10, + smooth1=3, + smooth2=3, + fillna=fillna, + ), + [(f"{colprefix}trend_stc", "stc")] + ), + # Relative Strength Index (RSI) + ( + "RSIIndicator", + dict(close=close, window=14, fillna=fillna), + [(f"{colprefix}momentum_rsi", "rsi")], + ), + # Stoch RSI (StochRSI) + ( + "StochRSIIndicator", + dict(close=close, window=14, smooth1=3, smooth2=3, fillna=fillna), + [ + (f"{colprefix}momentum_stoch_rsi", "stochrsi"), + (f"{colprefix}momentum_stoch_rsi_k", "stochrsi_k"), + (f"{colprefix}momentum_stoch_rsi_d", "stochrsi_d"), + ] + ), + # TSI Indicator + ( + "TSIIndicator", + dict(close=close, window_slow=25, window_fast=13, fillna=fillna), + [(f"{colprefix}momentum_tsi", "tsi")], + ), + # Ultimate Oscillator + ( + "UltimateOscillator", + dict( + high=high, + low=low, + close=close, + window1=7, + window2=14, + window3=28, + weight1=4.0, + weight2=2.0, + weight3=1.0, + fillna=fillna, + ), + [(f"{colprefix}momentum_uo", "ultimate_oscillator")] + ), + # Stoch Indicator + ( + "StochasticOscillator", + dict( + high=high, + low=low, + close=close, + window=14, + smooth_window=3, + fillna=fillna, + ), + [ + (f"{colprefix}momentum_stoch", "stoch"), + (f"{colprefix}momentum_stoch_signal", "stoch_signal"), + ] + ), + # Williams R Indicator + ( + "WilliamsRIndicator", + dict(high=high, low=low, close=close, lbp=14, fillna=fillna), + [(f"{colprefix}momentum_wr", "williams_r")], + ), + # Awesome Oscillator + ( + "AwesomeOscillatorIndicator", + dict(high=high, low=low, window1=5, window2=34, fillna=fillna), + [(f"{colprefix}momentum_ao", "awesome_oscillator")], + ), + # Rate Of Change + ( + "ROCIndicator", + dict(close=close, window=12, fillna=fillna), + [(f"{colprefix}momentum_roc", "roc")], + ), + # Percentage Price Oscillator + ( + "PercentagePriceOscillator", + dict(close=close, window_slow=26, window_fast=12, window_sign=9, fillna=fillna), + [ + (f"{colprefix}momentum_ppo", "ppo"), + (f"{colprefix}momentum_ppo_signal", "ppo_signal"), + (f"{colprefix}momentum_ppo_hist", "ppo_hist"), + ] + ), + # Percentage Volume Oscillator + ( + "PercentageVolumeOscillator", + dict(volume=volume, window_slow=26, window_fast=12, window_sign=9, fillna=fillna), + [ + (f"{colprefix}momentum_pvo", "pvo"), + (f"{colprefix}momentum_pvo_signal", "pvo_signal"), + (f"{colprefix}momentum_pvo_hist", "pvo_hist"), + ] + ), + # Daily Return + ( + "DailyReturnIndicator", + dict(close=close, fillna=fillna), + [(f"{colprefix}others_dr", "daily_return")], + ), + # Daily Log Return + ( + "DailyLogReturnIndicator", + dict(close=close, fillna=fillna), + [(f"{colprefix}others_dlr", "daily_log_return")], + ), + # Cumulative Return + ( + "CumulativeReturnIndicator", + dict(close=close, fillna=fillna), + [(f"{colprefix}others_cr", "cumulative_return")], + ), + ] + + if not vectorized: + init += [ + # Money Flow Indicator + ( + "MFIIndicator", + dict(high=high, low=low, close=close, volume=volume, window=14, fillna=fillna), + [(f"{colprefix}volume_mfi", "money_flow_index")] + ), + # Negative Volume Index + ( + "NegativeVolumeIndexIndicator", + dict(close=close, volume=volume, fillna=fillna), + [(f"{colprefix}volume_nvi", "negative_volume_index")] + ), + # Average True Range + ( + "AverageTrueRange", + dict(close=close, high=high, low=low, window=10, fillna=fillna), + [(f"{colprefix}volatility_atr", "average_true_range")] + ), + # Ulcer Index + ( + "UlcerIndex", + dict(close=close, window=14, fillna=fillna), + [(f"{colprefix}volatility_ui", "ulcer_index")] + ), + # Average Directional Movement Index (ADX) + ( + "ADXIndicator", + dict(high=high, low=low, close=close, window=14, fillna=fillna), + [ + (f"{colprefix}trend_adx", "adx"), + (f"{colprefix}trend_adx_pos", "adx_pos"), + (f"{colprefix}trend_adx_neg", "adx_neg"), + ] + ), + # CCI Indicator + ( + "CCIIndicator", + dict( + high=high, + low=low, + close=close, + window=20, + constant=0.015, + fillna=fillna, + ), + [(f"{colprefix}trend_cci", "cci")] + ), + # Ichimoku Visual Indicator + ( + "IchimokuIndicator", + dict( + high=high, + low=low, + window1=9, + window2=26, + window3=52, + visual=True, + fillna=fillna, + ), + [ + (f"{colprefix}trend_visual_ichimoku_a", "ichimoku_a"), + (f"{colprefix}trend_visual_ichimoku_b", "ichimoku_b"), + ] + ), + # Aroon Indicator + ( + "AroonIndicator", + dict(close=close, window=25, fillna=fillna), + [ + (f"{colprefix}trend_aroon_up", "aroon_up"), + (f"{colprefix}trend_aroon_down", "aroon_down"), + (f"{colprefix}trend_aroon_ind", "aroon_indicator"), + ] + ), + # PSAR Indicator + ( + "PSARIndicator", + dict( + high=high, + low=low, + close=close, + step=0.02, + max_step=0.20, + fillna=fillna, + ), + [ + # (f"{colprefix}trend_psar", "psar"), + (f"{colprefix}trend_psar_up", "psar_up"), + (f"{colprefix}trend_psar_down", "psar_down"), + (f"{colprefix}trend_psar_up_indicator", "psar_up_indicator"), + (f"{colprefix}trend_psar_down_indicator", "psar_down_indicator"), + ] + ), + # KAMA + ( + "KAMAIndicator", + dict(close=close, window=10, pow1=2, pow2=30, fillna=fillna), + [(f"{colprefix}momentum_kama", "kama")] + ), + ] + + columns_order = df.columns.tolist() + [column for *_, jobs in init for column, _ in jobs] + + results = Parallel(n_jobs=n_jobs)(delayed(_job)(df, *args) for args in init) + + for columns in results: + for column, values in columns.items(): + df[column] = values + + return df[columns_order]