File size: 1,585 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """Load market features from qlib."""
from __future__ import annotations
from typing import Iterable
import pandas as pd
from qlib.data import D
QLIB_TO_GP_FIELD = {
"$open": "开盘价",
"$close": "收盘价",
"$high": "最高价",
"$low": "最低价",
"$volume": "成交量",
"$vwap": "vwap",
"$amount": "成交额",
}
def load_instruments(market: str = "csi300") -> dict:
return D.instruments(market)
def load_market_features(
instruments,
fields: Iterable[str],
start_time: str,
end_time: str,
freq: str = "day",
rename_for_gp: bool = True,
) -> pd.DataFrame:
"""
Load OHLCV features via qlib D.features and optionally rename columns
to match the GP factor engine's Chinese field names.
"""
df = D.features(
instruments,
list(fields),
start_time=start_time,
end_time=end_time,
freq=freq,
)
df = df.sort_index()
if rename_for_gp:
rename_map = {k: v for k, v in QLIB_TO_GP_FIELD.items() if k in df.columns}
df = df.rename(columns=rename_map)
return df
def pivot_to_symbol_time(df: pd.DataFrame, field: str) -> pd.DataFrame:
"""Pivot multi-index (instrument, datetime) frame to symbols x time."""
if field not in df.columns:
raise KeyError(f"Field {field} not in dataframe columns: {df.columns.tolist()}")
series = df[field]
if not isinstance(series.index, pd.MultiIndex):
raise ValueError("Expected MultiIndex with instrument and datetime levels")
return series.unstack(level="datetime")
|