Update regime.py

regime.py

@@ -1,14 +1,19 @@
-import pandas as pd
-import numpy as np
 from typing import Dict, Any
-from config import ATR_PERIOD, STRUCTURE_LOOKBACK, VOLATILITY_EXPANSION_MULTIPLIER
+
+import numpy as np
+import pandas as pd
+
+from config import (
+    ATR_PERIOD,
+    STRUCTURE_LOOKBACK,
+    STRUCTURE_CONFIRM_BARS,
+    VOLATILITY_EXPANSION_MULT,
+    VOLATILITY_CONTRACTION_MULT,
+)
 
 
 def compute_atr(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
-    high = df["high"]
-    low = df["low"]
-    close = df["close"]
-    prev_close = close.shift(1)
+    high, low, prev_close = df["high"], df["low"], df["close"].shift(1)
     tr = pd.concat(
         [
             high - low,
@@ -17,67 +22,79 @@ def compute_atr(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
         ],
         axis=1,
     ).max(axis=1)
-    atr = tr.ewm(span=period, adjust=False).mean()
-    return atr
-
-
-def detect_structure(df: pd.DataFrame, lookback: int = STRUCTURE_LOOKBACK) -> pd.Series:
-    highs = df["high"].rolling(lookback).max()
-    lows = df["low"].rolling(lookback).min()
-    prev_highs = highs.shift(lookback // 2)
-    prev_lows = lows.shift(lookback // 2)
-    hh = highs > prev_highs
-    hl = lows > prev_lows
-    lh = highs < prev_highs
-    ll = lows < prev_lows
+    return tr.ewm(span=period, adjust=False).mean()
+
+
+def compute_structure(df: pd.DataFrame, lookback: int = STRUCTURE_LOOKBACK) -> pd.Series:
+    roll_high = df["high"].rolling(lookback).max()
+    roll_low = df["low"].rolling(lookback).min()
+    prev_roll_high = roll_high.shift(lookback // 2)
+    prev_roll_low = roll_low.shift(lookback // 2)
+
+    hh = roll_high > prev_roll_high
+    hl = roll_low > prev_roll_low
+    lh = roll_high < prev_roll_high
+    ll = roll_low < prev_roll_low
+
     structure = pd.Series(0, index=df.index)
     structure[hh & hl] = 1
     structure[lh & ll] = -1
     return structure
 
 
-def compute_volatility_regime(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
+def compute_vol_ratio(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
     atr = compute_atr(df, period)
-    atr_ma = atr.rolling(period * 2).mean()
-    ratio = atr / atr_ma.replace(0, np.nan)
-    return ratio
+    atr_ma = atr.rolling(period * 2).mean().replace(0, np.nan)
+    return atr / atr_ma
+
+
+def classify_trend(structure_series: pd.Series, lookback: int = STRUCTURE_CONFIRM_BARS) -> str:
+    recent = structure_series.iloc[-lookback:]
+    bullish = (recent == 1).sum()
+    bearish = (recent == -1).sum()
+    if bullish > bearish and bullish >= max(1, lookback // 2):
+        return "bullish"
+    if bearish > bullish and bearish >= max(1, lookback // 2):
+        return "bearish"
+    return "ranging"
 
 
 def detect_regime(df: pd.DataFrame) -> Dict[str, Any]:
-    atr = compute_atr(df, ATR_PERIOD)
-    structure = detect_structure(df, STRUCTURE_LOOKBACK)
-    vol_ratio = compute_volatility_regime(df, ATR_PERIOD)
-
-    last_atr = atr.iloc[-1]
-    last_structure = structure.iloc[-1]
-    last_vol_ratio = vol_ratio.iloc[-1]
-
-    recent_structure = structure.iloc[-5:]
-    bullish_count = (recent_structure == 1).sum()
-    bearish_count = (recent_structure == -1).sum()
-
-    if bullish_count >= 3:
-        trend = "bullish"
-    elif bearish_count >= 3:
-        trend = "bearish"
-    else:
-        trend = "ranging"
+    atr_series = compute_atr(df, ATR_PERIOD)
+    structure_series = compute_structure(df, STRUCTURE_LOOKBACK)
+    vol_ratio_series = compute_vol_ratio(df, ATR_PERIOD)
+
+    last_atr = float(atr_series.iloc[-1])
+    last_structure = int(structure_series.iloc[-1])
+    last_vol_ratio = float(vol_ratio_series.iloc[-1]) if not np.isnan(vol_ratio_series.iloc[-1]) else 1.0
+    last_close = float(df["close"].iloc[-1])
 
-    volatility_expanding = last_vol_ratio > VOLATILITY_EXPANSION_MULTIPLIER
+    trend = classify_trend(structure_series, STRUCTURE_CONFIRM_BARS)
+    vol_expanding = last_vol_ratio > VOLATILITY_EXPANSION_MULT
+    vol_contracting = last_vol_ratio < VOLATILITY_CONTRACTION_MULT
+    atr_pct = last_atr / last_close if last_close > 0 else 0.0
 
-    if trend == "bullish" and not volatility_expanding:
+    if trend == "bullish" and not vol_expanding:
         regime_score = 1.0
-    elif trend == "bullish" and volatility_expanding:
-        regime_score = 0.6
+    elif trend == "bullish" and vol_expanding:
+        regime_score = 0.55
+    elif trend == "ranging" and not vol_expanding and not vol_contracting:
+        regime_score = 0.35
     elif trend == "ranging":
-        regime_score = 0.3
-    elif trend == "bearish" and not volatility_expanding:
-        regime_score = 0.2
+        regime_score = 0.25
+    elif trend == "bearish" and not vol_expanding:
+        regime_score = 0.15
     else:
-        regime_score = 0.1
+        regime_score = 0.05
+
+    if last_structure == 1:
+        regime_score = min(1.0, regime_score + 0.1)
+    elif last_structure == -1:
+        regime_score = max(0.0, regime_score - 0.1)
 
-    close = df["close"].iloc[-1]
-    atr_pct = last_atr / close if close > 0 else 0.0
+    atr_ma_20 = atr_series.rolling(20).mean().iloc[-1]
+    atr_ma_50 = atr_series.rolling(50).mean().iloc[-1] if len(df) >= 50 else atr_ma_20
+    atr_trend = "rising" if atr_ma_20 > atr_ma_50 else "falling"
 
     return {
         "atr": last_atr,
@@ -85,9 +102,11 @@ def detect_regime(df: pd.DataFrame) -> Dict[str, Any]:
         "structure": last_structure,
         "trend": trend,
         "vol_ratio": last_vol_ratio,
-        "volatility_expanding": volatility_expanding,
-        "regime_score": regime_score,
-        "atr_series": atr,
-        "structure_series": structure,
-        "vol_ratio_series": vol_ratio,
+        "vol_expanding": vol_expanding,
+        "vol_contracting": vol_contracting,
+        "atr_trend": atr_trend,
+        "regime_score": round(float(np.clip(regime_score, 0.0, 1.0)), 4),
+        "atr_series": atr_series,
+        "structure_series": structure_series,
+        "vol_ratio_series": vol_ratio_series,
     }