mlquantx 0.2.5__py3-none-any.whl

mlquant/data/akshare_loader.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import time
+from typing import Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+import torch
+
+from .panel import Panel
+
+
+def load_akshare_panel(
+    tickers: Sequence[str],
+    start: str,
+    end: str,
+    device: Union[str, torch.device] = "cpu",
+    adjust: str = "qfq",
+    proxy_vwap: bool = True,
+    retries: int = 3,
+    retry_sleep: float = 0.75,
+) -> Panel:
+    """Download A-share daily data from akshare and return a Panel.
+
+    Parameters
+    ----------
+    tickers : Sequence[str]
+        A-share ticker codes without exchange prefixes (for example, "600519").
+    start : str
+        Start date in "YYYY-MM-DD" format.
+    end : str
+        End date in "YYYY-MM-DD" format.
+    device : str or torch.device
+        Where to allocate the resulting tensors.
+    adjust : str
+        Price adjustment mode: "qfq", "hfq", or an empty string.
+    proxy_vwap : bool
+        If True, estimate VWAP as (Open + Close + High + Low) / 4.
+    retries : int
+        Number of attempts for each ticker before reporting a fetch failure.
+    retry_sleep : float
+        Base seconds to wait between retry attempts.
+    """
+    import akshare
+
+    if not tickers:
+        raise ValueError("Tickers list cannot be empty")
+
+    # 按输入顺序去重，避免下载和透视时产生重复列
+    unique_tickers = list(dict.fromkeys(tickers))
+
+    if not proxy_vwap:
+        raise ValueError(
+            "akshare cannot provide an adjustment-consistent raw VWAP; use proxy_vwap=True"
+        )
+
+    column_map = {
+        "日期": "date",
+        "开盘": "open",
+        "收盘": "close",
+        "最高": "high",
+        "最低": "low",
+        "成交量": "volume",
+        "成交额": "amount",
+    }
+    data_columns = ["date", "open", "high", "low", "close", "volume", "amount"]
+
+    all_data = []
+    for idx, ticker in enumerate(unique_tickers, start=1):
+        if idx == 1 or idx % 25 == 0 or idx == len(unique_tickers):
+            print(f"AkShare download progress: {idx}/{len(unique_tickers)} ({ticker})", flush=True)
+        ticker_df = None
+        last_error: Optional[Exception] = None
+        for attempt in range(1, max(1, retries) + 1):
+            try:
+                ticker_df = akshare.stock_zh_a_hist(
+                    symbol=ticker,
+                    period="daily",
+                    start_date=start.replace("-", ""),
+                    end_date=end.replace("-", ""),
+                    adjust=adjust,
+                )
+                break
+            except Exception as e:
+                last_error = e
+                if attempt < max(1, retries):
+                    time.sleep(retry_sleep * attempt)
+
+        try:
+            if ticker_df is None and last_error is not None:
+                raise last_error
+            if ticker_df is None or ticker_df.empty:
+                continue
+
+            ticker_df = ticker_df.rename(columns=column_map)
+            missing_columns = [column for column in data_columns if column not in ticker_df.columns]
+            if missing_columns:
+                raise ValueError(f"missing columns: {missing_columns}")
+
+            ticker_df = ticker_df[data_columns].copy()
+            ticker_df["code"] = ticker
+            all_data.append(ticker_df)
+        except Exception as e:
+            print(f"Warning: Failed to fetch {ticker}: {e}")
+            continue
+
+    if not all_data:
+        raise ValueError(
+            f"No data returned for tickers {unique_tickers} from {start} to {end}"
+        )
+
+    df = pd.concat(all_data, ignore_index=True)
+
+    # 将行情字段统一转为数值类型
+    num_cols = ["open", "high", "low", "close", "volume", "amount"]
+    for col in num_cols:
+        df[col] = pd.to_numeric(df[col], errors="coerce")
+
+    df["date"] = pd.to_datetime(df["date"])
+
+    # 将长表转换为日期乘股票的宽表
+    fields_wide = {}
+
+    def get_wide_df(col_name):
+        wide = df.pivot(index="date", columns="code", values=col_name)
+        # 为所有请求的股票补齐列，并保持原始顺序
+        for ticker in unique_tickers:
+            if ticker not in wide.columns:
+                wide[ticker] = np.nan
+        return wide[list(unique_tickers)]
+
+    open_df = get_wide_df("open")
+    high_df = get_wide_df("high")
+    low_df = get_wide_df("low")
+    close_df = get_wide_df("close")
+    volume_df = get_wide_df("volume")
+    amount_df = get_wide_df("amount")
+
+    fields_wide["open"] = open_df
+    fields_wide["high"] = high_df
+    fields_wide["low"] = low_df
+    fields_wide["close"] = close_df
+    fields_wide["volume"] = volume_df
+    fields_wide["amount"] = amount_df
+    fields_wide["vwap"] = (open_df + close_df + high_df + low_df) / 4.0
+
+    dates = open_df.index.to_numpy()
+    stocks = list(unique_tickers)
+
+    mask = (
+        ~open_df.isna()
+        & ~high_df.isna()
+        & ~low_df.isna()
+        & ~close_df.isna()
+        & ~volume_df.isna()
+    ).to_numpy()
+
+    tensors = {
+        name: torch.from_numpy(df_.fillna(0.0).to_numpy(dtype=np.float32).copy()).to(device)
+        for name, df_ in fields_wide.items()
+    }
+
+    return Panel.from_tensors(
+        dates=dates,
+        stocks=stocks,
+        fields=tensors,
+        mask=torch.from_numpy(mask.copy()).to(device),
+    )

