Rank2

main/utils/indicators.py

@@ -0,0 +1,130 @@
+import pandas as pd
+import numpy as np
+
+def calculate_sharpe_sortino_labels_efficient(
+    df: pd.DataFrame,
+    N_days: int = 5,  # 未来考察的天数，用于计算收益率序列
+    risk_free_rate_annual: float = 0.02, # 年化无风险利率 (例如 0.02 表示 2%)
+    annualization_factor: float = np.sqrt(252), # 年化因子，日数据通常用sqrt(252)
+    min_periods_for_std: int = 2 # 计算标准差所需的最小有效数据点数
+) -> pd.DataFrame:
+    """
+    高效计算每只股票在每个交易日未来 N 日的年化夏普比率和索提诺比率，作为训练的Label。
+    函数会获取未来 N 天的每日收益率序列，并基于此序列计算夏普/索提诺。
+
+    Args:
+        df (pd.DataFrame): 输入 DataFrame，需包含 'ts_code', 'trade_date', 'close' 列。
+                           'close' 列必须是股票收盘价。
+        N_days (int): 未来考察的天数。例如，N_days=5 表示考察未来5个交易日。
+        risk_free_rate_annual (float): 年化无风险利率。
+        annualization_factor (float): 用于将日收益率转化为年化收益率的因子 (例如 np.sqrt(252))。
+        min_periods_for_std (int): 计算标准差所需的最小有效数据点数。
+                                   如果未来N天有效数据少于此值，则 Label 为 NaN。
+
+    Returns:
+        pd.DataFrame: 原始DataFrame，新增 'sharpe_ratio_label' 和 'sortino_ratio_label' 列。
+                      注意：数据末尾 N_days 行的 Label 将为 NaN。
+    """
+    df_copy = df
+    
+    # 确保数据已按股票和日期排序，这对于高效计算至关重要
+    df_copy = df_copy.sort_values(by=['ts_code', 'trade_date'])
+
+    # # 计算每日收益率
+    # df_copy['pct_chg'] = df_copy.groupby('ts_code')['close'].pct_change()
+
+    # 将年化无风险利率转换为每日无风险利率
+    daily_risk_free_rate = risk_free_rate_annual / 252 # 假设每年252个交易日
+
+    def _calculate_metrics_for_window(returns_series):
+        """
+        辅助函数：计算单个滚动窗口内的夏普和索提诺比率。
+        这个函数将被 apply 调用于每个 (N_days) 收益率序列。
+        """
+        if len(returns_series.dropna()) < min_periods_for_std:
+            return np.nan, np.nan # 数据不足，返回NaN
+
+        excess_returns = returns_series - daily_risk_free_rate
+
+        # --- 夏普比率 ---
+        mean_excess_return = excess_returns.mean()
+        std_dev_returns = returns_series.std() # 夏普比率使用总收益率的标准差
+
+        sharpe = np.nan
+        if std_dev_returns != 0:
+            sharpe = (mean_excess_return / std_dev_returns) * annualization_factor
+
+        # --- 索提诺比率 ---
+        downside_returns = returns_series[returns_series < daily_risk_free_rate]
+        sortino = np.nan
+        if not downside_returns.empty:
+            downside_deviation = np.sqrt(np.mean((downside_returns - daily_risk_free_rate)**2))
+            if downside_deviation != 0:
+                sortino = (mean_excess_return / downside_deviation) * annualization_factor
+
+        return sharpe, sortino
+
+    # 使用 groupby().rolling() 结合 apply 来高效计算
+    # 注意: 这里使用 shift(-N_days) 来获取未来 N_days 的窗口
+    # 同时 rolling 窗口需要是 N_days 大小，并对 pct_chg 序列进行操作
+    # 结果会被放置在窗口的末尾，但我们实际上需要它对应到窗口的起始位置
+    # 因此，我们先计算，然后将其向上 shift N_days + 1 （或者根据 rolling 的行为调整）
+    
+    # 获取未来 N_days 的每日收益率序列
+    # 为了避免在每一行中循环，我们构造一个辅助的DataFrame，其中包含未来N天的收益率
+    # 这里使用一个更高效的思路：对每只股票，生成一个N天的收益率列表
+    
+    # 使用 groupby().apply() 和 list comprehension 来构建未来N天的收益率列表
+    # 为了获得未来N天的收益率，需要对pct_chg进行shift(-N_days)
+    # 然后再对每个N_days的窗口进行处理
+    
+    # 方法一：先获取未来N天的收益率列，然后进行滑动窗口计算 （更推荐，效率高）
+    # 创建一个辅助列，表示未来 N 天的每日收益率列表
+    # 注意：这里需要 N_days 长度的窗口，且数据点至少为 min_periods_for_std
+    
+    # 对每只股票进行分组，然后计算未来 N 天的夏普/索提诺
+    # 这仍然涉及到对每个时间点获取未来N天的收益序列，Pandas 的 rolling 默认是向后看的
+    # 为了实现向后看 N_days (即从当前日期 t 预测 t+1 到 t+N_days)，我们需要一些技巧
+    
+    # 最直接且通常高效的方法是：
+    # 1. 计算日收益率。
+    # 2. 对每个股票，对日收益率进行反向滚动，然后计算指标。
+    #    或者更简单地，使用 shift(-N_days) 创建“未来的”日收益率列，然后进行滚动。
+    
+    # 为了简化且保持效率，我们可以这样做：
+    # 1. 计算每个股票的每日收益率。
+    # 2. 对于每个 (ts_code, trade_date) 行，我们将查看其未来的 N_days 每日收益率。
+    #    这在 Pandas 中不直接通过 .rolling() 实现，因为 rolling 是“当前及之前”的窗口。
+    #    我们可以用 groupby().apply() 配合自定义函数来模拟这个“未来窗口”的行为。
+
+    def _apply_metrics_per_stock(stock_df_group):
+        """
+        对单只股票的数据帧进行处理，计算其每个日期的未来 N 日夏普/索提诺比率。
+        """
+        sharpe_labels = [np.nan] * len(stock_df_group)
+        sortino_labels = [np.nan] * len(stock_df_group)
+
+        # 遍历当前股票的每个日期，获取未来的 N_days 收益率序列
+        for i in range(len(stock_df_group) - N_days):
+            # 获取未来 N_days 的每日收益率序列
+            # 从当前日期 t 的下一天 (t+1) 开始，到 t+N_days
+            future_pct_chgs = stock_df_group['pct_chg'].iloc[i+1 : i+1+N_days].dropna()
+            
+            # 计算夏普和索提诺
+            sharpe, sortino = _calculate_metrics_for_window(future_pct_chgs)
+            
+            # 赋值到对应日期
+            sharpe_labels[i] = sharpe
+            sortino_labels[i] = sortino
+        
+        stock_df_group['sharpe_ratio_label'] = sharpe_labels
+        stock_df_group['sortino_ratio_label'] = sortino_labels
+        return stock_df_group
+
+    # 对每个股票分组应用这个函数
+    df_result = df_copy.groupby('ts_code', group_keys=False).apply(_apply_metrics_per_stock)
+
+    # # 清理不再需要的中间列
+    # df_result = df_result.drop(columns=['pct_chg'])
+
+    return df_result