尝试增加arbr指标

quant_strategy.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import talib
-
+from util import calculate_brar
 
 # =====================================================================
 # 1. 指标引擎：未来所有新发掘的指标，全写在这里
@@ -57,8 +57,14 @@ class FeatureEngine:
         up, mid, low = talib.BBANDS(
             self.close, timeperiod=20, nbdevup=2, nbdevdn=2, matype=0
         )
+
         self.df["bb_up"], self.df["bb_mid"], self.df["bb_low"] = up, mid, low
 
+        df_tmp = calculate_brar(self.df)
+
+        self.df["ar"] = df_tmp["ar"]
+        self.df['br'] = df_tmp["br"]
+
         self.df["percent_b"] = (self.df["close"] - low) / (up - low)
         self.df["bandwidth"] = (up - low) / mid
         self.df["amount_avg_20d"] = (
@@ -90,6 +96,11 @@ class FeatureEngine:
         )
         features["is_not_new_low_10d"] = latest["close"] > latest["low_min_10d"]
 
+        features["ar"] = latest["ar"]
+        features["prev_ar"] = prev["ar"]
+        features["br"] = latest["br"]
+        features["prev_br"] = prev["br"]
+
         history_bw = self.df["bandwidth"].iloc[-250:]
         features["bw_quantile"] = (history_bw < latest["bandwidth"]).mean()
 
@@ -113,13 +124,21 @@ class StateRuleBook:
         # 【判断是否从主升浪跌破】
         if (
             current_state == "STATE_3_MAIN_WAVE"
-            and f["close"] < f["ma10"]
-            and f["prev_close"] > f["prev_ma10"]
+            and ((f["close"] < f["ma10"]
+            and f["prev_close"] > f["prev_ma10"])
+            or ((f["prev_ar"] >= 150 and f["ar"] < 150)
+            and (f["prev_br"] >= 300 and f["br"] < 300)))
         ):
-            return (
-                "STATE_4_WAVE_END",
-                "⚠️ 主升浪确认结束！清空做T仓，准备重新激活做T。",
-            )
+            if (f["close"] < f["ma10"] and f["prev_close"] > f["prev_ma10"]):
+                return (
+                    "STATE_4_WAVE_END",
+                    "⚠️ 主升浪确认结束！清空做T仓，准备重新激活做T。(股价跌破10日均线)",
+                )
+            elif ((f["prev_ar"] >= 150 and f["ar"] < 150) and (f["prev_br"] >= 300 and f["br"] < 300)):
+                return (
+                    "STATE_4_WAVE_END",
+                    "⚠️ 主升浪确认结束！清空做T仓，准备重新激活做T。(ARBR跌破警戒线)",
+                )
 
         # 【判断是否爆发主升浪】
         if (

util.py

@@ -19,6 +19,45 @@ CN_CODE_TYPE = {
     }
 
 
+def calculate_brar(df: pd.DataFrame, N: int = 26) -> pd.DataFrame:
+    """
+    计算AR和BR情绪指标
+
+    参数:
+        df (pd.DataFrame): 包含 'high', 'low', 'open', 'close' 列的DataFrame
+        N (int): 计算周期，默认通常为26
+    返回:
+        pd.DataFrame: 包含原始数据及计算出的 'ar' 和 'br' 列的DataFrame
+    """
+    # 创建副本，避免修改原始数据
+    data = df.copy()
+
+    # --- 计算 AR (人气指标) ---
+    # 公式: AR = (N日内 (最高价 - 开盘价) 之和) / (N日内 (开盘价 - 最低价) 之和) * 100
+    high_open = data['high'] - data['open']
+    open_low = data['open'] - data['low']
+
+    data['ar'] = (high_open.rolling(window=N, min_periods=1).sum() /
+                  open_low.rolling(window=N, min_periods=1).sum()) * 100
+
+    # --- 计算 BR (意愿指标) ---
+    # 公式: BR = (N日内 MAX(0, 最高价 - 前一日收盘价) 之和) /
+    #           (N日内 MAX(0, 前一日收盘价 - 最低价) 之和) * 100
+    prev_close = data['close'].shift(1)
+
+    # 计算多头意愿：当日最高价 - 前一日收盘价，若小于0则记为0
+    high_prev_close = data['high'] - prev_close
+    high_prev_close = high_prev_close.where(high_prev_close > 0, 0)
+
+    # 计算空头意愿：前一日收盘价 - 当日最低价，若小于0则记为0
+    prev_close_low = prev_close - data['low']
+    prev_close_low = prev_close_low.where(prev_close_low > 0, 0)
+
+    data['br'] = (high_prev_close.rolling(window=N, min_periods=1).sum() /
+                  prev_close_low.rolling(window=N, min_periods=1).sum()) * 100
+
+    return data
+
 
 def get_history_k(code, start_date='19700101', end_date=None):
 
@@ -36,8 +75,7 @@ def get_history_k(code, start_date='19700101', end_date=None):
         url = "https://api.zhituapi.com/hs/history/%s/d/f?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)
 
         if end_date is not None:
-            url += "&end_date=%s"%(end_date)
-
+            url += "&et=%s"%(end_date)
         response = requests.get(url)
 
         data = response.json()