diff --git a/sanguo_factor/universe.py b/sanguo_factor/universe.py index 7194c91..b52f00a 100644 --- a/sanguo_factor/universe.py +++ b/sanguo_factor/universe.py @@ -11,7 +11,7 @@ from datetime import datetime, timedelta import polars as pl STOCK_PREFIXES = ("60", "00", "30") -WARMUP_BARS = 60 +WARMUP_BARS = 120 _CHUNK = 300 # 分块 IN 查询每块 symbol 数(控瞬时内存) _LOOKBACK_DAYS = 300 # start 前缓冲(覆盖最大 60 日窗口 + 节假日) _FORWARD_DAYS = 45 # end 后缓冲(覆盖 10 日前瞻收益) @@ -98,11 +98,10 @@ def load_universe_bars( def evaluation_filter(df: pl.DataFrame, start: str, end: str) -> pl.DataFrame: """IC 评估行集:窗口内 + 预热期已过.""" - # Convert datetime to string (YYYY-MM-DD) for string comparison - # This avoids polars Datetime comparison issues + from datetime import datetime as _dt return df.filter( - (pl.col("datetime").dt.strftime("%Y-%m-%d") >= pl.lit(start)) - & (pl.col("datetime").dt.strftime("%Y-%m-%d") <= pl.lit(end)) + (pl.col("datetime") >= _dt.strptime(start, "%Y-%m-%d")) + & (pl.col("datetime") <= _dt.strptime(end, "%Y-%m-%d")) & (pl.col("bar_idx") >= WARMUP_BARS) ) diff --git a/tests/factor/test_universe.py b/tests/factor/test_universe.py index 7a52427..31defc1 100644 --- a/tests/factor/test_universe.py +++ b/tests/factor/test_universe.py @@ -32,10 +32,12 @@ def _row(sym, ex, day, close, volume=100.0, turnover=None): @pytest.fixture() def db(tmp_path): + from datetime import date, timedelta rows = [] - # 600000:60根(2017-10~2017-12) + 评估窗内 3 根 - for i in range(60): - rows.append(_row("600000", "SSE", f"2017-10-{(i % 28) + 1:02d}", 10.0 + i * 0.01)) + # 600000:130根预热(2017-08~2017-12,唯一日期) + 评估窗内 3 根 + for i in range(130): + day = (date(2017, 8, 1) + timedelta(days=i)).isoformat() + rows.append(_row("600000", "SSE", day, 10.0 + i * 0.01)) for d in ("2018-01-02", "2018-01-03", "2018-01-04"): rows.append(_row("600000", "SSE", d, 11.0)) # 000001:只有 2 根(次新,bar_idx