diff --git a/src/indicators/three_atr_ema.py b/src/indicators/three_atr_ema.py new file mode 100644 index 0000000..ab94810 --- /dev/null +++ b/src/indicators/three_atr_ema.py @@ -0,0 +1,79 @@ +import pandas as pd + +class ThreeATREMAIndicator: + def __init__(self, ema_length: int = 21): + """ + Initialize Three ATR EMA Indicator + + Args: + ema_length (int): Length for EMA calculation (default: 21) + """ + self.ema_length = ema_length + + def calculate_atr(self, high: pd.Series, low: pd.Series, close: pd.Series) -> pd.Series: + """Calculate Average True Range""" + tr1 = high - low + tr2 = abs(high - close.shift(1)) + tr3 = abs(low - close.shift(1)) + tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1) + return tr + + def calculate(self, df: pd.DataFrame) -> pd.DataFrame: + """ + Calculate Three ATR EMA indicator + + Args: + df (pd.DataFrame): DataFrame with 'high', 'low', 'close' columns + + Returns: + pd.DataFrame: DataFrame with added columns for EMA and ATR bands + """ + # Calculate EMA + ema = df['close'].ewm(span=self.ema_length, adjust=False).mean() + + # Calculate ATR + atr = self.calculate_atr(df['high'], df['low'], df['close']) + upper_band = ema + (atr * 3) + lower_band = ema - (atr * 3) + + # Generate signals based on conditions: + # 1. Price is below EMA + # 2. Price starts moving up from lower ATR band towards upper ATR band + condition = ( + df['close'] < ema & + df['close'].shift(1) <= lower_band & + df['close'] > df['close'].shift(1) + ) + + results = pd.DataFrame({ + 'ema': ema, + 'upper_band': upper_band, + 'lower_band': lower_band, + 'signal': condition + }) + + return results + + def get_signals(self, df: pd.DataFrame) -> dict: + """ + Get the current trading signals based on Three ATR EMA + + Args: + df (pd.DataFrame): DataFrame with price data + + Returns: + dict: Dictionary containing current signal and band values + """ + results = self.calculate(df) + + # Get latest values + if not results.empty: + current = results.iloc[-1] + return { + 'ema': current['ema'], + 'upper_band': current['upper_band'], + 'lower_band': current['lower_band'], + 'is_signal': current['signal'] + } + else: + return {}