8972d1058f
添加 compute_factors 方法到 AlphaLabSession: - 在 __init__ 添加 _loaded_symbols 和 _loaded_bars 缓存 - load_symbols 现在缓存 bar 数据供 compute_factors 使用 - compute_factors 使用缓存的 bars 调用 AlphaDataset - 使用懒导入避免本地 Python 3.14 缺少 polars/vnpy.alpha 的 ImportError 测试 (container only): - test_compute_factors_calls_prepare_and_fetch 验证 AlphaDataset 调用 Co-Authored-By: Claude <noreply@anthropic.com>
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
"""AlphaLab session management - lazy import vnpy.alpha to avoid ImportError when alphalens missing."""
|
|
import sys
|
|
import os
|
|
_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0"))
|
|
if _VNPY_SRC not in sys.path:
|
|
sys.path.insert(0, _VNPY_SRC)
|
|
|
|
|
|
class AlphaLabSession:
|
|
"""Session manager for vnpy.alpha AlphaLab operations."""
|
|
|
|
def __init__(self, lab_path: str):
|
|
"""
|
|
Initialize AlphaLab session.
|
|
|
|
Args:
|
|
lab_path: Path to AlphaLab directory
|
|
"""
|
|
from vnpy.alpha.lab import AlphaLab
|
|
|
|
self.lab_path = lab_path
|
|
self.lab = AlphaLab(lab_path)
|
|
self._loaded_symbols: list[str] = []
|
|
self._loaded_bars: dict[str, list] = {}
|
|
|
|
def load_symbols(self, symbols: list[str], start: str, end: str, cfg) -> None:
|
|
"""
|
|
Load symbol data from database and save to AlphaLab.
|
|
|
|
Args:
|
|
symbols: List of vt_symbols to load
|
|
start: Start date (YYYY-MM-DD)
|
|
end: End date (YYYY-MM-DD)
|
|
cfg: Database configuration object
|
|
"""
|
|
from sanguo_data.datareader import read_db_daily
|
|
from .data_adapter import save_alpha_lab_data
|
|
|
|
for symbol in symbols:
|
|
bars = read_db_daily(symbol, start, end, cfg)
|
|
if bars:
|
|
save_alpha_lab_data(bars, self.lab_path)
|
|
# Cache bars for compute_factors
|
|
if symbol not in self._loaded_symbols:
|
|
self._loaded_symbols.append(symbol)
|
|
self._loaded_bars[symbol] = bars
|
|
|
|
def compute_factors(self, factor_names: list[str], train_period: tuple, valid_period: tuple, test_period: tuple):
|
|
"""
|
|
Compute factors using cached bars and vnpy.alpha AlphaDataset.
|
|
|
|
Args:
|
|
factor_names: List of factor names to compute
|
|
train_period: Training period tuple (start, end)
|
|
valid_period: Validation period tuple (start, end)
|
|
test_period: Test period tuple (start, end)
|
|
|
|
Returns:
|
|
polars DataFrame with computed factors for test period
|
|
"""
|
|
# Lazy imports to avoid ImportError on local Python 3.14 without polars/vnpy.alpha
|
|
import polars as pl
|
|
from vnpy.alpha.dataset import AlphaDataset, Segment
|
|
from .registry import get_factor
|
|
from .data_adapter import convert_bars_to_alpha_df
|
|
|
|
# Gather all cached bars across loaded symbols
|
|
all_bars = []
|
|
for symbol in self._loaded_symbols:
|
|
all_bars.extend(self._loaded_bars.get(symbol, []))
|
|
|
|
# Convert bars to AlphaLab DataFrame format
|
|
df = convert_bars_to_alpha_df(all_bars)
|
|
|
|
# Create AlphaDataset with the specified periods
|
|
ds = AlphaDataset(df, train_period, valid_period, test_period)
|
|
|
|
# Add each factor to the dataset
|
|
for name in factor_names:
|
|
factor = get_factor(name)
|
|
if factor is None:
|
|
continue # Skip unknown factors
|
|
ds.add_feature(name, factor["expression"])
|
|
|
|
# Prepare data (compute features)
|
|
ds.prepare_data(max_workers=1)
|
|
|
|
# Return test period data
|
|
return ds.fetch_raw(Segment.TEST)
|