2.15

code/train/utils/utils.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pandas as pd
+
+def read_and_merge_h5_data(h5_filename, key, columns, df=None):
+    """
+    读取 HDF5 文件中的数据，根据指定的 columns 筛选数据，
+    如果传入 df 参数，则将其与读取的数据根据 ts_code 和 trade_date 合并。
+
+    参数:
+    - h5_filename: HDF5 文件名
+    - key: 数据存储在 HDF5 文件中的 key
+    - columns: 要读取的列名列表
+    - df: 需要合并的 DataFrame（如果为空，则不进行合并）
+
+    返回:
+    - 合并后的 DataFrame
+    """
+    # 处理 _ 开头的列名
+    processed_columns = []
+    for col in columns:
+        if col.startswith('_'):
+            processed_columns.append(col[1:])  # 去掉下划线
+        else:
+            processed_columns.append(col)
+
+    # 从 HDF5 文件读取数据，选择需要的列
+    data = pd.read_hdf(h5_filename, key=key, columns=processed_columns)
+
+    # 修改列名，如果列名以前有 _，加上 _
+    for col in data.columns:
+        if col not in columns:  # 只有不在 columns 中的列才需要加下划线
+            new_col = f'_{col}'
+            data.rename(columns={col: new_col}, inplace=True)
+
+    # 如果传入的 df 不为空，则进行合并
+    if df is not None and not df.empty:
+        # 确保两个 DataFrame 都有 ts_code 和 trade_date 列
+        df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
+        data['trade_date'] = pd.to_datetime(data['trade_date'], format='%Y%m%d')
+
+        # 根据 ts_code 和 trade_date 合并
+        merged_df = pd.merge(df, data, on=['ts_code', 'trade_date'], how='left')
+    else:
+        # 如果 df 为空，则直接返回读取的数据
+        merged_df = data
+
+    return merged_df
+
+
+def calculate_risk_adjusted_return(df, days=1, method='ratio', lambda_=0.5, eps=1e-8):
+    """
+    计算单只股票的风险调整收益。
+
+    参数：
+    - df: DataFrame，包含 'ts_code' 和 'close' 列，按日期排序（假设 'trade_date' 已排序）。
+    - days: 预测未来多少天的收益，默认1天。
+    - method: 'ratio'（收益/波动率） 或 'difference'（收益 - λ * 波动率）。
+    - lambda_: 风险惩罚系数，仅当 method='difference' 时有效。
+    - eps: 防止除零的小常数。
+
+    返回：
+    - df：添加 'risk_adj_return' 列的 DataFrame，表示风险调整后的收益。
+    """
+    # 确保数据按 ts_code 和 trade_date 排序
+    df = df.sort_values(by=['ts_code', 'trade_date'])
+
+    # 计算未来的对数收益率
+    df['future_return'] = np.log(df.groupby('ts_code')['close'].shift(-days) / df['close'])
+
+    # 计算历史收益（对数收益率）
+    df['historical_return'] = np.log(df.groupby('ts_code')['close'].shift(1) / df['close'])
+
+    # 计算波动率（历史收益的标准差）
+    df['volatility'] = df.groupby('ts_code')['historical_return'].rolling(window=days).std().reset_index(level=0,
+                                                                                                         drop=True)
+
+    # 根据选择的 method 计算风险调整收益
+    if method == 'ratio':
+        # 收益/波动率（防止除零）
+        df['risk_adj_return'] = df['future_return'] / (df['volatility'] + eps)
+    elif method == 'difference':
+        # 收益 - λ * 波动率
+        df['risk_adj_return'] = df['future_return'] - lambda_ * df['volatility']
+    else:
+        raise ValueError("Invalid method. Use 'ratio' or 'difference'.")
+
+    return df