quant-os/tools/build_qlib_tiny_fixture.py

"""Build a tiny synthetic Qlib binary provider with only the standard library."""

from __future__ import annotations

import argparse
import datetime as dt
import json
import math
import struct
from pathlib import Path
from typing import Dict, Iterable, Mapping, Optional, Sequence


def _qlib_name(symbol: str) -> str:
    value = str(symbol).strip().upper()
    if len(value) != 8 or value[:2] not in {"SH", "SZ", "BJ"}:
        raise ValueError("fixture symbols must use Qlib SH600000 form")
    if not value[2:].isdigit():
        raise ValueError("invalid Qlib symbol")
    return value


def _write_lines(path: Path, lines: Iterable[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _write_feature(path: Path, start_index: int, values: Sequence[float]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    numeric = [float(start_index)] + [float(value) for value in values]
    path.write_bytes(struct.pack("<" + "f" * len(numeric), *numeric))


def write_qlib_fixture(
    root: Path,
    *,
    calendar: Sequence[str],
    features: Mapping[str, Mapping[str, Sequence[float]]],
    markets: Optional[Mapping[str, Sequence[str]]] = None,
) -> Dict[str, object]:
    """Write calendar/instruments/features in Qlib's documented bin layout."""
    output = root.expanduser().resolve()
    if not calendar:
        raise ValueError("calendar must not be empty")
    if len(set(calendar)) != len(calendar):
        raise ValueError("calendar contains duplicates")
    ordered_calendar = sorted(str(value) for value in calendar)
    if ordered_calendar != list(calendar):
        raise ValueError("calendar must be sorted")
    _write_lines(output / "calendars/day.txt", ordered_calendar)

    normalized: Dict[str, Mapping[str, Sequence[float]]] = {}
    for raw_symbol, fields in features.items():
        symbol = _qlib_name(raw_symbol)
        if not fields:
            raise ValueError(f"{symbol} has no fields")
        for field, values in fields.items():
            if len(values) != len(calendar):
                raise ValueError(
                    f"{symbol}/{field} has {len(values)} values; "
                    f"expected {len(calendar)}"
                )
            for value in values:
                if not math.isfinite(float(value)):
                    raise ValueError(f"{symbol}/{field} contains non-finite data")
            _write_feature(
                output
                / "features"
                / symbol.lower()
                / (str(field).lower().lstrip("$") + ".day.bin"),
                0,
                values,
            )
        normalized[symbol] = fields

    start, end = calendar[0], calendar[-1]
    all_lines = [
        f"{symbol}\t{start}\t{end}" for symbol in sorted(normalized)
    ]
    _write_lines(output / "instruments/all.txt", all_lines)
    for market, symbols in (markets or {}).items():
        selected = [_qlib_name(symbol) for symbol in symbols]
        unknown = set(selected).difference(normalized)
        if unknown:
            raise ValueError(f"{market} contains missing symbols: {sorted(unknown)}")
        _write_lines(
            output / f"instruments/{market.lower()}.txt",
            [f"{symbol}\t{start}\t{end}" for symbol in sorted(selected)],
        )
    return {
        "provider_uri": str(output),
        "calendar_count": len(calendar),
        "instruments": sorted(normalized),
        "markets": sorted((markets or {}).keys()),
        "format_note": (
            "each *.day.bin is little-endian float32; element 0 is calendar "
            "start index, followed by field values"
        ),
    }


def _business_days(start: dt.date, count: int) -> list[str]:
    output = []
    current = start
    while len(output) < count:
        if current.weekday() < 5:
            output.append(current.isoformat())
        current += dt.timedelta(days=1)
    return output


def _fields(closes: Sequence[float], volume: float) -> Dict[str, Sequence[float]]:
    change = [0.0]
    for previous, current in zip(closes, closes[1:]):
        change.append(float(current) / float(previous) - 1.0)
    return {
        "open": [value * 0.998 for value in closes],
        "high": [value * 1.01 for value in closes],
        "low": [value * 0.99 for value in closes],
        "close": list(closes),
        "volume": [volume] * len(closes),
        "factor": [1.0] * len(closes),
        "change": change,
    }


def build_default_fixture(root: Path, days: int = 80) -> Dict[str, object]:
    """Create a deterministic two-stock market and one benchmark."""
    if int(days) < 30:
        raise ValueError("default fixture needs at least 30 business days")
    calendar = _business_days(dt.date(2024, 1, 2), int(days))
    closes_up = [10.0 * (1.004 ** index) for index in range(len(calendar))]
    closes_down = [20.0 * (0.999 ** index) for index in range(len(calendar))]
    closes_bench = [100.0 * (1.001 ** index) for index in range(len(calendar))]
    return write_qlib_fixture(
        root,
        calendar=calendar,
        features={
            "SH600000": _fields(closes_up, 2_000_000.0),
            "SZ000001": _fields(closes_down, 1_500_000.0),
            "SH000300": _fields(closes_bench, 100_000_000.0),
        },
        markets={"csi300": ["SH600000", "SZ000001"]},
    )


def _main(argv: Optional[list[str]] = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument("--days", type=int, default=80)
    args = parser.parse_args(argv)
    print(
        json.dumps(
            build_default_fixture(args.output_dir, args.days),
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(_main())