mlquant/data/baostock_loader.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import torch
+from typing import Sequence, Optional, Union
+
+from .panel import Panel
+
+def load_baostock_panel(
+    tickers: Sequence[str],
+    start: str,
+    end: str,
+    device: Union[str, torch.device] = "cpu",
+    proxy_vwap: bool = True
+) -> Panel:
+    """Download data from baostock and return a Panel.
+
+    Parameters
+    ----------
+    tickers : Sequence[str]
+        List of tickers to download. For A-shares, use baostock format (e.g., 'sh.600000', 'sz.000001').
+    start : str
+        Start date (e.g. "2020-01-01").
+    end : str
+        End date (e.g. "2023-12-31").
+    device : str or torch.device
+        Where to allocate the resulting tensors.
+    proxy_vwap : bool
+        If True, estimates VWAP as (Open + Close + High + Low) / 4.
+        Baostock does not expose an adjustment-consistent raw VWAP, so False
+        is rejected rather than mixing adjusted OHLC with unadjusted turnover.
+    """
+    if not proxy_vwap:
+        raise ValueError(
+            "Baostock cannot provide an adjustment-consistent raw VWAP; "
+            "use proxy_vwap=True"
+        )
+
+    import baostock as bs
+
+    if not tickers:
+        raise ValueError("Tickers list cannot be empty")
+
+    # Uniqueify tickers to avoid issues with duplicates during download or pivoting
+    unique_tickers = list(dict.fromkeys(tickers))
+
+    lg = bs.login()
+    if lg.error_code != '0':
+        raise RuntimeError(f"Baostock login failed: {lg.error_msg}")
+
+    try:
+        all_data = []
+        for ticker in unique_tickers:
+            rs = bs.query_history_k_data_plus(
+                ticker,
+                "date,code,open,high,low,close,preclose,volume,amount,tradestatus",
+                start_date=start,
+                end_date=end,
+                frequency="d",
+                adjustflag="2" # 2 for forward adjust
+            )
+            if rs.error_code != '0':
+                print(f"Warning: Failed to fetch {ticker}: {rs.error_msg}")
+                continue
+
+            while (rs.error_code == '0') & rs.next():
+                all_data.append(rs.get_row_data())
+    finally:
+        bs.logout()
+
+    if not all_data:
+        raise ValueError(f"No data returned for tickers {unique_tickers} from {start} to {end}")
+
+    df = pd.DataFrame(all_data, columns=["date", "code", "open", "high", "low", "close", "preclose", "volume", "amount", "tradestatus"])
+
+    # Convert numerical columns
+    num_cols = ["open", "high", "low", "close", "preclose", "volume", "amount"]
+    for col in num_cols:
+        df[col] = pd.to_numeric(df[col], errors='coerce')
+
+    df["date"] = pd.to_datetime(df["date"])
+
+    # Extract wide DataFrames
+    fields_wide = {}
+
+    def get_wide_df(col_name):
+        wide = df.pivot(index="date", columns="code", values=col_name)
+        # Ensure all requested tickers are present
+        for t in unique_tickers:
+            if t not in wide.columns:
+                wide[t] = np.nan
+        return wide[list(unique_tickers)]
+
+    open_df = get_wide_df("open")
+    high_df = get_wide_df("high")
+    low_df = get_wide_df("low")
+    close_df = get_wide_df("close")
+    volume_df = get_wide_df("volume")
+    amount_df = get_wide_df("amount")
+    preclose_df = get_wide_df("preclose")
+    status_df = get_wide_df("tradestatus") # 1 means tradable, 0 means halt
+
+    fields_wide["open"] = open_df
+    fields_wide["high"] = high_df
+    fields_wide["low"] = low_df
+    fields_wide["close"] = close_df
+    fields_wide["volume"] = volume_df
+    fields_wide["amount"] = amount_df
+    fields_wide["last_close"] = preclose_df
+
+    # OHLC are forward-adjusted (adjustflag="2"), while amount and volume are
+    # not. The typical-price proxy keeps every price field on the same basis.
+    fields_wide["vwap"] = (open_df + close_df + high_df + low_df) / 4.0
+
+    dates = open_df.index.to_numpy()
+    stocks = list(unique_tickers)
+
+    # status_df is string '1' or '0' natively before we pivoted, so convert it properly
+    # If a stock was missing on a day, it's NaN. Let's make tradable mask: status == '1'
+    status_df = status_df.fillna('0')
+    is_tradable = (status_df == '1') | (status_df == 1)
+
+    # Mask where close is not NaN and tradable
+    mask = (~open_df.isna() & ~high_df.isna() & ~low_df.isna() & ~close_df.isna() & ~volume_df.isna() & is_tradable).to_numpy()
+
+    tensors = {
+        name: torch.from_numpy(df_.fillna(0.0).to_numpy(dtype=np.float32).copy()).to(device)
+        for name, df_ in fields_wide.items()
+    }
+
+    return Panel.from_tensors(
+        dates=dates,
+        stocks=stocks,
+        fields=tensors,
+        mask=torch.from_numpy(mask.copy()).to(device),
+    )

