liaozhaorun
36a3ccbcc8
- 添加 TabMModel、TabPFNModel 深度学习模型实现 - 新增 DataQualityAnalyzer 进行训练前数据质量诊断 - 改进数据处理器 NaN/null 双重处理，增强数据鲁棒性 - 支持 train_skip_days 参数跳过训练初期数据不足期 - Pipeline 自动清理标签为 NaN 的样本
493 lines
17 KiB
Python
493 lines
17 KiB
Python
"""NaN 问题逐步诊断测试 - 精确定位问题环节
|
||
|
||
此测试会逐步检查 DataPipeline 的每个处理步骤，精确定位 NaN 产生的位置。
|
||
"""
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
import pytest
|
||
|
||
from src.factors import FactorEngine
|
||
from src.training import (
|
||
FactorManager,
|
||
DataPipeline,
|
||
NullFiller,
|
||
Winsorizer,
|
||
StandardScaler,
|
||
)
|
||
from src.training.components.filters import STFilter
|
||
from src.experiment.common import (
|
||
SELECTED_FACTORS,
|
||
FACTOR_DEFINITIONS,
|
||
LABEL_NAME,
|
||
LABEL_FACTOR,
|
||
stock_pool_filter,
|
||
STOCK_FILTER_REQUIRED_COLUMNS,
|
||
)
|
||
|
||
# 测试配置
|
||
EXCLUDED_FACTORS = [f"GTJA_alpha{i:03d}" for i in range(1, 50)] # 排除前50个加速测试
|
||
TEST_DATE_RANGE = {
|
||
"train": ("20200101", "20201231"), # 一整年数据
|
||
"val": ("20210101", "20210331"),
|
||
"test": ("20210401", "20210630"),
|
||
}
|
||
|
||
|
||
class TestNaNStepByStep:
|
||
"""逐步诊断 NaN 问题的测试类"""
|
||
|
||
@pytest.fixture(scope="class")
|
||
def base_data(self):
|
||
"""准备基础数据（未经过任何处理）"""
|
||
print("\n" + "=" * 80)
|
||
print("[Fixture] 准备基础数据...")
|
||
|
||
engine = FactorEngine()
|
||
factor_manager = FactorManager(
|
||
selected_factors=SELECTED_FACTORS,
|
||
factor_definitions=FACTOR_DEFINITIONS,
|
||
label_factor=LABEL_FACTOR,
|
||
excluded_factors=EXCLUDED_FACTORS,
|
||
)
|
||
|
||
# 注册因子
|
||
feature_cols = factor_manager.register_to_engine(engine, verbose=False)
|
||
print(f" 特征数: {len(feature_cols)}")
|
||
|
||
# 计算完整日期范围
|
||
all_start = min(
|
||
TEST_DATE_RANGE["train"][0],
|
||
TEST_DATE_RANGE["val"][0],
|
||
TEST_DATE_RANGE["test"][0],
|
||
)
|
||
all_end = max(
|
||
TEST_DATE_RANGE["train"][1],
|
||
TEST_DATE_RANGE["val"][1],
|
||
TEST_DATE_RANGE["test"][1],
|
||
)
|
||
|
||
# 计算因子
|
||
raw_data = engine.compute(
|
||
factor_names=feature_cols + [LABEL_NAME],
|
||
start_date=all_start,
|
||
end_date=all_end,
|
||
)
|
||
print(f" 原始数据形状: {raw_data.shape}")
|
||
|
||
return {
|
||
"engine": engine,
|
||
"factor_manager": factor_manager,
|
||
"feature_cols": feature_cols,
|
||
"raw_data": raw_data,
|
||
}
|
||
|
||
def test_step_0_raw_data(self, base_data):
|
||
"""步骤0: 检查原始数据中的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤0] 检查原始数据中的 NaN")
|
||
|
||
raw_data = base_data["raw_data"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
nan_stats = self._check_nan_in_df(raw_data, feature_cols, "原始数据")
|
||
|
||
# 记录有 NaN 的列
|
||
print(f" 含 NaN 的特征列数: {len(nan_stats['cols_with_nan'])}")
|
||
if nan_stats["cols_with_nan"]:
|
||
print(f" 示例: {nan_stats['cols_with_nan'][:5]}")
|
||
|
||
return nan_stats
|
||
|
||
def test_step_1_after_st_filter(self, base_data):
|
||
"""步骤1: 检查 STFilter 后的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤1] 检查 STFilter 后的 NaN")
|
||
|
||
raw_data = base_data["raw_data"]
|
||
feature_cols = base_data["feature_cols"]
|
||
engine = base_data["engine"]
|
||
|
||
st_filter = STFilter(data_router=engine.router)
|
||
filtered_data = st_filter.filter(raw_data)
|
||
|
||
print(f" 过滤后数据形状: {filtered_data.shape}")
|
||
print(f" 删除记录数: {len(raw_data) - len(filtered_data)}")
|
||
|
||
nan_stats = self._check_nan_in_df(filtered_data, feature_cols, "STFilter后")
|
||
|
||
# 对比步骤0，看是否有新增 NaN
|
||
step0_nan = self.test_step_0_raw_data(base_data)
|
||
if nan_stats["total_nan"] != step0_nan["total_nan"]:
|
||
print(
|
||
f" [警告] NaN 数量变化: {step0_nan['total_nan']} -> {nan_stats['total_nan']}"
|
||
)
|
||
|
||
return nan_stats
|
||
|
||
def test_step_2_after_stock_pool(self, base_data):
|
||
"""步骤2: 检查股票池筛选后的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤2] 检查股票池筛选后的 NaN")
|
||
|
||
raw_data = base_data["raw_data"]
|
||
feature_cols = base_data["feature_cols"]
|
||
engine = base_data["engine"]
|
||
|
||
# 先应用 STFilter
|
||
st_filter = STFilter(data_router=engine.router)
|
||
filtered_data = st_filter.filter(raw_data)
|
||
|
||
# 再应用股票池筛选
|
||
from src.training.core.stock_pool_manager import StockPoolManager
|
||
|
||
pool_manager = StockPoolManager(
|
||
filter_func=stock_pool_filter,
|
||
required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
data_router=engine.router,
|
||
)
|
||
pool_data = pool_manager.filter_and_select_daily(filtered_data)
|
||
|
||
print(f" 筛选后数据形状: {pool_data.shape}")
|
||
print(f" 删除记录数: {len(filtered_data) - len(pool_data)}")
|
||
|
||
nan_stats = self._check_nan_in_df(pool_data, feature_cols, "股票池筛选后")
|
||
return nan_stats
|
||
|
||
def test_step_3_train_split_without_skip(self, base_data):
|
||
"""步骤3: 检查训练集划分后的 NaN（不跳过天数）"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤3] 检查训练集划分后的 NaN（不跳过天数）")
|
||
|
||
raw_data = base_data["raw_data"]
|
||
feature_cols = base_data["feature_cols"]
|
||
engine = base_data["engine"]
|
||
|
||
# 应用过滤器
|
||
st_filter = STFilter(data_router=engine.router)
|
||
filtered_data = st_filter.filter(raw_data)
|
||
|
||
from src.training.core.stock_pool_manager import StockPoolManager
|
||
|
||
pool_manager = StockPoolManager(
|
||
filter_func=stock_pool_filter,
|
||
required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
|
||
data_router=engine.router,
|
||
)
|
||
pool_data = pool_manager.filter_and_select_daily(filtered_data)
|
||
|
||
# 划分训练集
|
||
train_start, train_end = TEST_DATE_RANGE["train"]
|
||
train_mask = (pool_data["trade_date"] >= train_start) & (
|
||
pool_data["trade_date"] <= train_end
|
||
)
|
||
train_df = pool_data.filter(train_mask)
|
||
|
||
print(f" 训练集形状: {train_df.shape}")
|
||
|
||
# 统计交易日数量
|
||
unique_dates = train_df["trade_date"].unique().sort()
|
||
print(f" 训练集交易日数量: {len(unique_dates)}")
|
||
print(f" 日期范围: {unique_dates[0]} ~ {unique_dates[-1]}")
|
||
|
||
nan_stats = self._check_nan_in_df(train_df, feature_cols, "训练集(不跳过)")
|
||
|
||
# 返回训练集供后续测试使用
|
||
return {
|
||
"nan_stats": nan_stats,
|
||
"train_df": train_df,
|
||
"unique_dates": unique_dates,
|
||
}
|
||
|
||
def test_step_4_train_split_with_skip(self, base_data):
|
||
"""步骤4: 检查训练集划分后的 NaN（跳过前252天）"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤4] 检查训练集划分后的 NaN（跳过前252天）")
|
||
|
||
step3_result = self.test_step_3_train_split_without_skip(base_data)
|
||
train_df = step3_result["train_df"]
|
||
unique_dates = step3_result["unique_dates"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
# 跳过前252天
|
||
skip_days = 252
|
||
if len(unique_dates) > skip_days:
|
||
start_date = unique_dates[skip_days]
|
||
train_df_skipped = train_df.filter(pl.col("trade_date") >= start_date)
|
||
print(f" 跳过前{skip_days}天后，从 {start_date} 开始")
|
||
print(f" 跳过后训练集形状: {train_df_skipped.shape}")
|
||
print(f" 跳过记录数: {len(train_df) - len(train_df_skipped)}")
|
||
else:
|
||
train_df_skipped = train_df
|
||
print(
|
||
f" [警告] 训练集交易日数({len(unique_dates)})少于跳过天数({skip_days})，未跳过"
|
||
)
|
||
|
||
nan_stats = self._check_nan_in_df(
|
||
train_df_skipped, feature_cols, "训练集(跳过252天)"
|
||
)
|
||
|
||
return {
|
||
"nan_stats": nan_stats,
|
||
"train_df": train_df_skipped,
|
||
}
|
||
|
||
def test_step_5_after_null_filler(self, base_data):
|
||
"""步骤5: 检查 NullFiller 后的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤5] 检查 NullFiller 后的 NaN")
|
||
|
||
step4_result = self.test_step_4_train_split_with_skip(base_data)
|
||
train_df = step4_result["train_df"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
print(f" 处理前数据形状: {train_df.shape}")
|
||
|
||
# 应用 NullFiller
|
||
null_filler = NullFiller(
|
||
feature_cols=feature_cols, strategy="mean", by_date=True
|
||
)
|
||
after_null = null_filler.fit_transform(train_df)
|
||
|
||
print(f" 处理后数据形状: {after_null.shape}")
|
||
|
||
nan_stats = self._check_nan_in_df(after_null, feature_cols, "NullFiller后")
|
||
|
||
# 检查哪些列还有 NaN
|
||
if nan_stats["cols_with_nan"]:
|
||
print(
|
||
f" [错误] NullFiller 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
|
||
)
|
||
for col in nan_stats["cols_with_nan"][:10]:
|
||
count = after_null[col].null_count()
|
||
dtype = after_null[col].dtype
|
||
print(f" {col}: {count} 个 NaN, dtype={dtype}")
|
||
|
||
return {
|
||
"nan_stats": nan_stats,
|
||
"after_null": after_null,
|
||
}
|
||
|
||
def test_step_6_after_winsorizer(self, base_data):
|
||
"""步骤6: 检查 Winsorizer 后的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤6] 检查 Winsorizer 后的 NaN")
|
||
|
||
step5_result = self.test_step_5_after_null_filler(base_data)
|
||
after_null = step5_result["after_null"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
# 应用 Winsorizer
|
||
winsorizer = Winsorizer(
|
||
feature_cols=feature_cols, lower=0.01, upper=0.99, by_date=False
|
||
)
|
||
after_winsor = winsorizer.fit_transform(after_null)
|
||
|
||
nan_stats = self._check_nan_in_df(after_winsor, feature_cols, "Winsorizer后")
|
||
|
||
# 检查哪些列还有 NaN
|
||
if nan_stats["cols_with_nan"]:
|
||
print(
|
||
f" [错误] Winsorizer 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
|
||
)
|
||
for col in nan_stats["cols_with_nan"][:10]:
|
||
count = after_winsor[col].null_count()
|
||
dtype = after_winsor[col].dtype
|
||
print(f" {col}: {count} 个 NaN, dtype={dtype}")
|
||
|
||
return {
|
||
"nan_stats": nan_stats,
|
||
"after_winsor": after_winsor,
|
||
}
|
||
|
||
def test_step_7_after_standard_scaler(self, base_data):
|
||
"""步骤7: 检查 StandardScaler 后的 NaN"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤7] 检查 StandardScaler 后的 NaN")
|
||
|
||
step6_result = self.test_step_6_after_winsorizer(base_data)
|
||
after_winsor = step6_result["after_winsor"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
# 在应用 StandardScaler 之前，检查那些后来出问题的列
|
||
print("\n [预检查] StandardScaler 前，检查关键列...")
|
||
problematic_cols = [
|
||
"GTJA_alpha062",
|
||
"GTJA_alpha073",
|
||
"GTJA_alpha085",
|
||
"GTJA_alpha087",
|
||
"GTJA_alpha092",
|
||
"GTJA_alpha103",
|
||
"GTJA_alpha104",
|
||
"GTJA_alpha117",
|
||
"GTJA_alpha124",
|
||
"GTJA_alpha131",
|
||
]
|
||
for col in problematic_cols:
|
||
if col in after_winsor.columns:
|
||
null_count = after_winsor[col].null_count()
|
||
dtype = after_winsor[col].dtype
|
||
min_val = after_winsor[col].min()
|
||
max_val = after_winsor[col].max()
|
||
print(
|
||
f" {col}: null={null_count}, dtype={dtype}, min={min_val}, max={max_val}"
|
||
)
|
||
|
||
# 应用 StandardScaler
|
||
scaler = StandardScaler(feature_cols=feature_cols)
|
||
after_scaler = scaler.fit_transform(after_winsor)
|
||
|
||
# 检查 StandardScaler 学到的统计量
|
||
print("\n [统计量检查] StandardScaler 学到的统计量...")
|
||
for col in problematic_cols:
|
||
if col in scaler.mean_:
|
||
print(f" {col}: mean={scaler.mean_[col]}, std={scaler.std_[col]}")
|
||
else:
|
||
print(f" {col}: [未学到统计量]")
|
||
|
||
nan_stats = self._check_nan_in_df(
|
||
after_scaler, feature_cols, "StandardScaler后"
|
||
)
|
||
|
||
# 检查哪些列还有 NaN
|
||
if nan_stats["cols_with_nan"]:
|
||
print(
|
||
f" [错误] StandardScaler 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
|
||
)
|
||
for col in nan_stats["cols_with_nan"][:10]:
|
||
count = after_scaler[col].null_count()
|
||
dtype = after_scaler[col].dtype
|
||
print(f" {col}: {count} 个 NaN, dtype={dtype}")
|
||
|
||
# 检查这列是否学到了统计量
|
||
if col in scaler.mean_:
|
||
print(
|
||
f" mean={scaler.mean_[col]:.4f}, std={scaler.std_[col]:.4f}"
|
||
)
|
||
else:
|
||
print(f" [警告] 未学到统计量！")
|
||
|
||
return {
|
||
"nan_stats": nan_stats,
|
||
"after_scaler": after_scaler,
|
||
"scaler": scaler,
|
||
}
|
||
|
||
def test_step_8_extract_X(self, base_data):
|
||
"""步骤8: 检查提取 X 后的 NaN（转换为 numpy）"""
|
||
print("\n" + "=" * 80)
|
||
print("[步骤8] 检查提取 X 后的 NaN")
|
||
|
||
step7_result = self.test_step_7_after_standard_scaler(base_data)
|
||
after_scaler = step7_result["after_scaler"]
|
||
feature_cols = base_data["feature_cols"]
|
||
|
||
# 提取 X
|
||
X_df = after_scaler.select(feature_cols)
|
||
print(f" X DataFrame 形状: {X_df.shape}")
|
||
|
||
# 对比 DataFrame 和 select 后的 null 数量
|
||
print("\n [对比] DataFrame vs select 后的 null 数量:")
|
||
mismatched = []
|
||
for col in feature_cols[:20]: # 只检查前20个
|
||
null_in_df = after_scaler[col].null_count()
|
||
null_in_x = X_df[col].null_count()
|
||
if null_in_df != null_in_x:
|
||
mismatched.append((col, null_in_df, null_in_x))
|
||
|
||
if mismatched:
|
||
print(f" [警告] 发现 {len(mismatched)} 列不匹配:")
|
||
for col, df_null, x_null in mismatched[:10]:
|
||
print(f" {col}: DataFrame={df_null}, X={x_null}")
|
||
else:
|
||
print(f" [通过] 所有列的 null 数量一致")
|
||
|
||
# 转换为 numpy
|
||
X_np = X_df.to_numpy()
|
||
print(f"\n X numpy 形状: {X_np.shape}")
|
||
|
||
nan_count = np.isnan(X_np).sum()
|
||
print(f" X 中 NaN 总数: {nan_count}")
|
||
|
||
if nan_count > 0:
|
||
# 找出哪些列有 NaN
|
||
nan_by_col = []
|
||
for i, col in enumerate(feature_cols):
|
||
col_nan = np.isnan(X_np[:, i]).sum()
|
||
if col_nan > 0:
|
||
nan_by_col.append((col, col_nan))
|
||
|
||
print(f"\n [错误] 含 NaN 的特征列数: {len(nan_by_col)}")
|
||
for col, count in nan_by_col[:10]:
|
||
# 检查原始 DataFrame 中的情况
|
||
df_null = after_scaler[col].null_count()
|
||
dtype = after_scaler[col].dtype
|
||
|
||
# 检查是否有 Infinity
|
||
inf_count_pos = (after_scaler[col] == float("inf")).sum()
|
||
inf_count_neg = (after_scaler[col] == float("-inf")).sum()
|
||
|
||
# 检查 min/max
|
||
col_min = after_scaler[col].min()
|
||
col_max = after_scaler[col].max()
|
||
|
||
print(
|
||
f" {col}: numpy中{count}个NaN, DataFrame中{df_null}个null, dtype={dtype}"
|
||
)
|
||
print(f" min={col_min}, max={col_max}")
|
||
print(f" +inf={inf_count_pos}, -inf={inf_count_neg}")
|
||
|
||
# 如果有 inf，显示一些样本值
|
||
if inf_count_pos > 0 or inf_count_neg > 0:
|
||
sample_vals = after_scaler[col].drop_nulls().tail(5).to_list()
|
||
print(f" 样本值: {sample_vals}")
|
||
|
||
# 断言失败，显示详细信息
|
||
assert False, f"X 中含 {nan_count} 个 NaN，涉及 {len(nan_by_col)} 个特征列"
|
||
else:
|
||
print("\n [通过] X 中无 NaN！")
|
||
|
||
def _check_nan_in_df(
|
||
self, df: pl.DataFrame, feature_cols: list, step_name: str
|
||
) -> dict:
|
||
"""检查 DataFrame 中的 NaN 统计信息
|
||
|
||
Returns:
|
||
dict: {
|
||
'total_nan': 总NaN数,
|
||
'cols_with_nan': 含NaN的列名列表,
|
||
'nan_by_col': {列名: NaN数} 的字典
|
||
}
|
||
"""
|
||
nan_by_col = {}
|
||
total_nan = 0
|
||
|
||
for col in feature_cols:
|
||
null_count = df[col].null_count()
|
||
if null_count > 0:
|
||
nan_by_col[col] = null_count
|
||
total_nan += null_count
|
||
|
||
cols_with_nan = list(nan_by_col.keys())
|
||
|
||
print(f" {step_name}:")
|
||
print(f" 总记录数: {len(df)}")
|
||
print(f" 特征列数: {len(feature_cols)}")
|
||
print(f" 总NaN数: {total_nan}")
|
||
print(f" 含NaN的列数: {len(cols_with_nan)}")
|
||
|
||
if cols_with_nan and len(cols_with_nan) <= 5:
|
||
print(f" 含NaN的列: {cols_with_nan}")
|
||
elif cols_with_nan:
|
||
print(f" 含NaN的列(前5): {cols_with_nan[:5]}...")
|
||
|
||
return {
|
||
"total_nan": total_nan,
|
||
"cols_with_nan": cols_with_nan,
|
||
"nan_by_col": nan_by_col,
|
||
}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
pytest.main([__file__, "-v", "-s"])
|