Update indicater.py

indicater.py

@@ -3,76 +3,70 @@ import pandas as pd
 import numpy as np
 import talib
 
-# -------------------------------
-# Custom SuperTrend implementation
-# -------------------------------
-def supertrend(df, period=10, multiplier=3):
-    """
-    Compute SuperTrend
-    """
-    hl2 = (df['High'] + df['Low']) / 2
-    atr = talib.ATR(df['High'], df['Low'], df['Close'], timeperiod=period)
-    
-    upperband = hl2 + (multiplier * atr)
-    lowerband = hl2 - (multiplier * atr)
-    supertrend = pd.Series(index=df.index, dtype=float)
-    direction = pd.Series(index=df.index, dtype=int)
-    
-    for i in range(len(df)):
-        if i == 0:
-            supertrend.iloc[i] = upperband.iloc[i]
-            direction.iloc[i] = 1
-            continue
-        if df['Close'].iloc[i] > supertrend.iloc[i-1]:
-            direction.iloc[i] = 1
-            supertrend.iloc[i] = lowerband.iloc[i]
-        elif df['Close'].iloc[i] < supertrend.iloc[i-1]:
-            direction.iloc[i] = -1
-            supertrend.iloc[i] = upperband.iloc[i]
-        else:
-            direction.iloc[i] = direction.iloc[i-1]
-            supertrend.iloc[i] = supertrend.iloc[i-1]
-    return supertrend
-
-# -------------------------------
-# Main indicator calculation
-# -------------------------------
 def calculate_indicators(df):
     """
-    Compute major TA-Lib indicators and custom indicators.
-    Return dictionary {indicator_name: series_or_df}
+    Calculate various indicators.
+    df: OHLCV dataframe with columns: Open, High, Low, Close, Volume
+    Returns dict of indicator name -> DataFrame
     """
     indicators = {}
 
-    # --- Price-based moving averages ---
-    indicators['SMA5'] = talib.SMA(df['Close'], timeperiod=5)
-    indicators['SMA20'] = talib.SMA(df['Close'], timeperiod=20)
-    indicators['SMA50'] = talib.SMA(df['Close'], timeperiod=50)
-    indicators['SMA200'] = talib.SMA(df['Close'], timeperiod=200)
-    
-    indicators['EMA5'] = talib.EMA(df['Close'], timeperiod=5)
-    indicators['EMA20'] = talib.EMA(df['Close'], timeperiod=20)
-    indicators['EMA50'] = talib.EMA(df['Close'], timeperiod=50)
-    indicators['EMA200'] = talib.EMA(df['Close'], timeperiod=200)
+    close = df['Close']
+    high = df['High']
+    low = df['Low']
+    volume = df['Volume'] if 'Volume' in df else pd.Series([1]*len(df), index=df.index)
+
+    # --- TA-Lib indicators ---
+    try:
+        # Moving averages
+        indicators['SMA20'] = pd.DataFrame({'SMA20': talib.SMA(close, timeperiod=20)}, index=df.index)
+        indicators['SMA50'] = pd.DataFrame({'SMA50': talib.SMA(close, timeperiod=50)}, index=df.index)
+        indicators['EMA20'] = pd.DataFrame({'EMA20': talib.EMA(close, timeperiod=20)}, index=df.index)
+        indicators['EMA50'] = pd.DataFrame({'EMA50': talib.EMA(close, timeperiod=50)}, index=df.index)
+
+        # MACD
+        macd, macdsignal, macdhist = talib.MACD(close)
+        indicators['MACD'] = pd.DataFrame({'MACD': macd, 'Signal': macdsignal, 'Hist': macdhist}, index=df.index)
 
-    # --- Momentum indicators ---
-    indicators['RSI'] = talib.RSI(df['Close'], timeperiod=14)
-    indicators['STOCH'], indicators['STOCH_signal'] = talib.STOCH(
-        df['High'], df['Low'], df['Close'], fastk_period=14, slowk_period=3, slowk_matype=0,
-        slowd_period=3, slowd_matype=0
-    )
-    indicators['MACD'], indicators['MACD_signal'], indicators['MACD_hist'] = talib.MACD(df['Close'])
-    indicators['ADX'] = talib.ADX(df['High'], df['Low'], df['Close'], timeperiod=14)
-    indicators['CCI'] = talib.CCI(df['High'], df['Low'], df['Close'], timeperiod=14)
-    indicators['OBV'] = talib.OBV(df['Close'], df['Volume'])
+        # RSI
+        indicators['RSI14'] = pd.DataFrame({'RSI14': talib.RSI(close, timeperiod=14)}, index=df.index)
 
-    # --- Volatility indicators ---
-    indicators['ATR'] = talib.ATR(df['High'], df['Low'], df['Close'], timeperiod=14)
-    indicators['BB_upper'], indicators['BB_middle'], indicators['BB_lower'] = talib.BBANDS(
-        df['Close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0
-    )
+        # Bollinger Bands
+        upper, middle, lower = talib.BBANDS(close)
+        indicators['BB_upper'] = pd.DataFrame({'BB_upper': upper}, index=df.index)
+        indicators['BB_middle'] = pd.DataFrame({'BB_middle': middle}, index=df.index)
+        indicators['BB_lower'] = pd.DataFrame({'BB_lower': lower}, index=df.index)
 
-    # --- Custom indicators ---
-    indicators['SuperTrend'] = supertrend(df)
+        # ADX
+        indicators['ADX14'] = pd.DataFrame({'ADX14': talib.ADX(high, low, close, timeperiod=14)}, index=df.index)
+    except Exception as e:
+        print("TA-Lib indicators error:", e)
+
+    # --- Custom indicators if not in TA-Lib ---
+    # SuperTrend
+    indicators['SuperTrend'] = calculate_supertrend(df)
 
     return indicators
+
+def calculate_supertrend(df, period=10, multiplier=3):
+    """
+    Basic SuperTrend calculation
+    """
+    hl2 = (df['High'] + df['Low']) / 2
+    atr = talib.ATR(df['High'], df['Low'], df['Close'], timeperiod=period)
+    st_upper = hl2 + multiplier * atr
+    st_lower = hl2 - multiplier * atr
+    supertrend = pd.Series(index=df.index, dtype=float)
+    trend = True  # True = up, False = down
+
+    for i in range(len(df)):
+        if i == 0:
+            supertrend.iloc[i] = st_upper.iloc[i]
+        else:
+            if df['Close'].iloc[i] > supertrend.iloc[i-1]:
+                trend = True
+                supertrend.iloc[i] = st_lower.iloc[i]
+            else:
+                trend = False
+                supertrend.iloc[i] = st_upper.iloc[i]
+    return pd.DataFrame({'SuperTrend': supertrend}, index=df.index)