mlquant/data/loaders.py

@@ -0,0 +1,125 @@
+"""Loaders for real OCHLV CSV / parquet files.
+
+The paper uses Wind exports; this module is intentionally generic so
+you can plug in any tab-separated OCHLV dump that has at least these
+columns:
+
+    TRADE_DT  S_INFO_WINDCODE  open  close  high  low  vwap  volume
+
+Any extra columns — ``S_DQ_AMOUNT`` (turnover), ``LIMIT_UP``,
+``LIMIT_DOWN``, ``LAST_CLOSE``, ``FHZS_FLAG`` (ex-dividend / split
+flag) — are picked up automatically when present and surfaced on the
+:class:`Panel` so the bias-correction and market-breadth pipelines can
+use the *real* exchange figures rather than ±10 % proxies.
+
+The loader is the single entry point for any external file path; the
+synthetic generator in :mod:`mlquant.data.synthetic` produces the same
+shape, so downstream code never branches on data source.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Mapping, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+import torch
+
+from .panel import Panel
+
+
+# Each entry maps the canonical Panel field name to a tuple of column
+# aliases we accept on disk. The first match wins, so the most specific
+# Wind name should come first.
+_REQUIRED_ALIASES: Mapping[str, tuple[str, ...]] = {
+    "open":   ("open",   "S_FWDS_ADJOPEN",  "OPEN"),
+    "high":   ("high",   "S_FWDS_ADJHIGH",  "HIGH"),
+    "low":    ("low",    "S_FWDS_ADJLOW",   "LOW"),
+    "close":  ("close",  "S_FWDS_ADJCLOSE", "CLOSE"),
+    "volume": ("volume", "S_DQ_VOLUME"),
+    "vwap":   ("vwap",   "S_DQ_AVGPRICE"),
+}
+
+_OPTIONAL_ALIASES: Mapping[str, tuple[str, ...]] = {
+    "amount":     ("amount",     "S_DQ_AMOUNT",   "AMOUNT"),
+    "limit_up":   ("limit_up",   "LIMIT_UP",      "S_DQ_LIMIT"),
+    "limit_down": ("limit_down", "LIMIT_DOWN",    "S_DQ_STOPPING"),
+    "last_close": ("last_close", "LAST_CLOSE",    "S_DQ_PRECLOSE"),
+}
+
+
+def _resolve(df: pd.DataFrame, name: str, aliases: tuple[str, ...]) -> Optional[pd.Series]:
+    for alias in aliases:
+        if alias in df.columns:
+            return df[alias].astype(np.float32)
+    return None
+
+
+def load_ochlv_csv(
+    path: str | Path,
+    *,
+    sep: str = "\t",
+    universe: Optional[Sequence[str]] = None,
+    device: str | torch.device = "cpu",
+) -> Panel:
+    """Load a long-format OCHLV file into a :class:`Panel`.
+
+    Parameters
+    ----------
+    path : str | Path
+        Path to a tab-separated (or ``sep`` separated) file with at least
+        ``TRADE_DT`` and ``S_INFO_WINDCODE`` columns plus OCHLV fields.
+    universe : sequence of str, optional
+        If given, restrict to this set of tickers before pivoting.
+    device : torch device or str
+        Where to allocate the resulting tensors.
+    """
+    df = pd.read_csv(path, sep=sep, low_memory=False)
+    if "TRADE_DT" not in df.columns or "S_INFO_WINDCODE" not in df.columns:
+        raise KeyError("file must contain TRADE_DT and S_INFO_WINDCODE columns")
+
+    df["TRADE_DT"] = pd.to_datetime(df["TRADE_DT"].astype(str))
+    df["S_INFO_WINDCODE"] = df["S_INFO_WINDCODE"].astype(str).str[:6]
+    if universe is not None:
+        df = df[df["S_INFO_WINDCODE"].isin(set(universe))]
+    df = df.drop_duplicates(subset=["TRADE_DT", "S_INFO_WINDCODE"]).sort_values(
+        ["TRADE_DT", "S_INFO_WINDCODE"]
+    )
+
+    # Build a wide [date, stock] frame for every field we recognise.
+    fields_long: dict[str, pd.Series] = {}
+    for name, aliases in {**_REQUIRED_ALIASES, **_OPTIONAL_ALIASES}.items():
+        ser = _resolve(df, name, aliases)
+        if ser is not None:
+            fields_long[name] = ser
+        elif name in _REQUIRED_ALIASES:
+            raise KeyError(
+                f"None of {aliases} present in dataframe columns {list(df.columns)}"
+            )
+
+    fields_wide: dict[str, pd.DataFrame] = {}
+    for fld, ser in fields_long.items():
+        wide = (
+            df.assign(**{fld: ser})
+              .pivot(index="TRADE_DT", columns="S_INFO_WINDCODE", values=fld)
+              .sort_index()
+        )
+        fields_wide[fld] = wide
+
+    base = fields_wide["close"]
+    dates = base.index.to_numpy()
+    stocks = base.columns.to_numpy()
+
+    mask = (~base.isna()).to_numpy()
+    tensors = {
+        name: torch.from_numpy(df_.fillna(0.0).to_numpy(dtype=np.float32)).to(device)
+        for name, df_ in fields_wide.items()
+    }
+
+    panel = Panel.from_tensors(
+        dates=dates,
+        stocks=stocks,
+        fields=tensors,
+        mask=torch.from_numpy(mask).to(device),
+    )
+    return panel

mlquant/data/panel.py

@@ -0,0 +1,206 @@
+"""The :class:`Panel` — a thin, mask-aware container for OCHLV panels.
+
+Design rules
+------------
+* Every numeric field has shape ``[T, N]`` (date × stock) and dtype
+  ``float32``.
+* ``mask[t, i] == True`` iff stock *i* was tradable on date *t*. Limit-up,
+  limit-down, halted, and pre-IPO cells are masked out — this is the
+  "bias correction" plumbing the paper relies on.
+* Panels are *immutable*: factor pipelines return new tensors rather
+  than mutating in place. This keeps reasoning straightforward and is
+  a hard requirement for autograd-backed factors.
+
+Optional A-share microstructure fields
+--------------------------------------
+The legacy ``cuda_features.py`` operates on Wind dumps that carry a few
+extra columns beyond OCHLV: traded ``amount`` (turnover in CNY), the
+exchange-published ``limit_up`` / ``limit_down`` prices and the previous
+session's ``last_close``. They are *optional* on this dataclass — the
+synthetic generator emits them, real-data loaders fill them when the
+file has the columns, and downstream code falls back to derived
+quantities when they are absent (e.g. ±10 % return proxy for limit
+moves).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import Iterable, Mapping, Optional, Sequence
+
+import numpy as np
+import torch
+
+
+_REQUIRED_FIELDS = ("open", "high", "low", "close", "volume", "vwap")
+_OPTIONAL_FIELDS = ("amount", "limit_up", "limit_down", "last_close")
+
+
+@dataclass(frozen=True)
+class Panel:
+    """OCHLV panel + tradability mask, all on the same device."""
+
+    dates:  np.ndarray            # shape [T], dtype object/datetime64
+    stocks: np.ndarray            # shape [N], dtype object (ticker str)
+
+    open:   torch.Tensor          # [T, N]
+    high:   torch.Tensor
+    low:    torch.Tensor
+    close:  torch.Tensor
+    volume: torch.Tensor
+    vwap:   torch.Tensor
+    mask:   torch.Tensor          # [T, N] bool
+
+    # ------- optional A-share microstructure fields -----------------------
+    # All [T, N] float32 when present; ``None`` means "data source did not
+    # supply this column — fall back to derived defaults".
+    amount:     Optional[torch.Tensor] = field(default=None)
+    limit_up:   Optional[torch.Tensor] = field(default=None)
+    limit_down: Optional[torch.Tensor] = field(default=None)
+    last_close: Optional[torch.Tensor] = field(default=None)
+
+    # ------------------------------------------------------------------
+    # Construction helpers
+    # ------------------------------------------------------------------
+    @classmethod
+    def from_tensors(
+        cls,
+        dates: Sequence,
+        stocks: Sequence[str],
+        fields: Mapping[str, torch.Tensor],
+        mask: torch.Tensor,
+    ) -> "Panel":
+        missing = set(_REQUIRED_FIELDS) - set(fields)
+        if missing:
+            raise KeyError(f"Panel.from_tensors missing fields: {sorted(missing)}")
+        T, N = mask.shape
+        promoted = {}
+        for k, v in fields.items():
+            if k not in _REQUIRED_FIELDS and k not in _OPTIONAL_FIELDS:
+                # Allow forward-compat fields to round-trip silently.
+                continue
+            if tuple(v.shape) != (T, N):
+                raise ValueError(f"field {k!r} has shape {tuple(v.shape)}, expected {(T, N)}")
+            promoted[k] = v.float() if v.dtype != torch.float32 else v
+        kwargs = {k: promoted[k] for k in _REQUIRED_FIELDS}
+        for k in _OPTIONAL_FIELDS:
+            if k in promoted:
+                kwargs[k] = promoted[k]
+        return cls(
+            dates=np.asarray(dates),
+            stocks=np.asarray(stocks),
+            mask=mask.bool(),
+            **kwargs,
+        )
+
+    # ------------------------------------------------------------------
+    # Convenience
+    # ------------------------------------------------------------------
+    @property
+    def n_dates(self) -> int:
+        return int(self.mask.shape[0])
+
+    @property
+    def n_stocks(self) -> int:
+        return int(self.mask.shape[1])
+
+    @property
+    def device(self) -> torch.device:
+        return self.close.device
+
+    @property
+    def returns(self) -> torch.Tensor:
+        """Daily simple returns, masked at non-tradable cells.
+
+        When :attr:`last_close` is present we use it directly so the day-0
+        return is meaningful and adjusted-vs-unadjusted close mismatches
+        do not inject NaNs. Otherwise we fall back to the previous-row
+        close.
+        """
+        if self.last_close is not None:
+            prev = self.last_close.clamp_min(1e-9)
+            ret = self.close / prev - 1.0
+            valid = self.mask & (self.last_close > 0)
+            return ret.masked_fill(~valid, 0.0)
+        prev = torch.roll(self.close, shifts=1, dims=0)
+        prev[0] = self.close[0]                     # avoid div-by-zero on day 0
+        ret = self.close / prev - 1.0
+        ret[0] = 0.0
+        valid = self.mask & torch.roll(self.mask, shifts=1, dims=0)
+        valid[0] = False
+        return ret.masked_fill(~valid, 0.0)
+
+    @property
+    def has_real_limits(self) -> bool:
+        """True iff exchange-published limit-up / limit-down prices are present."""
+        return self.limit_up is not None and self.limit_down is not None
+
+    def field(self, name: str) -> Optional[torch.Tensor]:
+        """Generic field accessor — returns ``None`` if the field is absent."""
+        return getattr(self, name, None)
+
+    def to(self, device: torch.device | str) -> "Panel":
+        def _mv(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+            return None if t is None else t.to(device)
+        return replace(
+            self,
+            open=self.open.to(device),
+            high=self.high.to(device),
+            low=self.low.to(device),
+            close=self.close.to(device),
+            volume=self.volume.to(device),
+            vwap=self.vwap.to(device),
+            mask=self.mask.to(device),
+            amount=_mv(self.amount),
+            limit_up=_mv(self.limit_up),
+            limit_down=_mv(self.limit_down),
+            last_close=_mv(self.last_close),
+        )
+
+    def slice_dates(self, start: int, end: int) -> "Panel":
+        def _sl(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+            return None if t is None else t[start:end]
+        return replace(
+            self,
+            dates=self.dates[start:end],
+            open=self.open[start:end],
+            high=self.high[start:end],
+            low=self.low[start:end],
+            close=self.close[start:end],
+            volume=self.volume[start:end],
+            vwap=self.vwap[start:end],
+            mask=self.mask[start:end],
+            amount=_sl(self.amount),
+            limit_up=_sl(self.limit_up),
+            limit_down=_sl(self.limit_down),
+            last_close=_sl(self.last_close),
+        )
+
+    # ------------------------------------------------------------------
+    # Sanity checks (used by tests + CLI ``mlquant check``)
+    # ------------------------------------------------------------------
+    def assert_consistent(self) -> None:
+        T, N = self.mask.shape
+        assert self.dates.shape == (T,), f"dates {self.dates.shape}"
+        assert self.stocks.shape == (N,), f"stocks {self.stocks.shape}"
+        for name in _REQUIRED_FIELDS:
+            f_ = getattr(self, name)
+            assert f_.shape == (T, N), f"{name} {f_.shape}"
+        for name in _OPTIONAL_FIELDS:
+            f_ = getattr(self, name)
+            if f_ is not None:
+                assert f_.shape == (T, N), f"{name} {f_.shape}"
+        # H >= O,C,L and L <= O,C,H on tradable cells
+        m = self.mask
+        ohlc = torch.stack([self.open, self.high, self.low, self.close], dim=0)
+        assert ((self.high >= ohlc[[0, 2, 3]].max(dim=0).values) | ~m).all(), "high < o/l/c"
+        assert ((self.low  <= ohlc[[0, 1, 3]].min(dim=0).values) | ~m).all(), "low  > o/h/c"
+
+
+def fields_iter(panel: Panel) -> Iterable[tuple[str, torch.Tensor]]:
+    """Yield ``(name, tensor)`` for every numeric panel field that is set."""
+    for name in _REQUIRED_FIELDS:
+        yield name, getattr(panel, name)
+    for name in _OPTIONAL_FIELDS:
+        v = getattr(panel, name)
+        if v is not None:
+            yield name, v