data_pipeline/load_data.py

"""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")