chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,30 @@
|
||||
# Shared utilities
|
||||
|
||||
Library code imported across every chapter and case study. Notebooks run from the
|
||||
repository root, so `import utils` and `from utils.x import y` resolve with no
|
||||
installation. This is code you **import**, not run — for command-line tools, see
|
||||
[`scripts/`](../scripts).
|
||||
|
||||
**Configuration and paths.** `config.py` loads and validates the paths and settings
|
||||
in `.env` (and sorts out CUDA library paths); `paths.py` holds the chapter registry
|
||||
and resolves chapter, case-study, and output directories so notebooks never hard-code
|
||||
a location.
|
||||
|
||||
**Figures.** `style.py` defines the ML4T color palette and the matplotlib / Plotly
|
||||
defaults that give every figure in the book one consistent look.
|
||||
|
||||
**Data.** `data_quality.py` summarizes coverage, checks OHLC invariants, and subsets
|
||||
symbols for fast test runs; `downloading.py` is the shared backbone of the `data/`
|
||||
download scripts (argument parsing, path/YAML resolution, atomic writes);
|
||||
`artifact_specs.py` loads the per-case-study YAML sidecars that describe market data,
|
||||
labels, and features.
|
||||
|
||||
**Modeling and cross-validation.** `modeling.py` is the workhorse — it loads a
|
||||
modeling dataset, parses model configs, prepares folds, and detects the schema;
|
||||
`cv_splits.py` builds the walk-forward splits (calendar-aware, leakage-safe);
|
||||
`predictions_cache.py` caches long-form prediction frames so the teaching notebooks
|
||||
don't recompute them.
|
||||
|
||||
**Reproducibility.** `reproducibility.py` seeds Python, NumPy, and Torch (CPU + CUDA)
|
||||
in a single call; `storage_benchmarks.py` provides the synthetic data and timing
|
||||
harness behind the Chapter 2 storage benchmarks.
|
||||
@@ -0,0 +1,54 @@
|
||||
"""ML4T utilities - configuration, paths, styling, and data quality.
|
||||
|
||||
This package provides:
|
||||
- Configuration management (utils.config)
|
||||
- Path utilities (utils.paths)
|
||||
- Visualization styling (utils.style)
|
||||
- Data quality checks (utils.data_quality)
|
||||
|
||||
Usage:
|
||||
>>> from utils import ML4T_PATH, ML4T_DATA_PATH, DATA_DIR
|
||||
>>> from utils.paths import CHAPTERS, get_output_dir
|
||||
>>> from utils.style import COLORS
|
||||
>>> from utils.data_quality import describe_coverage, check_ohlc_invariants
|
||||
"""
|
||||
|
||||
from utils.config import (
|
||||
ALPACA_API_KEY,
|
||||
ALPACA_SECRET_KEY,
|
||||
CASE_STUDIES_DIR,
|
||||
DATABENTO_API_KEY,
|
||||
ML4T_DATA_PATH,
|
||||
ML4T_PATH,
|
||||
OANDA_API_KEY,
|
||||
REPO_ROOT,
|
||||
)
|
||||
|
||||
# Backward compatibility alias
|
||||
DATA_DIR = ML4T_DATA_PATH
|
||||
|
||||
# Plotly: include PNG in cell output so GitHub can render .ipynb figures.
|
||||
# Override with PLOTLY_RENDERER env var (e.g. "json" for headless CI).
|
||||
try:
|
||||
import os as _os
|
||||
|
||||
import plotly.io as _pio
|
||||
|
||||
if not _os.environ.get("PLOTLY_RENDERER"):
|
||||
_pio.renderers.default = "plotly_mimetype+png"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__all__ = [
|
||||
# Core configuration
|
||||
"ML4T_PATH",
|
||||
"ML4T_DATA_PATH",
|
||||
"DATA_DIR", # Backward compatibility
|
||||
"CASE_STUDIES_DIR",
|
||||
"REPO_ROOT",
|
||||
# API keys
|
||||
"DATABENTO_API_KEY",
|
||||
"OANDA_API_KEY",
|
||||
"ALPACA_API_KEY",
|
||||
"ALPACA_SECRET_KEY",
|
||||
]
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Compatibility loaders for case-study artifact sidecars."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from copy import deepcopy
|
||||
from functools import cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from utils.paths import get_case_study_dir
|
||||
|
||||
try:
|
||||
from ml4t.backtest.spec_io import load_spec as _load_shared_spec
|
||||
except ImportError: # pragma: no cover - depends on install state
|
||||
_load_shared_spec = None
|
||||
|
||||
|
||||
def _artifact_root(case_study_id: str) -> Path:
|
||||
return get_case_study_dir(case_study_id, create=False) / "config" / "artifacts"
|
||||
|
||||
|
||||
def _to_mapping(spec: Any) -> dict[str, Any]:
|
||||
if isinstance(spec, Mapping):
|
||||
return dict(spec)
|
||||
if hasattr(spec, "to_dict"):
|
||||
return dict(spec.to_dict())
|
||||
raise TypeError(f"Unsupported artifact spec type: {type(spec).__name__}")
|
||||
|
||||
|
||||
def _load_spec(path: Path) -> dict[str, Any] | None:
|
||||
if not path.exists():
|
||||
return None
|
||||
if _load_shared_spec is not None:
|
||||
return _to_mapping(_load_shared_spec(path))
|
||||
with path.open() as f:
|
||||
data = yaml.safe_load(f)
|
||||
return dict(data) if data else None
|
||||
|
||||
|
||||
@cache
|
||||
def _load_market_data_spec_cached(case_study_id: str) -> dict[str, Any] | None:
|
||||
return _load_spec(_artifact_root(case_study_id) / "market_data.yaml")
|
||||
|
||||
|
||||
@cache
|
||||
def _load_label_spec_cached(case_study_id: str, label: str) -> dict[str, Any] | None:
|
||||
return _load_spec(_artifact_root(case_study_id) / "labels" / f"{label}.yaml")
|
||||
|
||||
|
||||
@cache
|
||||
def _load_feature_spec_cached(case_study_id: str, feature_set: str) -> dict[str, Any] | None:
|
||||
return _load_spec(_artifact_root(case_study_id) / "features" / f"{feature_set}.yaml")
|
||||
|
||||
|
||||
def load_market_data_spec(case_study_id: str) -> dict[str, Any] | None:
|
||||
spec = _load_market_data_spec_cached(case_study_id)
|
||||
return deepcopy(spec) if spec is not None else None
|
||||
|
||||
|
||||
def load_label_spec(case_study_id: str, label: str) -> dict[str, Any] | None:
|
||||
spec = _load_label_spec_cached(case_study_id, label)
|
||||
return deepcopy(spec) if spec is not None else None
|
||||
|
||||
|
||||
def load_feature_spec(case_study_id: str, feature_set: str) -> dict[str, Any] | None:
|
||||
spec = _load_feature_spec_cached(case_study_id, feature_set)
|
||||
return deepcopy(spec) if spec is not None else None
|
||||
|
||||
|
||||
def resolve_storage_path(
|
||||
case_study_id: str,
|
||||
spec: dict[str, Any] | None,
|
||||
default_relative_path: str,
|
||||
) -> Path:
|
||||
if spec is None:
|
||||
return get_case_study_dir(case_study_id, create=False) / default_relative_path
|
||||
storage = spec.get("storage", {})
|
||||
rel_path = storage.get("path", default_relative_path)
|
||||
return get_case_study_dir(case_study_id, create=False) / str(rel_path)
|
||||
|
||||
|
||||
def resolve_label_buffer(
|
||||
case_study_id: str,
|
||||
label: str,
|
||||
setup: Mapping[str, Any] | None = None,
|
||||
) -> str | None:
|
||||
label_spec = load_label_spec(case_study_id, label)
|
||||
if label_spec is not None:
|
||||
definition = label_spec.get("definition", {})
|
||||
if definition.get("buffer"):
|
||||
return str(definition["buffer"])
|
||||
|
||||
labels = (setup or {}).get("labels", {})
|
||||
if labels.get("primary") == label and labels.get("buffer"):
|
||||
return str(labels["buffer"])
|
||||
variant_buffers = labels.get("variant_buffers", {})
|
||||
if label in variant_buffers:
|
||||
return str(variant_buffers[label])
|
||||
return None
|
||||
|
||||
|
||||
def resolve_market_semantics(
|
||||
case_study_id: str,
|
||||
setup: Mapping[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
market_spec = load_market_data_spec(case_study_id) or {}
|
||||
semantics = market_spec.get("semantics", {})
|
||||
evaluation = (setup or {}).get("evaluation", {})
|
||||
return {
|
||||
"calendar": semantics.get("calendar") or evaluation.get("calendar"),
|
||||
"timezone": semantics.get("timezone"),
|
||||
"data_frequency": semantics.get("data_frequency"),
|
||||
"timestamp_semantics": semantics.get("timestamp_semantics"),
|
||||
"session_start_time": semantics.get("session_start_time"),
|
||||
"bar_type": semantics.get("bar_type"),
|
||||
}
|
||||
|
||||
|
||||
def resolve_market_runtime(case_study_id: str) -> dict[str, Any]:
|
||||
market_spec = load_market_data_spec(case_study_id) or {}
|
||||
runtime = market_spec.get("runtime", {})
|
||||
return dict(runtime) if isinstance(runtime, Mapping) else {}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"load_feature_spec",
|
||||
"load_label_spec",
|
||||
"load_market_data_spec",
|
||||
"resolve_label_buffer",
|
||||
"resolve_market_runtime",
|
||||
"resolve_market_semantics",
|
||||
"resolve_storage_path",
|
||||
]
|
||||
+121
@@ -0,0 +1,121 @@
|
||||
"""Centralized configuration using python-dotenv.
|
||||
|
||||
This module loads all configuration from the .env file in the repository root.
|
||||
It provides explicit, fail-fast configuration with clear error messages.
|
||||
|
||||
Configuration priority:
|
||||
1. .env file (ONLY source - no fallbacks)
|
||||
2. Validation ensures paths exist or provides clear instructions
|
||||
|
||||
Usage:
|
||||
from utils import ML4T_PATH, ML4T_DATA_PATH, REPO_ROOT
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Find repository root (where .env lives)
|
||||
REPO_ROOT = Path(__file__).parent.parent.resolve()
|
||||
ENV_FILE = REPO_ROOT / ".env"
|
||||
|
||||
# Auto-create .env from .env.example if missing (CI, first-time setup)
|
||||
if not ENV_FILE.exists():
|
||||
example = REPO_ROOT / ".env.example"
|
||||
if example.exists():
|
||||
import contextlib
|
||||
import shutil
|
||||
|
||||
with contextlib.suppress(OSError):
|
||||
shutil.copy(example, ENV_FILE)
|
||||
else:
|
||||
raise FileNotFoundError(
|
||||
f".env file not found at {ENV_FILE}\n\n"
|
||||
f"Copy .env.example to .env and configure paths:\n"
|
||||
f" cd {REPO_ROOT}\n"
|
||||
f" cp .env.example .env\n"
|
||||
f" # Edit .env with your paths:\n"
|
||||
f" # ML4T_PATH={REPO_ROOT}\n"
|
||||
f" # ML4T_DATA_PATH=/path/to/data\n"
|
||||
)
|
||||
|
||||
# Load environment variables from .env
|
||||
# override=False means environment variables take precedence over .env file
|
||||
load_dotenv(ENV_FILE, override=False)
|
||||
|
||||
# ============================================================================
|
||||
# Required Paths
|
||||
# ============================================================================
|
||||
|
||||
# Priority: 1. Environment variable, 2. .env file, 3. Default
|
||||
ML4T_PATH = Path(os.getenv("ML4T_PATH", str(REPO_ROOT))).expanduser().resolve()
|
||||
data_path = os.getenv("ML4T_DATA_PATH")
|
||||
if data_path is None:
|
||||
ML4T_DATA_PATH = ML4T_PATH / "data"
|
||||
else:
|
||||
ML4T_DATA_PATH = Path(data_path).expanduser().resolve()
|
||||
|
||||
# Case studies directory - centralized artifact store for all strategies
|
||||
# Default: case_studies/ in repo root (git-tracked configs, gitignored binaries)
|
||||
case_studies_path = os.getenv("CASE_STUDIES_DIR")
|
||||
if case_studies_path is None:
|
||||
CASE_STUDIES_DIR = REPO_ROOT / "case_studies"
|
||||
else:
|
||||
CASE_STUDIES_DIR = Path(case_studies_path).expanduser().resolve()
|
||||
|
||||
# ============================================================================
|
||||
# Validation
|
||||
# ============================================================================
|
||||
|
||||
# Validate ML4T_PATH exists
|
||||
if not ML4T_PATH.exists():
|
||||
raise FileNotFoundError(
|
||||
f"ML4T_PATH does not exist: {ML4T_PATH}\n\n"
|
||||
f"Update ML4T_PATH in .env to point to the repository root."
|
||||
)
|
||||
|
||||
# Validate ML4T_DATA_PATH exists
|
||||
if not ML4T_DATA_PATH.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Data directory not found: {ML4T_DATA_PATH}\n\n"
|
||||
f"Options:\n"
|
||||
f" 1. Download data:\n"
|
||||
f" python data/download_all.py\n"
|
||||
f" 2. Update ML4T_DATA_PATH in .env to point to existing data\n"
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# CUDA Library Path (local uv environments)
|
||||
# ============================================================================
|
||||
|
||||
# PyTorch bundles its own CUDA libraries which may be newer than system ones.
|
||||
# Without this, torch imports fail with "undefined symbol" errors on systems
|
||||
# where the system libcudart.so is older than what torch expects.
|
||||
# Docker images don't need this (CUDA is installed system-wide).
|
||||
if not os.environ.get("LD_LIBRARY_PATH", "").startswith("/usr/local/cuda"):
|
||||
try:
|
||||
import torch as _torch
|
||||
|
||||
_torch_root = Path(_torch.__file__).parent
|
||||
_cuda_paths = [str(_torch_root / "lib")]
|
||||
_nvidia_dir = _torch_root.parent / "nvidia"
|
||||
if _nvidia_dir.exists():
|
||||
_cuda_paths.extend(str(p) for p in _nvidia_dir.glob("*/lib"))
|
||||
_existing_ld = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
os.environ["LD_LIBRARY_PATH"] = ":".join(
|
||||
_cuda_paths + ([_existing_ld] if _existing_ld else [])
|
||||
)
|
||||
del _torch, _torch_root, _cuda_paths, _nvidia_dir, _existing_ld
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# ============================================================================
|
||||
# API Keys (optional - only needed for data downloads)
|
||||
# ============================================================================
|
||||
|
||||
# These are optional and only needed if you're downloading data
|
||||
DATABENTO_API_KEY = os.getenv("DATABENTO_API_KEY", "")
|
||||
OANDA_API_KEY = os.getenv("OANDA_API_KEY", "")
|
||||
ALPACA_API_KEY = os.getenv("ALPACA_API_KEY", "")
|
||||
ALPACA_SECRET_KEY = os.getenv("ALPACA_SECRET_KEY", "")
|
||||
@@ -0,0 +1,338 @@
|
||||
"""Cross-validation split generation for case study pipelines.
|
||||
|
||||
Reads the ``evaluation`` section from ``setup.yaml`` and generates
|
||||
walk-forward date boundaries by delegating to ml4t-diagnostic's
|
||||
``WalkForwardCV``. This is the single source of truth for CV splits
|
||||
used by all case studies (Ch11+).
|
||||
|
||||
Usage:
|
||||
from utils.cv_splits import generate_cv_splits, load_evaluation_config, make_walk_forward_config
|
||||
|
||||
# Date-boundary splits
|
||||
splits = generate_cv_splits(dataset, case_study_id="etfs", label_buffer="21D")
|
||||
for split in splits:
|
||||
train_mask = (df[date_col] >= split["train_start"]) & (df[date_col] <= split["train_end"])
|
||||
val_mask = (df[date_col] >= split["val_start"]) & (df[date_col] <= split["val_end"])
|
||||
|
||||
# WalkForwardConfig for library integration
|
||||
config = make_walk_forward_config("etfs", label_horizon="21D")
|
||||
|
||||
Design decisions:
|
||||
- Delegates fold generation to ml4t-diagnostic's WalkForwardCV
|
||||
- Calendar-aware splitting (NYSE, CME, etc.) replaces broken ppd arithmetic
|
||||
- Operates on unique dates (handles panel data correctly)
|
||||
- Rolling training windows (respects train_size from config)
|
||||
- Backward stepping from holdout boundary
|
||||
- label_buffer is provided at call time (depends on label, not config)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import polars as pl
|
||||
import yaml
|
||||
|
||||
from utils.artifact_specs import resolve_market_semantics
|
||||
from utils.paths import get_case_study_dir
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ml4t.diagnostic.splitters.config import WalkForwardConfig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Calendar name mapping: setup.yaml → pandas_market_calendars exchange names
|
||||
# ---------------------------------------------------------------------------
|
||||
_CALENDAR_MAP: dict[str, str | None] = {
|
||||
"NYSE": "NYSE",
|
||||
"CME": "CME_Equity",
|
||||
"FX": "CME_FX",
|
||||
"crypto": None, # 24/7 trading, no calendar
|
||||
}
|
||||
|
||||
|
||||
def _map_calendar_id(calendar: str | None) -> str | None:
|
||||
"""Map setup.yaml calendar name to pandas_market_calendars exchange name.
|
||||
|
||||
Returns None for 24/7 markets (crypto) to disable calendar-aware splitting.
|
||||
Unknown names are passed through unchanged (will error in the library if invalid).
|
||||
"""
|
||||
if calendar is None:
|
||||
return None
|
||||
return _CALENDAR_MAP.get(calendar, calendar)
|
||||
|
||||
|
||||
def _normalize_duration(s: str) -> str:
|
||||
"""Strip ISO 8601 prefix (P, PT) and normalize unit aliases.
|
||||
|
||||
Examples: P5Y → 5Y, P1Y → 1Y, PT8H → 8h, 21D → 21D (unchanged).
|
||||
"""
|
||||
s = re.sub(r"^P?T?", "", s)
|
||||
s = re.sub(r"(\d+)H$", r"\1h", s)
|
||||
s = re.sub(r"(\d+)T$", r"\1min", s)
|
||||
return s
|
||||
|
||||
|
||||
def _normalize_label_buffer(s: str) -> str:
|
||||
"""Normalize label buffer for pd.Timedelta compatibility.
|
||||
|
||||
Strips ISO prefix, normalizes units, and converts month-based
|
||||
durations to day equivalents since pd.Timedelta rejects 'M' as ambiguous.
|
||||
"""
|
||||
s = _normalize_duration(s)
|
||||
m = re.match(r"^(\d+)M$", s)
|
||||
if m:
|
||||
return f"{int(m.group(1)) * 30}D"
|
||||
return s
|
||||
|
||||
|
||||
def load_evaluation_config(case_study_id: str) -> dict[str, Any]:
|
||||
"""Read the evaluation section from setup.yaml.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
case_study_id : str
|
||||
Case study identifier (e.g., "etfs", "crypto_perps_funding").
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
Evaluation config with keys: n_splits, train_size, val_size,
|
||||
holdout_start, holdout_end, calendar.
|
||||
"""
|
||||
import os
|
||||
|
||||
setup_path = get_case_study_dir(case_study_id) / "config" / "setup.yaml"
|
||||
setup: dict[str, Any] = {}
|
||||
if setup_path.exists():
|
||||
with open(setup_path) as f:
|
||||
setup = yaml.safe_load(f) or {}
|
||||
if "evaluation" not in setup:
|
||||
# Under ML4T_OUTPUT_DIR isolation, the redirected setup.yaml may
|
||||
# be absent or lack hand-curated sections. Fall back to source.
|
||||
test_output = os.environ.get("ML4T_OUTPUT_DIR")
|
||||
if test_output:
|
||||
from utils import CASE_STUDIES_DIR
|
||||
|
||||
source_path = CASE_STUDIES_DIR / case_study_id / "config" / "setup.yaml"
|
||||
if source_path.exists():
|
||||
with open(source_path) as f:
|
||||
setup = yaml.safe_load(f) or {}
|
||||
if "evaluation" not in setup:
|
||||
raise KeyError(
|
||||
f"No 'evaluation' section in {setup_path}. "
|
||||
f"Expected keys: n_splits, train_size, val_size, holdout_start, holdout_end, calendar."
|
||||
)
|
||||
evaluation = dict(setup["evaluation"])
|
||||
market_semantics = resolve_market_semantics(case_study_id, setup)
|
||||
if market_semantics.get("calendar") and not evaluation.get("calendar"):
|
||||
evaluation["calendar"] = market_semantics["calendar"]
|
||||
return evaluation
|
||||
|
||||
|
||||
def make_walk_forward_config(
|
||||
case_study_id: str,
|
||||
label_horizon: str = "0D",
|
||||
date_col: str = "timestamp",
|
||||
) -> WalkForwardConfig:
|
||||
"""Create a WalkForwardConfig from a case study's setup.yaml.
|
||||
|
||||
Bridges the setup.yaml evaluation section to the ml4t-diagnostic
|
||||
library's WalkForwardConfig, using its built-in aliases
|
||||
(val_size→test_size, holdout_start→test_start, etc.).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
case_study_id : str
|
||||
Case study identifier (e.g., "etfs").
|
||||
label_horizon : str, default "0D"
|
||||
Label buffer as duration string (e.g., "21D" for fwd_ret_21d).
|
||||
date_col : str, default "timestamp"
|
||||
Timestamp column name for the dataset.
|
||||
|
||||
Returns
|
||||
-------
|
||||
WalkForwardConfig
|
||||
Configured for the case study's walk-forward protocol.
|
||||
"""
|
||||
from ml4t.diagnostic.splitters import WalkForwardConfig
|
||||
|
||||
eval_config = load_evaluation_config(case_study_id)
|
||||
calendar_id = _map_calendar_id(eval_config.get("calendar"))
|
||||
|
||||
# For D-unit buffers with a calendar, pass as int (trading days)
|
||||
normalized_horizon: int | str = _normalize_label_buffer(label_horizon)
|
||||
if calendar_id is not None and isinstance(normalized_horizon, str):
|
||||
d_match = re.match(r"^(\d+)D$", normalized_horizon)
|
||||
if d_match:
|
||||
normalized_horizon = int(d_match.group(1))
|
||||
|
||||
return WalkForwardConfig(
|
||||
n_splits=eval_config["n_splits"],
|
||||
train_size=_normalize_duration(str(eval_config["train_size"])),
|
||||
val_size=_normalize_duration(str(eval_config["val_size"])),
|
||||
holdout_start=eval_config.get("holdout_start"),
|
||||
holdout_end=eval_config.get("holdout_end"),
|
||||
label_horizon=normalized_horizon,
|
||||
calendar_id=calendar_id,
|
||||
timestamp_col=date_col,
|
||||
fold_direction="backward",
|
||||
)
|
||||
|
||||
|
||||
def make_wf_config(
|
||||
case_study_id: str,
|
||||
label_horizon: str = "0D",
|
||||
date_col: str = "timestamp",
|
||||
) -> WalkForwardConfig:
|
||||
"""Backward-compatible alias for make_walk_forward_config."""
|
||||
return make_walk_forward_config(
|
||||
case_study_id=case_study_id,
|
||||
label_horizon=label_horizon,
|
||||
date_col=date_col,
|
||||
)
|
||||
|
||||
|
||||
def generate_cv_splits(
|
||||
dataset: pl.DataFrame | pd.DataFrame,
|
||||
case_study_id: str | None = None,
|
||||
setup_path: Path | None = None,
|
||||
label_buffer: str = "0D",
|
||||
date_col: str = "timestamp",
|
||||
*,
|
||||
cv_config: dict[str, Any] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Generate walk-forward date splits from evaluation config.
|
||||
|
||||
Delegates to ml4t-diagnostic's ``WalkForwardCV`` for calendar-aware
|
||||
fold generation. Reads the ``evaluation`` section from ``setup.yaml``
|
||||
(via ``case_study_id`` or ``setup_path``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dataset : pl.DataFrame or pd.DataFrame
|
||||
Dataset with a date/timestamp column. Only used to extract unique
|
||||
timestamps -- the full panel rows are not needed.
|
||||
case_study_id : str, optional
|
||||
Case study identifier. Used to locate setup.yaml.
|
||||
setup_path : Path, optional
|
||||
Explicit path to setup.yaml. Takes precedence over case_study_id.
|
||||
label_buffer : str, default "0D"
|
||||
Gap between train_end and val_start sized to the label horizon.
|
||||
Determined by the label being trained on (e.g., "21D" for fwd_ret_21d).
|
||||
date_col : str, default "timestamp"
|
||||
Name of the date/timestamp column.
|
||||
cv_config : dict, optional
|
||||
Pass a cv_config dict directly (e.g. from cv_config.json).
|
||||
If provided, case_study_id and setup_path are ignored.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict]
|
||||
List of split dicts with keys: ``fold``, ``train_start``,
|
||||
``train_end``, ``val_start``, ``val_end``.
|
||||
"""
|
||||
from ml4t.diagnostic.splitters import WalkForwardCV
|
||||
from ml4t.diagnostic.splitters.config import WalkForwardConfig as LibWalkForwardConfig
|
||||
|
||||
# Legacy path: pre-computed explicit splits
|
||||
if cv_config is not None and "splits" in cv_config:
|
||||
return cv_config["splits"]
|
||||
|
||||
# Normalize label buffer (strip ISO prefix, convert M → days)
|
||||
label_buffer = _normalize_label_buffer(label_buffer)
|
||||
|
||||
# Load evaluation config
|
||||
if cv_config is not None:
|
||||
# Legacy cv_config dict
|
||||
test_size_key = "val_size" if "val_size" in cv_config else "test_size"
|
||||
holdout_start_key = "holdout_start" if "holdout_start" in cv_config else "test_start"
|
||||
holdout_end_key = "holdout_end" if "holdout_end" in cv_config else "test_end"
|
||||
eval_config = {
|
||||
"n_splits": cv_config["n_splits"],
|
||||
"train_size": str(cv_config["train_size"]),
|
||||
"val_size": str(cv_config[test_size_key]),
|
||||
"holdout_start": cv_config.get(holdout_start_key),
|
||||
"holdout_end": cv_config.get(holdout_end_key),
|
||||
"calendar": cv_config.get("calendar"),
|
||||
}
|
||||
elif setup_path is not None:
|
||||
with open(setup_path) as f:
|
||||
setup = yaml.safe_load(f)
|
||||
eval_config = dict(setup["evaluation"])
|
||||
elif case_study_id is not None:
|
||||
eval_config = load_evaluation_config(case_study_id)
|
||||
else:
|
||||
raise ValueError("Provide either case_study_id, setup_path, or cv_config")
|
||||
|
||||
# Map calendar name to library exchange name
|
||||
calendar_id = _map_calendar_id(eval_config.get("calendar"))
|
||||
|
||||
# For D-unit buffers with a calendar, pass label_horizon as int so the
|
||||
# library interprets it as trading days (not calendar days). This fixes
|
||||
# the under-buffering where "21D" → pd.Timedelta("21 days") → ~15 trading
|
||||
# days instead of the intended 21 trading days.
|
||||
label_horizon: int | str = label_buffer
|
||||
if calendar_id is not None:
|
||||
d_match = re.match(r"^(\d+)D$", label_buffer)
|
||||
if d_match:
|
||||
label_horizon = int(d_match.group(1))
|
||||
|
||||
# Build WalkForwardConfig (library Pydantic model)
|
||||
config = LibWalkForwardConfig(
|
||||
n_splits=eval_config["n_splits"],
|
||||
train_size=_normalize_duration(str(eval_config["train_size"])),
|
||||
val_size=_normalize_duration(str(eval_config["val_size"])),
|
||||
holdout_start=eval_config.get("holdout_start"),
|
||||
holdout_end=eval_config.get("holdout_end"),
|
||||
label_horizon=label_horizon,
|
||||
calendar_id=calendar_id,
|
||||
fold_direction="backward",
|
||||
)
|
||||
|
||||
# Extract sorted unique timestamps from the dataset
|
||||
if isinstance(dataset, pl.DataFrame):
|
||||
unique_ts = dataset.select(date_col).unique().sort(date_col).to_series().to_pandas()
|
||||
else:
|
||||
unique_ts = pd.Series(sorted(dataset[date_col].dropna().unique()))
|
||||
|
||||
if len(unique_ts) == 0:
|
||||
raise ValueError("No timestamps found in dataset")
|
||||
|
||||
# Build a single-column DataFrame with DatetimeIndex for the splitter
|
||||
ts_index = pd.DatetimeIndex(unique_ts)
|
||||
input_tz_naive = ts_index.tz is None
|
||||
if input_tz_naive:
|
||||
ts_index = ts_index.tz_localize("UTC")
|
||||
ts_df = pd.DataFrame(
|
||||
{"_dummy": np.zeros(len(ts_index), dtype=np.int8)},
|
||||
index=ts_index,
|
||||
)
|
||||
|
||||
# Create WalkForwardCV with rolling window (expanding=False)
|
||||
cv = WalkForwardCV(config=config)
|
||||
cv.expanding = False
|
||||
|
||||
# Generate splits and extract date boundaries.
|
||||
# Match tz-awareness to the input data so comparisons work.
|
||||
def _ts(idx):
|
||||
t = ts_index[idx]
|
||||
return t.tz_localize(None) if input_tz_naive else t
|
||||
|
||||
splits = []
|
||||
for fold_i, (train_idx, val_idx) in enumerate(cv.split(ts_df)):
|
||||
splits.append(
|
||||
{
|
||||
"fold": fold_i,
|
||||
"train_start": _ts(train_idx[0]),
|
||||
"train_end": _ts(train_idx[-1]),
|
||||
"val_start": _ts(val_idx[0]),
|
||||
"val_end": _ts(val_idx[-1]),
|
||||
}
|
||||
)
|
||||
|
||||
return splits
|
||||
@@ -0,0 +1,567 @@
|
||||
"""Data quality and filtering utilities for data loading.
|
||||
|
||||
This module provides centralized functions for:
|
||||
- Coverage summaries (rows, symbols, date range)
|
||||
- OHLC invariant checks
|
||||
- Null rate analysis
|
||||
- Gap detection in time series
|
||||
- Symbol subsetting for test-mode execution
|
||||
|
||||
Usage:
|
||||
>>> from utils.data_quality import describe_coverage, check_ohlc_invariants
|
||||
>>> coverage = describe_coverage(df, time_col="timestamp", asset_col="symbol")
|
||||
>>> invariants = check_ohlc_invariants(df)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import polars as pl
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
def apply_max_symbols(
|
||||
data: pl.DataFrame | pl.LazyFrame,
|
||||
max_symbols: int,
|
||||
symbol_col: str = "symbol",
|
||||
seed: int = 42,
|
||||
) -> pl.DataFrame | pl.LazyFrame:
|
||||
"""Limit data to a random subset of symbols for fast-path testing.
|
||||
|
||||
Selects a reproducible random sample of symbols using a fixed seed.
|
||||
Returns data unchanged if max_symbols <= 0 or >= total symbols.
|
||||
"""
|
||||
if max_symbols <= 0:
|
||||
return data
|
||||
|
||||
if isinstance(data, pl.LazyFrame):
|
||||
all_symbols = data.select(pl.col(symbol_col).unique()).collect()[symbol_col].to_list()
|
||||
else:
|
||||
all_symbols = data[symbol_col].unique().to_list()
|
||||
|
||||
if max_symbols >= len(all_symbols):
|
||||
return data
|
||||
|
||||
rng = random.Random(seed)
|
||||
selected = rng.sample(sorted(all_symbols), max_symbols)
|
||||
return data.filter(pl.col(symbol_col).is_in(selected))
|
||||
|
||||
|
||||
def describe_coverage(
|
||||
df: pl.DataFrame,
|
||||
time_col: str = "timestamp",
|
||||
asset_col: str = "symbol",
|
||||
) -> dict:
|
||||
"""Return coverage summary for a dataset.
|
||||
|
||||
Args:
|
||||
df: DataFrame with time and asset columns
|
||||
time_col: Name of the timestamp/date column
|
||||
asset_col: Name of the asset identifier column
|
||||
|
||||
Returns:
|
||||
Dictionary with rows, assets, time_min, time_max, unique_times
|
||||
"""
|
||||
return {
|
||||
"rows": df.height,
|
||||
"assets": df[asset_col].n_unique() if asset_col in df.columns else 0,
|
||||
"time_min": df[time_col].min(),
|
||||
"time_max": df[time_col].max(),
|
||||
"unique_times": df[time_col].n_unique(),
|
||||
}
|
||||
|
||||
|
||||
def print_coverage(
|
||||
df: pl.DataFrame,
|
||||
time_col: str = "timestamp",
|
||||
asset_col: str = "symbol",
|
||||
dataset_name: str = "Dataset",
|
||||
) -> None:
|
||||
"""Print formatted coverage summary."""
|
||||
cov = describe_coverage(df, time_col, asset_col)
|
||||
print(f"=== {dataset_name} Coverage ===")
|
||||
print(f" Rows: {cov['rows']:,}")
|
||||
print(f" Assets: {cov['assets']:,}")
|
||||
print(f" Time range: {cov['time_min']} to {cov['time_max']}")
|
||||
print(f" Unique times: {cov['unique_times']:,}")
|
||||
|
||||
|
||||
def check_ohlc_invariants(
|
||||
df: pl.DataFrame,
|
||||
open_col: str = "open",
|
||||
high_col: str = "high",
|
||||
low_col: str = "low",
|
||||
close_col: str = "close",
|
||||
volume_col: str = "volume",
|
||||
) -> pl.DataFrame:
|
||||
"""Check OHLC data quality invariants.
|
||||
|
||||
Validates:
|
||||
- high >= low
|
||||
- high >= open
|
||||
- high >= close
|
||||
- low <= open
|
||||
- low <= close
|
||||
- volume >= 0 (if volume column exists)
|
||||
|
||||
For each check, only rows where all relevant columns are non-null are
|
||||
considered. This prevents null comparisons from distorting percentages
|
||||
(important for TAQ data where trade columns may be null for no-trade bars).
|
||||
|
||||
Args:
|
||||
df: DataFrame with OHLC columns
|
||||
open_col, high_col, low_col, close_col: Column names for OHLC
|
||||
volume_col: Column name for volume (optional)
|
||||
|
||||
Returns:
|
||||
DataFrame with check names and valid_pct columns
|
||||
"""
|
||||
results = []
|
||||
total_rows = df.height
|
||||
cols = set(df.columns)
|
||||
|
||||
def _check_invariant(name: str, condition: pl.Expr, required_cols: list[str]) -> None:
|
||||
"""Check an invariant on rows where all required columns are non-null."""
|
||||
# Filter to rows where all required columns are non-null
|
||||
not_null_filter = pl.all_horizontal([pl.col(c).is_not_null() for c in required_cols])
|
||||
applicable = df.filter(not_null_filter)
|
||||
n_applicable = applicable.height
|
||||
|
||||
if n_applicable == 0:
|
||||
return # Skip if no applicable rows
|
||||
|
||||
valid_pct = applicable.select(condition.mean()).item() * 100
|
||||
results.append(
|
||||
{
|
||||
"check": name,
|
||||
"valid_pct": valid_pct,
|
||||
"applicable_rows": n_applicable,
|
||||
"total_rows": total_rows,
|
||||
}
|
||||
)
|
||||
|
||||
# Define checks with their required columns
|
||||
if {high_col, low_col}.issubset(cols):
|
||||
_check_invariant(
|
||||
"high_gte_low",
|
||||
pl.col(high_col) >= pl.col(low_col),
|
||||
[high_col, low_col],
|
||||
)
|
||||
|
||||
if {high_col, open_col}.issubset(cols):
|
||||
_check_invariant(
|
||||
"high_gte_open",
|
||||
pl.col(high_col) >= pl.col(open_col),
|
||||
[high_col, open_col],
|
||||
)
|
||||
|
||||
if {high_col, close_col}.issubset(cols):
|
||||
_check_invariant(
|
||||
"high_gte_close",
|
||||
pl.col(high_col) >= pl.col(close_col),
|
||||
[high_col, close_col],
|
||||
)
|
||||
|
||||
if {low_col, open_col}.issubset(cols):
|
||||
_check_invariant(
|
||||
"low_lte_open",
|
||||
pl.col(low_col) <= pl.col(open_col),
|
||||
[low_col, open_col],
|
||||
)
|
||||
|
||||
if {low_col, close_col}.issubset(cols):
|
||||
_check_invariant(
|
||||
"low_lte_close",
|
||||
pl.col(low_col) <= pl.col(close_col),
|
||||
[low_col, close_col],
|
||||
)
|
||||
|
||||
if volume_col in cols:
|
||||
_check_invariant(
|
||||
"volume_non_negative",
|
||||
pl.col(volume_col) >= 0,
|
||||
[volume_col],
|
||||
)
|
||||
|
||||
if not results:
|
||||
return pl.DataFrame({"check": [], "valid_pct": [], "applicable_rows": [], "total_rows": []})
|
||||
|
||||
return pl.DataFrame(results)
|
||||
|
||||
|
||||
def print_ohlc_invariants(
|
||||
df: pl.DataFrame,
|
||||
open_col: str = "open",
|
||||
high_col: str = "high",
|
||||
low_col: str = "low",
|
||||
close_col: str = "close",
|
||||
volume_col: str = "volume",
|
||||
show_coverage: bool = False,
|
||||
) -> None:
|
||||
"""Print OHLC invariant check results.
|
||||
|
||||
Args:
|
||||
show_coverage: If True, show how many rows each check applies to
|
||||
"""
|
||||
result = check_ohlc_invariants(df, open_col, high_col, low_col, close_col, volume_col)
|
||||
print("=== OHLC Invariants ===")
|
||||
for row in result.iter_rows(named=True):
|
||||
status = "[OK]" if row["valid_pct"] >= 99.99 else "[WARN]"
|
||||
coverage = ""
|
||||
if show_coverage and row["applicable_rows"] < row["total_rows"]:
|
||||
coverage = f" ({row['applicable_rows']:,}/{row['total_rows']:,} rows)"
|
||||
print(f" {status} {row['check']}: {row['valid_pct']:.2f}%{coverage}")
|
||||
|
||||
|
||||
def null_rate(
|
||||
df: pl.DataFrame,
|
||||
cols: Sequence[str] | None = None,
|
||||
) -> pl.DataFrame:
|
||||
"""Calculate null rates for specified columns.
|
||||
|
||||
Args:
|
||||
df: DataFrame to analyze
|
||||
cols: Columns to check (default: all columns)
|
||||
|
||||
Returns:
|
||||
DataFrame with column names and null_pct
|
||||
"""
|
||||
if cols is None:
|
||||
cols = df.columns
|
||||
else:
|
||||
cols = [c for c in cols if c in df.columns]
|
||||
|
||||
if not cols:
|
||||
return pl.DataFrame({"column": [], "null_pct": []})
|
||||
|
||||
rates = df.select([pl.col(c).is_null().mean().alias(c) for c in cols])
|
||||
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"column": list(rates.columns),
|
||||
"null_pct": [rates[col].item() * 100 for col in rates.columns],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def print_null_rates(
|
||||
df: pl.DataFrame,
|
||||
cols: Sequence[str] | None = None,
|
||||
threshold: float = 0.0,
|
||||
) -> None:
|
||||
"""Print null rates for columns exceeding threshold.
|
||||
|
||||
Args:
|
||||
df: DataFrame to analyze
|
||||
cols: Columns to check (default: all columns)
|
||||
threshold: Only print columns with null_pct > threshold
|
||||
"""
|
||||
result = null_rate(df, cols)
|
||||
result = result.filter(pl.col("null_pct") > threshold)
|
||||
print("=== Null Rates ===")
|
||||
if result.height == 0:
|
||||
print(" No nulls detected")
|
||||
else:
|
||||
for row in result.iter_rows(named=True):
|
||||
print(f" {row['column']}: {row['null_pct']:.2f}%")
|
||||
|
||||
|
||||
def gap_summary(
|
||||
df: pl.DataFrame,
|
||||
time_col: str = "timestamp",
|
||||
group_col: str | None = "symbol",
|
||||
expected_delta: timedelta | None = None,
|
||||
) -> pl.DataFrame:
|
||||
"""Identify gaps in time series data.
|
||||
|
||||
Args:
|
||||
df: DataFrame with time series data
|
||||
time_col: Name of timestamp column
|
||||
group_col: Column to group by (e.g., symbol). None for ungrouped.
|
||||
expected_delta: Expected time between rows (e.g., timedelta(hours=1))
|
||||
|
||||
Returns:
|
||||
DataFrame with gap statistics per group (if grouped) or overall
|
||||
"""
|
||||
df_sorted = df.sort([group_col, time_col] if group_col else [time_col])
|
||||
|
||||
# Calculate time differences
|
||||
if group_col:
|
||||
df_gaps = df_sorted.with_columns(pl.col(time_col).diff().over(group_col).alias("time_diff"))
|
||||
else:
|
||||
df_gaps = df_sorted.with_columns(pl.col(time_col).diff().alias("time_diff"))
|
||||
|
||||
# If expected_delta provided, filter to gaps exceeding it
|
||||
if expected_delta is not None:
|
||||
df_gaps = df_gaps.filter(
|
||||
(pl.col("time_diff") > expected_delta) | pl.col("time_diff").is_null()
|
||||
)
|
||||
|
||||
# Aggregate
|
||||
if group_col:
|
||||
return (
|
||||
df_gaps.filter(pl.col("time_diff").is_not_null())
|
||||
.group_by(group_col)
|
||||
.agg(
|
||||
pl.len().alias("gap_count"),
|
||||
pl.col("time_diff").max().alias("max_gap"),
|
||||
)
|
||||
.sort(group_col)
|
||||
)
|
||||
else:
|
||||
gaps = df_gaps.filter(pl.col("time_diff").is_not_null())
|
||||
if gaps.height == 0:
|
||||
return pl.DataFrame({"gap_count": [0], "max_gap": [None]})
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"gap_count": [gaps.height],
|
||||
"max_gap": [gaps["time_diff"].max()],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def per_asset_stats(
|
||||
df: pl.DataFrame,
|
||||
time_col: str = "timestamp",
|
||||
asset_col: str = "symbol",
|
||||
price_col: str = "close",
|
||||
volume_col: str | None = "volume",
|
||||
) -> pl.DataFrame:
|
||||
"""Calculate per-asset summary statistics.
|
||||
|
||||
Args:
|
||||
df: DataFrame with time series data
|
||||
time_col: Timestamp column name
|
||||
asset_col: Asset identifier column name
|
||||
price_col: Price column for mean calculation
|
||||
volume_col: Volume column (optional)
|
||||
|
||||
Returns:
|
||||
DataFrame with rows, start, end, avg_price per asset
|
||||
"""
|
||||
aggs = [
|
||||
pl.len().alias("rows"),
|
||||
pl.col(time_col).min().alias("start"),
|
||||
pl.col(time_col).max().alias("end"),
|
||||
pl.col(price_col).mean().alias("avg_price"),
|
||||
]
|
||||
|
||||
if volume_col and volume_col in df.columns:
|
||||
aggs.append(pl.col(volume_col).mean().alias("avg_volume"))
|
||||
|
||||
return df.group_by(asset_col).agg(aggs).sort(asset_col)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Modeling pipeline quality gates
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def validate_prices(
|
||||
df: pl.DataFrame,
|
||||
price_cols: Sequence[str] = ("open", "high", "low", "close"),
|
||||
asset_col: str = "symbol",
|
||||
time_col: str = "timestamp",
|
||||
) -> list[str]:
|
||||
"""Check price columns for negative values, infinities, and NaN.
|
||||
|
||||
Returns a list of warning/error strings. Empty list = all clean.
|
||||
"""
|
||||
issues: list[str] = []
|
||||
cols_present = [c for c in price_cols if c in df.columns]
|
||||
|
||||
for col in cols_present:
|
||||
n_neg = df.filter(pl.col(col) < 0).height
|
||||
n_inf = df.filter(pl.col(col).is_infinite()).height
|
||||
n_nan = df.filter(pl.col(col).is_nan()).height
|
||||
|
||||
if n_neg > 0:
|
||||
# Show which assets have negative prices
|
||||
neg_assets = df.filter(pl.col(col) < 0).select(asset_col).unique().to_series().to_list()
|
||||
issues.append(
|
||||
f"CRITICAL: {col} has {n_neg} negative values "
|
||||
f"(assets: {neg_assets[:5]}{'...' if len(neg_assets) > 5 else ''})"
|
||||
)
|
||||
if n_inf > 0:
|
||||
issues.append(f"CRITICAL: {col} has {n_inf} infinite values")
|
||||
if n_nan > 0:
|
||||
issues.append(f"WARNING: {col} has {n_nan} NaN values")
|
||||
|
||||
return issues
|
||||
|
||||
|
||||
def validate_labels(
|
||||
df: pl.DataFrame,
|
||||
label_col: str,
|
||||
max_abs_return: float = 0.5,
|
||||
) -> list[str]:
|
||||
"""Check forward return labels for data quality issues.
|
||||
|
||||
Args:
|
||||
df: DataFrame containing the label column
|
||||
label_col: Name of the forward return column
|
||||
max_abs_return: Maximum plausible absolute return (e.g., 0.5 = 50%)
|
||||
|
||||
Returns list of warning/error strings.
|
||||
"""
|
||||
issues: list[str] = []
|
||||
vals = df[label_col].drop_nulls()
|
||||
|
||||
n_inf = vals.filter(vals.is_infinite()).len()
|
||||
n_nan = vals.filter(vals.is_nan()).len()
|
||||
n_extreme = vals.filter(vals.abs() > max_abs_return).len()
|
||||
n_total = vals.len()
|
||||
|
||||
if n_inf > 0:
|
||||
issues.append(f"CRITICAL: {label_col} has {n_inf} infinite values")
|
||||
if n_nan > 0:
|
||||
issues.append(f"CRITICAL: {label_col} has {n_nan} NaN values")
|
||||
if n_extreme > 0:
|
||||
pct = n_extreme / n_total * 100
|
||||
issues.append(
|
||||
f"WARNING: {label_col} has {n_extreme} values with |ret| > {max_abs_return:.0%} "
|
||||
f"({pct:.2f}% of {n_total:,} rows)"
|
||||
)
|
||||
|
||||
return issues
|
||||
|
||||
|
||||
def validate_features(
|
||||
df: pl.DataFrame,
|
||||
feature_cols: Sequence[str],
|
||||
max_abs_value: float = 1e6,
|
||||
) -> list[str]:
|
||||
"""Check feature columns for infinities, all-null, and extreme values.
|
||||
|
||||
Args:
|
||||
df: DataFrame containing feature columns
|
||||
feature_cols: List of feature column names to validate
|
||||
max_abs_value: Threshold for flagging extreme values
|
||||
|
||||
Returns list of warning/error strings.
|
||||
"""
|
||||
issues: list[str] = []
|
||||
n_rows = df.height
|
||||
|
||||
inf_cols = []
|
||||
null_cols = []
|
||||
extreme_cols = []
|
||||
|
||||
for col in feature_cols:
|
||||
if col not in df.columns:
|
||||
continue
|
||||
|
||||
series = df[col]
|
||||
n_null = series.null_count()
|
||||
non_null = series.drop_nulls()
|
||||
|
||||
if n_null == n_rows:
|
||||
null_cols.append(col)
|
||||
continue
|
||||
|
||||
if non_null.len() > 0:
|
||||
n_inf = non_null.filter(non_null.is_infinite()).len()
|
||||
if n_inf > 0:
|
||||
inf_cols.append((col, n_inf))
|
||||
|
||||
n_extreme = non_null.filter(non_null.abs() > max_abs_value).len()
|
||||
if n_extreme > 0:
|
||||
extreme_cols.append((col, n_extreme))
|
||||
|
||||
if inf_cols:
|
||||
details = ", ".join(f"{c}({n})" for c, n in inf_cols[:10])
|
||||
issues.append(f"CRITICAL: {len(inf_cols)} features have infinite values: {details}")
|
||||
|
||||
if null_cols:
|
||||
issues.append(
|
||||
f"WARNING: {len(null_cols)} features are entirely null: "
|
||||
f"{null_cols[:10]}{'...' if len(null_cols) > 10 else ''}"
|
||||
)
|
||||
|
||||
if extreme_cols:
|
||||
details = ", ".join(f"{c}({n})" for c, n in extreme_cols[:10])
|
||||
issues.append(
|
||||
f"WARNING: {len(extreme_cols)} features have values |x| > {max_abs_value:.0e}: {details}"
|
||||
)
|
||||
|
||||
return issues
|
||||
|
||||
|
||||
def validate_modeling_inputs(
|
||||
features_df: pl.DataFrame,
|
||||
label_df: pl.DataFrame,
|
||||
feature_cols: Sequence[str],
|
||||
label_col: str,
|
||||
join_cols: Sequence[str] = ("timestamp", "symbol"),
|
||||
price_cols: Sequence[str] = (),
|
||||
asset_col: str = "symbol",
|
||||
max_abs_return: float = 0.5,
|
||||
max_abs_feature: float = 1e6,
|
||||
fail_on_critical: bool = True,
|
||||
) -> dict:
|
||||
"""Run all data quality checks before modeling.
|
||||
|
||||
This is the gate between data preparation (labels + features) and
|
||||
model training. Call this at the start of evaluation notebooks.
|
||||
|
||||
Args:
|
||||
features_df: Feature DataFrame
|
||||
label_df: Label DataFrame with forward returns
|
||||
feature_cols: Feature column names to validate
|
||||
label_col: Forward return column name
|
||||
join_cols: Columns used to join features and labels
|
||||
price_cols: Price columns to check (if present in features_df)
|
||||
asset_col: Asset identifier column name
|
||||
max_abs_return: Max plausible absolute return for labels
|
||||
max_abs_feature: Max plausible absolute feature value
|
||||
fail_on_critical: If True, raise ValueError on CRITICAL issues
|
||||
|
||||
Returns:
|
||||
Dict with 'issues' (list of strings), 'n_critical', 'n_warning'
|
||||
|
||||
Raises:
|
||||
ValueError: If fail_on_critical=True and any CRITICAL issues found
|
||||
"""
|
||||
all_issues: list[str] = []
|
||||
|
||||
# 1. Price checks (if price columns present)
|
||||
if price_cols:
|
||||
all_issues.extend(validate_prices(features_df, price_cols, asset_col=asset_col))
|
||||
|
||||
# 2. Label checks
|
||||
all_issues.extend(validate_labels(label_df, label_col, max_abs_return))
|
||||
|
||||
# 3. Feature checks
|
||||
all_issues.extend(validate_features(features_df, feature_cols, max_abs_feature))
|
||||
|
||||
# Summarize
|
||||
n_critical = sum(1 for i in all_issues if i.startswith("CRITICAL"))
|
||||
n_warning = sum(1 for i in all_issues if i.startswith("WARNING"))
|
||||
|
||||
# Print results
|
||||
if all_issues:
|
||||
print(f"Data Quality Gate: {n_critical} CRITICAL, {n_warning} WARNING")
|
||||
for issue in all_issues:
|
||||
marker = "[X]" if issue.startswith("CRITICAL") else "[!]"
|
||||
print(f" {marker} {issue}")
|
||||
else:
|
||||
print("Data Quality Gate: ALL CLEAR")
|
||||
|
||||
result = {
|
||||
"issues": all_issues,
|
||||
"n_critical": n_critical,
|
||||
"n_warning": n_warning,
|
||||
}
|
||||
|
||||
if fail_on_critical and n_critical > 0:
|
||||
raise ValueError(
|
||||
f"Data quality gate FAILED: {n_critical} critical issues. "
|
||||
f"Fix upstream data before modeling."
|
||||
)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,369 @@
|
||||
"""Shared utilities for ML4T data download scripts.
|
||||
|
||||
Provides:
|
||||
- Standardized argument parsing with --dry-run, --data-path, --force, --verbose
|
||||
- Standardized path resolution via utils.config
|
||||
- YAML config helpers for dataset download scripts
|
||||
- Import checking with helpful errors
|
||||
- Consistent output formatting
|
||||
- Atomic file writes
|
||||
- Download summary reporting
|
||||
- DataBento cost acknowledgment
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_repo_root = Path(__file__).parent.parent
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def resolve_data_dir(cli_arg: Path | None = None) -> Path:
|
||||
"""Resolve data directory with standardized precedence.
|
||||
|
||||
Priority:
|
||||
1. CLI argument (highest)
|
||||
2. ML4T_DATA_PATH environment variable (from utils.config)
|
||||
3. <repo>/data (default)
|
||||
"""
|
||||
if cli_arg is not None:
|
||||
path = Path(cli_arg).expanduser().resolve()
|
||||
print(f"Using data path (CLI): {path}")
|
||||
return path
|
||||
|
||||
try:
|
||||
from utils.config import ML4T_DATA_PATH
|
||||
|
||||
print(f"Using data path (ML4T_DATA_PATH): {ML4T_DATA_PATH}")
|
||||
return ML4T_DATA_PATH
|
||||
except (ImportError, FileNotFoundError):
|
||||
pass
|
||||
|
||||
env_dir = os.environ.get("ML4T_DATA_PATH")
|
||||
if env_dir:
|
||||
path = Path(env_dir).expanduser().resolve()
|
||||
print(f"Using data path (env): {path}")
|
||||
return path
|
||||
|
||||
default_path = _repo_root / "data"
|
||||
print(f"Using data path (default): {default_path}")
|
||||
default_path.mkdir(parents=True, exist_ok=True)
|
||||
return default_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dotenv(env_file: Path | None = None):
|
||||
"""Load environment variables from .env file."""
|
||||
if env_file is None:
|
||||
env_file = _repo_root / ".env"
|
||||
|
||||
if not env_file.exists():
|
||||
return
|
||||
|
||||
with open(env_file) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
|
||||
key, value = line.split("=", 1)
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
|
||||
if value and not os.getenv(key):
|
||||
os.environ[key] = value
|
||||
|
||||
|
||||
def require_env(var: str, hint: str | None = None) -> str:
|
||||
"""Get required environment variable or exit with helpful message."""
|
||||
value = os.getenv(var)
|
||||
if not value:
|
||||
print(f"ERROR: {var} not set")
|
||||
if hint:
|
||||
print(f" {hint}")
|
||||
print(f" Add to .env file: {var}=your-value")
|
||||
sys.exit(1)
|
||||
assert value is not None
|
||||
return value
|
||||
|
||||
|
||||
def check_import(module: str, install_hint: str):
|
||||
"""Check if module can be imported, exit with helpful message if not."""
|
||||
try:
|
||||
__import__(module)
|
||||
except ImportError as e:
|
||||
print(f"ERROR: {module} not available")
|
||||
print(f" Run: {install_hint}")
|
||||
print(f" ({e})")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# YAML config helpers (from book_config)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_section(config_path: str | Path, section: str) -> dict[str, Any]:
|
||||
"""Load a top-level YAML section from a config file."""
|
||||
import yaml
|
||||
|
||||
path = Path(config_path).expanduser()
|
||||
with open(path) as f:
|
||||
raw = yaml.safe_load(f) or {}
|
||||
return raw.get(section, {})
|
||||
|
||||
|
||||
def resolve_storage_path(data_root: Path, configured_path: str | None, fallback: str) -> Path:
|
||||
"""Resolve storage path relative to the selected ML4T data root."""
|
||||
raw_path = configured_path or fallback
|
||||
path = Path(raw_path).expanduser()
|
||||
return path if path.is_absolute() else data_root / path
|
||||
|
||||
|
||||
def flatten_group_values(groups: dict[str, Any], values_key: str) -> list[str]:
|
||||
"""Flatten grouped config values like symbols or pairs into a unique ordered list."""
|
||||
values: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for group in groups.values():
|
||||
if not isinstance(group, dict):
|
||||
continue
|
||||
for value in group.get(values_key, []):
|
||||
if value not in seen:
|
||||
values.append(value)
|
||||
seen.add(value)
|
||||
|
||||
return values
|
||||
|
||||
|
||||
def save_dataset_profile(
|
||||
df,
|
||||
data_path: str | Path,
|
||||
*,
|
||||
source: str,
|
||||
timestamp_col: str = "timestamp",
|
||||
symbol_col: str | None = "symbol",
|
||||
) -> Path:
|
||||
"""Generate and save a dataset profile next to a data file."""
|
||||
from ml4t.data.storage.data_profile import generate_profile, get_profile_path, save_profile
|
||||
|
||||
path = Path(data_path)
|
||||
profile = generate_profile(
|
||||
df, source=source, timestamp_col=timestamp_col, symbol_col=symbol_col
|
||||
)
|
||||
profile_path = get_profile_path(path)
|
||||
save_profile(profile, profile_path)
|
||||
return profile_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DataBento
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def patch_databento_symbology():
|
||||
"""Patch databento 0.72.0 bug where insert_metadata expects 'asset' key."""
|
||||
try:
|
||||
import databento.common.symbology as sym
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
if getattr(sym.InstrumentMap.insert_metadata, "_ml4t_patched", False):
|
||||
return
|
||||
|
||||
_orig = sym.InstrumentMap.insert_metadata
|
||||
|
||||
def _patched(self, metadata):
|
||||
mappings = metadata.mappings
|
||||
if mappings:
|
||||
for _symbol_in, entries in mappings.items():
|
||||
for entry in entries:
|
||||
if "asset" not in entry and "symbol" in entry:
|
||||
entry["asset"] = entry["symbol"]
|
||||
|
||||
class _PatchedMeta:
|
||||
"""Thin wrapper that returns fixed mappings."""
|
||||
|
||||
def __init__(self, orig, fixed_mappings):
|
||||
self._orig = orig
|
||||
self._fixed = fixed_mappings
|
||||
|
||||
@property
|
||||
def mappings(self):
|
||||
return self._fixed
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._orig, name)
|
||||
|
||||
return _orig(self, _PatchedMeta(metadata, mappings))
|
||||
|
||||
_patched._ml4t_patched = True
|
||||
sym.InstrumentMap.insert_metadata = _patched
|
||||
|
||||
|
||||
DATABENTO_WARNING = """
|
||||
================================================================================
|
||||
DATABENTO API - PAID SERVICE WARNING
|
||||
================================================================================
|
||||
|
||||
This download uses the DataBento API which is a PAID service.
|
||||
|
||||
IMPORTANT INFORMATION:
|
||||
- DataBento requires registration with a credit card on file
|
||||
- As of February 2026, DataBento offers a $125 sign-up credit for new accounts
|
||||
- This credit is sufficient to cover the ML4T book datasets:
|
||||
* CME Futures (continuous contracts): ~$75
|
||||
* MBO Tick Data (3 symbols, 10 days): ~$10-15
|
||||
|
||||
WARNINGS:
|
||||
- If you have already used your sign-up credit on other downloads,
|
||||
this download WILL BE CHARGED to your credit card
|
||||
- DataBento may change or remove the sign-up credit at any time
|
||||
- Cost estimates are approximate; actual costs may vary
|
||||
- YOU are responsible for managing your DataBento account and costs
|
||||
|
||||
ESTIMATED COST FOR THIS DOWNLOAD: ${estimated_cost:.2f}
|
||||
|
||||
================================================================================
|
||||
"""
|
||||
|
||||
|
||||
def databento_acknowledge(estimated_cost: float, force: bool = False) -> bool:
|
||||
"""Display DataBento cost warning and require explicit acknowledgment."""
|
||||
print(DATABENTO_WARNING.format(estimated_cost=estimated_cost))
|
||||
|
||||
if force:
|
||||
print("--force flag set: Proceeding without interactive confirmation.")
|
||||
print("By using --force, you acknowledge the above warnings.")
|
||||
return True
|
||||
|
||||
print("To proceed with this download, type exactly: I UNDERSTAND")
|
||||
print("To cancel, press Ctrl+C or type anything else.")
|
||||
print()
|
||||
|
||||
try:
|
||||
response = input("Your response: ").strip()
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nDownload cancelled.")
|
||||
return False
|
||||
|
||||
if response == "I UNDERSTAND":
|
||||
print("\nAcknowledgment received. Proceeding with download...")
|
||||
return True
|
||||
else:
|
||||
print(f"\nResponse '{response}' does not match 'I UNDERSTAND'.")
|
||||
print("Download cancelled for your protection.")
|
||||
return False
|
||||
|
||||
|
||||
def databento_estimate_only_notice(estimated_cost: float) -> None:
|
||||
"""Print cost estimate without prompting for download."""
|
||||
print("\n" + "=" * 70)
|
||||
print("DATABENTO COST ESTIMATE (No download - estimate only)")
|
||||
print("=" * 70)
|
||||
print(f"\n Estimated cost: ${estimated_cost:.2f}")
|
||||
print()
|
||||
print(" Note: As of Feb 2026, new DataBento accounts receive $125 credit.")
|
||||
print(" If your credit is exhausted, this amount will be charged.")
|
||||
print()
|
||||
print(" To proceed with download, run without --estimate-only flag.")
|
||||
print("=" * 70 + "\n")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output formatting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_section(title: str, char: str = "=", width: int = 60):
|
||||
"""Print a formatted section header."""
|
||||
print("\n" + char * width)
|
||||
print(title)
|
||||
print(char * width)
|
||||
|
||||
|
||||
def atomic_write_parquet(df, path: Path):
|
||||
"""Write Polars DataFrame to parquet with atomic rename."""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
tmp_file = path.parent / f".{path.name}.tmp"
|
||||
|
||||
try:
|
||||
df.write_parquet(tmp_file)
|
||||
tmp_file.replace(path)
|
||||
except Exception as e:
|
||||
if tmp_file.exists():
|
||||
tmp_file.unlink()
|
||||
raise e
|
||||
|
||||
|
||||
def get_repo_root() -> Path:
|
||||
"""Get repository root directory."""
|
||||
return _repo_root
|
||||
|
||||
|
||||
def create_base_parser(description: str) -> argparse.ArgumentParser:
|
||||
"""Create argument parser with standard ML4T download flags."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description=description,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data-path",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Data storage location (default: $ML4T_DATA_PATH or repo/data)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Show what would be done without downloading",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Force re-download even if data exists",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
"-v",
|
||||
action="store_true",
|
||||
help="Verbose output",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def print_download_summary(stats: dict, dry_run: bool = False) -> None:
|
||||
"""Print standardized download summary."""
|
||||
prefix = "[DRY RUN] " if dry_run else ""
|
||||
print_section(f"{prefix}SUMMARY")
|
||||
for key, value in stats.items():
|
||||
display_key = key.replace("_", " ").title()
|
||||
if isinstance(value, int) and value > 1000:
|
||||
print(f" {display_key}: {value:,}")
|
||||
elif isinstance(value, float):
|
||||
print(f" {display_key}: {value:.2f}")
|
||||
else:
|
||||
print(f" {display_key}: {value}")
|
||||
|
||||
|
||||
def print_dry_run_notice() -> None:
|
||||
"""Print notice that this is a dry run."""
|
||||
print("\n" + "=" * 60)
|
||||
print("DRY RUN - No data will be downloaded")
|
||||
print("Remove --dry-run to actually download")
|
||||
print("=" * 60 + "\n")
|
||||
+1312
File diff suppressed because it is too large
Load Diff
+392
@@ -0,0 +1,392 @@
|
||||
"""Chapter and data path management for ML4T notebooks.
|
||||
|
||||
This module provides:
|
||||
1. CHAPTERS - Canonical registry of chapter numbers to directory names
|
||||
2. get_chapter_dir() - Type-safe chapter path access
|
||||
3. get_case_study_dir() - Centralized case study artifact store
|
||||
4. get_output_dir() - Chapter-local outputs (non-case-study data)
|
||||
5. CH01-CH27 - Direct chapter path constants
|
||||
|
||||
Naming Convention:
|
||||
- Dataset IDs: etf, crypto, futures, fx, equities_us, equities_nasdaq100, options_sp500
|
||||
- Strategy IDs: etfs, crypto_perps_funding, cme_futures, fx_pairs, us_equities_panel,
|
||||
nasdaq100_microstructure, sp500_equity_option_analytics, sp500_options,
|
||||
us_firm_characteristics
|
||||
|
||||
Usage:
|
||||
from utils.paths import get_case_study_dir
|
||||
|
||||
# Write features in Ch7
|
||||
CASE_DIR = get_case_study_dir("etfs")
|
||||
features.write_parquet(CASE_DIR / "features" / "features.parquet")
|
||||
|
||||
# Read features in Ch9
|
||||
CASE_DIR = get_case_study_dir("etfs")
|
||||
features = pl.read_parquet(CASE_DIR / "features" / "features.parquet")
|
||||
predictions.write_parquet(CASE_DIR / "models" / "linear" / "predictions.parquet")
|
||||
|
||||
# Chapter-local outputs (benchmarks, demonstrations - NOT case studies)
|
||||
from utils.paths import get_output_dir
|
||||
OUTPUT_DIR = get_output_dir(7, "benchmark_results")
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
# =============================================================================
|
||||
# Core Paths
|
||||
# =============================================================================
|
||||
|
||||
REPO_ROOT = Path(__file__).parent.parent.resolve()
|
||||
|
||||
# =============================================================================
|
||||
# Chapter Registry (Single Source of Truth)
|
||||
# =============================================================================
|
||||
|
||||
CHAPTERS: dict[int, str] = {
|
||||
1: "01_process_is_edge",
|
||||
2: "02_financial_data_universe",
|
||||
3: "03_market_microstructure",
|
||||
4: "04_fundamental_alternative_data",
|
||||
5: "05_synthetic_data",
|
||||
6: "06_strategy_definition",
|
||||
7: "07_defining_the_learning_task",
|
||||
8: "08_financial_features",
|
||||
9: "09_model_based_features",
|
||||
10: "10_text_feature_engineering",
|
||||
11: "11_ml_pipeline",
|
||||
12: "12_gradient_boosting",
|
||||
13: "13_dl_time_series",
|
||||
14: "14_latent_factors",
|
||||
15: "15_causal_estimation",
|
||||
16: "16_strategy_simulation",
|
||||
17: "17_portfolio_construction",
|
||||
18: "18_transaction_costs",
|
||||
19: "19_risk_management",
|
||||
20: "20_strategy_synthesis",
|
||||
21: "21_rl_execution_hedging",
|
||||
22: "22_rag_financial_research",
|
||||
23: "23_knowledge_graphs",
|
||||
24: "24_autonomous_agents",
|
||||
25: "25_live_trading",
|
||||
26: "26_mlops_governance",
|
||||
27: "27_systematic_edge",
|
||||
}
|
||||
|
||||
# Sub-chapters (not in main sequence)
|
||||
SUB_CHAPTERS: dict[str, str] = {}
|
||||
|
||||
# Strategy IDs (Tier 3 naming)
|
||||
# Updated 2026-02-12 to match case_studies/ directory structure
|
||||
STRATEGY_IDS = frozenset(
|
||||
{
|
||||
"etfs",
|
||||
"crypto_perps_funding",
|
||||
"nasdaq100_microstructure",
|
||||
"sp500_equity_option_analytics",
|
||||
"us_firm_characteristics",
|
||||
"fx_pairs",
|
||||
"cme_futures",
|
||||
"sp500_options",
|
||||
"us_equities_panel",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Stage mapping: chapter number → subdirectory within case_studies/{strategy}/
|
||||
# Used as documentation reference; notebooks use string literals directly.
|
||||
# Updated 2026-02-25 for 27-chapter scheme (old Ch16+21 merged → Ch20)
|
||||
CHAPTER_STAGES: dict[int, str] = {
|
||||
6: "exploration",
|
||||
7: "labels",
|
||||
8: "features",
|
||||
9: "features", # Temporal features
|
||||
10: "features/text",
|
||||
11: "models/linear",
|
||||
12: "models/gbm",
|
||||
13: "models/deep_learning",
|
||||
14: "models/latent_factors",
|
||||
15: "models/causal",
|
||||
16: "backtest/simulation",
|
||||
17: "backtest/portfolio",
|
||||
18: "backtest/costs",
|
||||
19: "backtest/risk",
|
||||
20: "synthesis",
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Chapter Path Functions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def display_path(path: Path | str) -> str:
|
||||
"""Return `path` relative to REPO_ROOT when possible, else as-is.
|
||||
|
||||
Use in `print(...)` statements inside notebooks so the committed cell
|
||||
output never bakes in machine-specific absolute paths (e.g. `/home/<user>/...`).
|
||||
"""
|
||||
p = Path(path)
|
||||
try:
|
||||
return str(p.relative_to(REPO_ROOT))
|
||||
except ValueError:
|
||||
return str(p)
|
||||
|
||||
|
||||
def get_chapter_dir(chapter: int | str) -> Path:
|
||||
"""Get the directory path for a chapter.
|
||||
|
||||
Args:
|
||||
chapter: Chapter number (1-27)
|
||||
|
||||
Returns:
|
||||
Absolute path to chapter directory
|
||||
|
||||
Raises:
|
||||
ValueError: If chapter number is invalid
|
||||
|
||||
Examples:
|
||||
>>> get_chapter_dir(7)
|
||||
ML4T_PATH / "code/07_feature_engineering"
|
||||
"""
|
||||
if isinstance(chapter, int):
|
||||
if chapter not in CHAPTERS:
|
||||
raise ValueError(f"Invalid chapter number: {chapter}. Valid: 1-27")
|
||||
return REPO_ROOT / CHAPTERS[chapter]
|
||||
|
||||
# String: check sub-chapters
|
||||
if chapter in SUB_CHAPTERS:
|
||||
return REPO_ROOT / SUB_CHAPTERS[chapter]
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid chapter: {chapter!r}. "
|
||||
f"Use integers 1-27 or sub-chapter IDs: {list(SUB_CHAPTERS.keys())}"
|
||||
)
|
||||
|
||||
|
||||
def get_output_dir(
|
||||
chapter: int | str,
|
||||
strategy_id: str,
|
||||
*,
|
||||
create: bool = True,
|
||||
) -> Path:
|
||||
"""Get output directory for cross-chapter data flow.
|
||||
|
||||
This is the primary function for case study data that flows between chapters:
|
||||
- Ch7 creates features.parquet -> Ch9 reads for signal evaluation
|
||||
- Ch9 creates signals.parquet -> Ch12 reads for ML pipeline
|
||||
- Ch12 creates predictions.parquet -> Ch18 reads for backtest
|
||||
|
||||
When ML4T_OUTPUT_DIR is set (e.g., by pytest), outputs are redirected to
|
||||
a temporary directory to prevent tests from overwriting production data.
|
||||
|
||||
Args:
|
||||
chapter: Chapter number (1-27) or sub-chapter ID
|
||||
strategy_id: Strategy identifier (e.g., "etfs", "crypto_perps_funding")
|
||||
create: If True, create directory if it doesn't exist
|
||||
|
||||
Returns:
|
||||
Path to output directory: {chapter_dir}/output/{strategy_id}/
|
||||
|
||||
Examples:
|
||||
>>> get_output_dir(7, "etfs")
|
||||
PosixPath('.../07_defining_the_learning_task/output/etfs')
|
||||
|
||||
>>> # Read from previous chapter
|
||||
>>> features = pl.read_parquet(get_output_dir(7, "etfs") / "features.parquet")
|
||||
|
||||
>>> # In test mode with ML4T_OUTPUT_DIR=/tmp/test
|
||||
>>> get_output_dir(7, "etfs")
|
||||
PosixPath('/tmp/test/ch07_etfs')
|
||||
"""
|
||||
import os
|
||||
|
||||
# Test mode: allow chapter outputs to redirect independently from case studies.
|
||||
test_output = os.environ.get("ML4T_CHAPTER_OUTPUT_DIR") or os.environ.get("ML4T_OUTPUT_DIR")
|
||||
if test_output:
|
||||
# Get chapter number for directory naming
|
||||
if isinstance(chapter, int):
|
||||
ch_num = chapter
|
||||
else:
|
||||
# Sub-chapter: use string as-is
|
||||
ch_num = chapter
|
||||
output_dir = (
|
||||
Path(test_output) / f"ch{ch_num:02d}_{strategy_id}"
|
||||
if isinstance(ch_num, int)
|
||||
else Path(test_output) / f"{ch_num}_{strategy_id}"
|
||||
)
|
||||
if create:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
return output_dir
|
||||
|
||||
# Production mode: normal chapter output
|
||||
chapter_dir = get_chapter_dir(chapter)
|
||||
output_dir = chapter_dir / "output" / strategy_id
|
||||
|
||||
if create:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return output_dir
|
||||
|
||||
|
||||
def get_case_study_source_dir(strategy_id: str) -> Path:
|
||||
"""Get the source-controlled case study directory, preferring sibling dev assets."""
|
||||
dev_case_dir = REPO_ROOT.parent / "dev" / "case_studies" / strategy_id
|
||||
if dev_case_dir.exists():
|
||||
return dev_case_dir
|
||||
return REPO_ROOT / "case_studies" / strategy_id
|
||||
|
||||
|
||||
def get_case_study_dir(strategy_id: str, *, create: bool = True) -> Path:
|
||||
"""Get the case study directory for a strategy.
|
||||
|
||||
Case studies are centralized under CASE_STUDIES_DIR (default: repo_root/case_studies/).
|
||||
Each strategy has its own directory with stage-based subdirectories:
|
||||
|
||||
case_studies/{strategy_id}/
|
||||
├── exploration/ # Ch6 priors & EDA
|
||||
├── labels/ # Ch7 labels & evaluation
|
||||
├── features/ # Ch8-10 feature engineering
|
||||
├── models/ # Ch11-15 ML models
|
||||
│ ├── linear/ # Ch11
|
||||
│ ├── gbm/ # Ch12
|
||||
│ ├── deep_learning/ # Ch13
|
||||
│ ├── latent_factors/ # Ch14
|
||||
│ └── causal/ # Ch15
|
||||
├── backtest/ # Ch16-19 strategy implementation
|
||||
│ ├── simulation/ # Ch16
|
||||
│ ├── portfolio/ # Ch17
|
||||
│ ├── costs/ # Ch18
|
||||
│ └── risk/ # Ch19
|
||||
└── synthesis/ # Ch20 strategy development synthesis
|
||||
|
||||
When ML4T_OUTPUT_DIR is set (e.g., by pytest), outputs are redirected to
|
||||
a temporary directory to prevent tests from overwriting production data.
|
||||
|
||||
Args:
|
||||
strategy_id: Strategy identifier (e.g., "etfs", "crypto_perps_funding")
|
||||
create: If True, create directory if it doesn't exist
|
||||
|
||||
Returns:
|
||||
Path to case study directory: case_studies/{strategy_id}/
|
||||
|
||||
Examples:
|
||||
>>> get_case_study_dir("etfs")
|
||||
PosixPath('.../case_studies/etfs')
|
||||
|
||||
>>> # Typical usage with stage subdirectory
|
||||
>>> CASE_DIR = get_case_study_dir("etfs")
|
||||
>>> features.write_parquet(CASE_DIR / "features" / "features.parquet")
|
||||
|
||||
>>> # In test mode with ML4T_OUTPUT_DIR=/tmp/test
|
||||
>>> get_case_study_dir("etfs")
|
||||
PosixPath('/tmp/test/etfs')
|
||||
"""
|
||||
import os
|
||||
|
||||
# Test mode: redirect to temp directory (prevents overwriting production data)
|
||||
test_output = os.environ.get("ML4T_OUTPUT_DIR")
|
||||
if test_output:
|
||||
output_dir = Path(test_output) / strategy_id
|
||||
else:
|
||||
from utils import CASE_STUDIES_DIR
|
||||
|
||||
output_dir = CASE_STUDIES_DIR / strategy_id
|
||||
|
||||
if create:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return output_dir
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Dataset IDs (Tier 2 naming)
|
||||
# =============================================================================
|
||||
|
||||
DATASET_IDS = frozenset(
|
||||
{
|
||||
"etf",
|
||||
"crypto",
|
||||
"futures",
|
||||
"fx",
|
||||
"equities_us",
|
||||
"equities_nasdaq100",
|
||||
"equities_sp500",
|
||||
"options_sp500",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Chapter Path Constants (for direct imports)
|
||||
# =============================================================================
|
||||
|
||||
CH01 = REPO_ROOT / CHAPTERS[1]
|
||||
CH02 = REPO_ROOT / CHAPTERS[2]
|
||||
CH03 = REPO_ROOT / CHAPTERS[3]
|
||||
CH04 = REPO_ROOT / CHAPTERS[4]
|
||||
CH05 = REPO_ROOT / CHAPTERS[5]
|
||||
CH06 = REPO_ROOT / CHAPTERS[6]
|
||||
CH07 = REPO_ROOT / CHAPTERS[7]
|
||||
CH08 = REPO_ROOT / CHAPTERS[8]
|
||||
CH09 = REPO_ROOT / CHAPTERS[9]
|
||||
CH10 = REPO_ROOT / CHAPTERS[10]
|
||||
CH11 = REPO_ROOT / CHAPTERS[11]
|
||||
CH12 = REPO_ROOT / CHAPTERS[12]
|
||||
CH13 = REPO_ROOT / CHAPTERS[13]
|
||||
CH14 = REPO_ROOT / CHAPTERS[14]
|
||||
CH15 = REPO_ROOT / CHAPTERS[15]
|
||||
CH16 = REPO_ROOT / CHAPTERS[16]
|
||||
CH17 = REPO_ROOT / CHAPTERS[17]
|
||||
CH18 = REPO_ROOT / CHAPTERS[18]
|
||||
CH19 = REPO_ROOT / CHAPTERS[19]
|
||||
CH20 = REPO_ROOT / CHAPTERS[20]
|
||||
CH21 = REPO_ROOT / CHAPTERS[21]
|
||||
CH22 = REPO_ROOT / CHAPTERS[22]
|
||||
CH23 = REPO_ROOT / CHAPTERS[23]
|
||||
CH24 = REPO_ROOT / CHAPTERS[24]
|
||||
CH25 = REPO_ROOT / CHAPTERS[25]
|
||||
CH26 = REPO_ROOT / CHAPTERS[26]
|
||||
CH27 = REPO_ROOT / CHAPTERS[27]
|
||||
|
||||
__all__ = [
|
||||
# Registry
|
||||
"CHAPTERS",
|
||||
"SUB_CHAPTERS",
|
||||
"STRATEGY_IDS",
|
||||
"CHAPTER_STAGES",
|
||||
"CASE_STUDIES", # Legacy alias
|
||||
"DATASET_IDS",
|
||||
"REPO_ROOT",
|
||||
# Functions
|
||||
"get_chapter_dir",
|
||||
"get_output_dir",
|
||||
"get_case_study_dir",
|
||||
# Constants
|
||||
"CH01",
|
||||
"CH02",
|
||||
"CH03",
|
||||
"CH04",
|
||||
"CH05",
|
||||
"CH06",
|
||||
"CH07",
|
||||
"CH08",
|
||||
"CH09",
|
||||
"CH10",
|
||||
"CH11",
|
||||
"CH12",
|
||||
"CH13",
|
||||
"CH14",
|
||||
"CH15",
|
||||
"CH16",
|
||||
"CH17",
|
||||
"CH18",
|
||||
"CH19",
|
||||
"CH20",
|
||||
"CH21",
|
||||
"CH22",
|
||||
"CH23",
|
||||
"CH24",
|
||||
"CH25",
|
||||
"CH26",
|
||||
"CH27",
|
||||
]
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Predictions cache for chapter teaching notebooks.
|
||||
|
||||
Cache long-form prediction frames keyed by a content-addressed spec hash so
|
||||
re-running visualization / interpretation cells doesn't require re-training
|
||||
the upstream stages. Cache files live under
|
||||
|
||||
{chapter_dir}/output/predictions/{notebook_id}/{spec_hash}.parquet
|
||||
|
||||
which is gitignored.
|
||||
|
||||
Frame schema (required columns):
|
||||
|
||||
date -- value identifying the prediction period
|
||||
symbol -- value identifying the asset
|
||||
y_pred -- float, model prediction
|
||||
y_true -- float, realised forward return
|
||||
|
||||
Optional column:
|
||||
|
||||
forecaster -- string label, when one notebook stacks multiple forecasters
|
||||
(e.g. Constant / AR(1) / EWMA) into a single frame.
|
||||
|
||||
Typical use::
|
||||
|
||||
from utils.predictions_cache import load_predictions, save_predictions
|
||||
|
||||
spec = {
|
||||
"data": {"source": "etfs", "start": START_DATE, "end": END_DATE,
|
||||
"max_symbols": MAX_SYMBOLS},
|
||||
"model": {"name": "rp_pca", "n_factors": N_FACTORS,
|
||||
"focus_gamma": focus_gamma},
|
||||
"forecasters": ["Constant", "AR(1)", "EWMA"],
|
||||
}
|
||||
cached = load_predictions(chapter=14, notebook_id="rp_pca", spec=spec)
|
||||
if cached is None:
|
||||
# ... expensive Stage 1-3 work, build long-form `frame` ...
|
||||
save_predictions(chapter=14, notebook_id="rp_pca", spec=spec, frame=frame)
|
||||
cached = frame
|
||||
|
||||
# downstream code consumes `cached` for plotting / IC summaries.
|
||||
|
||||
The spec dict is the contract: any value that materially changes the
|
||||
predictions must appear in it. Two runs that share a spec hash will share a
|
||||
cache entry, so changing a hyperparameter without updating the spec will
|
||||
silently reuse stale predictions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
|
||||
from utils.paths import get_chapter_dir
|
||||
|
||||
REQUIRED_COLUMNS = ("date", "symbol", "y_pred", "y_true")
|
||||
KEY_LENGTH = 12
|
||||
|
||||
|
||||
def predictions_cache_key(spec: dict) -> str:
|
||||
"""SHA-1 hash of the canonicalised spec; first 12 hex digits."""
|
||||
payload = json.dumps(spec, sort_keys=True, default=str).encode("utf-8")
|
||||
return hashlib.sha1(payload).hexdigest()[:KEY_LENGTH]
|
||||
|
||||
|
||||
def predictions_cache_path(chapter: int | str, notebook_id: str, spec: dict) -> Path:
|
||||
"""Deterministic parquet path; does not create the file or its parents."""
|
||||
root = get_chapter_dir(chapter) / "output" / "predictions" / notebook_id
|
||||
return root / f"{predictions_cache_key(spec)}.parquet"
|
||||
|
||||
|
||||
def load_predictions(chapter: int | str, notebook_id: str, spec: dict) -> pl.DataFrame | None:
|
||||
"""Return the cached predictions frame, or None if no cache entry exists."""
|
||||
path = predictions_cache_path(chapter, notebook_id, spec)
|
||||
if not path.exists():
|
||||
return None
|
||||
return pl.read_parquet(path)
|
||||
|
||||
|
||||
def save_predictions(chapter: int | str, notebook_id: str, spec: dict, frame: pl.DataFrame) -> Path:
|
||||
"""Write predictions frame to its content-addressed cache location."""
|
||||
missing = [c for c in REQUIRED_COLUMNS if c not in frame.columns]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"predictions frame missing required columns {missing}; have {list(frame.columns)}"
|
||||
)
|
||||
path = predictions_cache_path(chapter, notebook_id, spec)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
frame.write_parquet(path)
|
||||
return path
|
||||
@@ -0,0 +1,47 @@
|
||||
"""One-call seed initialization for reproducible notebook runs.
|
||||
|
||||
Notebooks that produce any random output should call ``set_global_seeds()``
|
||||
in their preamble, between imports and the first computation. Monte Carlo
|
||||
demos that *want* per-run variability should still call it, with the seed
|
||||
declared in their parameters cell so readers can change it explicitly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import random
|
||||
|
||||
|
||||
def set_global_seeds(seed: int = 42) -> None:
|
||||
"""Seed Python ``random``, NumPy, Torch (CPU+CUDA), and ``PYTHONHASHSEED``.
|
||||
|
||||
Polars and pandas operations that need a seed accept it per-call
|
||||
(e.g. ``df.sample(seed=seed)``) — there is no global polars seed to set.
|
||||
Scikit-learn estimators take ``random_state=`` per-instance.
|
||||
|
||||
Returns ``None`` rather than the seed so that a bare
|
||||
``set_global_seeds(SEED)`` in a notebook preamble does not render a
|
||||
spurious ``42`` execute_result. Notebooks that want to echo the seed back
|
||||
to the reader should ``print(SEED)`` explicitly.
|
||||
|
||||
``PYTHONHASHSEED`` is set on ``os.environ`` for the benefit of subprocesses
|
||||
spawned after the call; it has **no effect** on hash randomization in the
|
||||
currently running interpreter (that value is read once at startup). For
|
||||
end-to-end hash determinism the kernel must be launched with
|
||||
``PYTHONHASHSEED=<seed>`` already set in the environment.
|
||||
"""
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
random.seed(seed)
|
||||
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(seed)
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
pass
|
||||
else:
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
@@ -0,0 +1,883 @@
|
||||
"""Shared utilities for storage benchmarks.
|
||||
|
||||
This module provides common functionality used by both local and server benchmarks:
|
||||
- Synthetic data generation (OHLCV, tick data)
|
||||
- Timing infrastructure with warm-up and percentiles
|
||||
- Result validation with forced materialization
|
||||
- Configuration via YAML or environment variables
|
||||
- Memory estimation utilities
|
||||
"""
|
||||
|
||||
import gc
|
||||
import os
|
||||
import time
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import plotly.graph_objects as go
|
||||
import polars as pl
|
||||
import yaml
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
# =============================================================================
|
||||
# CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
# Benchmark scale configuration
|
||||
# BENCHMARK_SCALE: XS, S, M, L, XL, XXL
|
||||
BENCHMARK_SCALE = os.environ.get("BENCHMARK_SCALE", "").upper()
|
||||
|
||||
# Chapter directory paths
|
||||
from utils.paths import get_chapter_dir
|
||||
|
||||
CHAPTER_DIR = get_chapter_dir(2)
|
||||
CODE_DIR = CHAPTER_DIR
|
||||
|
||||
# Load configuration from YAML if available
|
||||
CONFIG_PATH = CODE_DIR / "benchmark_config.yaml"
|
||||
|
||||
|
||||
def load_config() -> dict:
|
||||
"""Load benchmark configuration from YAML file."""
|
||||
if CONFIG_PATH.exists():
|
||||
with open(CONFIG_PATH) as f:
|
||||
return yaml.safe_load(f)
|
||||
return {}
|
||||
|
||||
|
||||
CONFIG = load_config()
|
||||
|
||||
# New scale configurations (from YAML or defaults)
|
||||
# Format: {scale: {ohlcv: {symbols, rows_per_symbol, total_rows}, tick: {...}}}
|
||||
SCALE_CONFIGS_NEW = {}
|
||||
if CONFIG.get("scales"):
|
||||
for scale_name, scale_cfg in CONFIG["scales"].items():
|
||||
ohlcv = scale_cfg.get("ohlcv", {})
|
||||
tick = scale_cfg.get("tick", {})
|
||||
SCALE_CONFIGS_NEW[scale_name] = {
|
||||
"ohlcv": {
|
||||
"symbols": ohlcv.get("symbols", 10),
|
||||
"rows_per_symbol": ohlcv.get("rows_per_symbol", 1000),
|
||||
"total_rows": ohlcv.get("total_rows", 10000),
|
||||
},
|
||||
"tick": {
|
||||
"symbols": tick.get("symbols", 5),
|
||||
"trades": tick.get("trades", 5000),
|
||||
"quotes": tick.get("quotes", 25000),
|
||||
},
|
||||
"target_memory": scale_cfg.get("target_memory", "1MB"),
|
||||
"description": scale_cfg.get("description", ""),
|
||||
}
|
||||
else:
|
||||
# Default scale configs if YAML not present
|
||||
SCALE_CONFIGS_NEW = {
|
||||
"XS": {
|
||||
"ohlcv": {"symbols": 5, "rows_per_symbol": 200, "total_rows": 1000},
|
||||
"tick": {"symbols": 3, "trades": 500, "quotes": 2500},
|
||||
"target_memory": "100KB",
|
||||
"description": "Unit tests",
|
||||
},
|
||||
"S": {
|
||||
"ohlcv": {"symbols": 10, "rows_per_symbol": 1000, "total_rows": 10000},
|
||||
"tick": {"symbols": 5, "trades": 5000, "quotes": 25000},
|
||||
"target_memory": "1MB",
|
||||
"description": "Quick validation",
|
||||
},
|
||||
"M": {
|
||||
"ohlcv": {"symbols": 50, "rows_per_symbol": 2000, "total_rows": 100000},
|
||||
"tick": {"symbols": 10, "trades": 50000, "quotes": 250000},
|
||||
"target_memory": "10MB",
|
||||
"description": "Development",
|
||||
},
|
||||
"L": {
|
||||
"ohlcv": {"symbols": 100, "rows_per_symbol": 10000, "total_rows": 1000000},
|
||||
"tick": {"symbols": 50, "trades": 500000, "quotes": 2500000},
|
||||
"target_memory": "100MB",
|
||||
"description": "Standard benchmark",
|
||||
},
|
||||
"XL": {
|
||||
"ohlcv": {"symbols": 500, "rows_per_symbol": 20000, "total_rows": 10000000},
|
||||
"tick": {"symbols": 100, "trades": 5000000, "quotes": 25000000},
|
||||
"target_memory": "1GB",
|
||||
"description": "Scale testing",
|
||||
},
|
||||
"XXL": {
|
||||
"ohlcv": {"symbols": 1000, "rows_per_symbol": 100000, "total_rows": 100000000},
|
||||
"tick": {"symbols": 500, "trades": 50000000, "quotes": 250000000},
|
||||
"target_memory": "10GB",
|
||||
"description": "Production-scale",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_scale_config(scale: str) -> dict:
|
||||
"""Get configuration for a scale level.
|
||||
|
||||
Args:
|
||||
scale: Scale name (XS, S, M, L, XL, XXL)
|
||||
|
||||
Returns:
|
||||
Dict with ohlcv and tick configuration
|
||||
"""
|
||||
if scale not in SCALE_CONFIGS_NEW:
|
||||
raise ValueError(f"Unknown scale {scale!r}, expected one of {list(SCALE_CONFIGS_NEW)}")
|
||||
return SCALE_CONFIGS_NEW[scale]
|
||||
|
||||
|
||||
# Determine active scale
|
||||
# BENCHMARK_VERBOSE controls whether to print on import (default: False for clean imports)
|
||||
BENCHMARK_VERBOSE = os.environ.get("BENCHMARK_VERBOSE", "0") == "1"
|
||||
|
||||
if BENCHMARK_SCALE and BENCHMARK_SCALE in SCALE_CONFIGS_NEW:
|
||||
ACTIVE_SCALE = BENCHMARK_SCALE
|
||||
scale_cfg = get_scale_config(BENCHMARK_SCALE)
|
||||
N_SYMBOLS = scale_cfg["ohlcv"]["symbols"]
|
||||
N_ROWS_PER_SYMBOL = scale_cfg["ohlcv"]["rows_per_symbol"]
|
||||
N_TICKS_TRADES = scale_cfg["tick"]["trades"]
|
||||
N_TICKS_QUOTES = scale_cfg["tick"]["quotes"]
|
||||
TIMING_RUNS = CONFIG.get("execution", {}).get("iterations", {}).get(ACTIVE_SCALE, 3)
|
||||
else:
|
||||
# Default: S scale for quick iteration
|
||||
ACTIVE_SCALE = "S"
|
||||
scale_cfg = get_scale_config("S")
|
||||
N_SYMBOLS = scale_cfg["ohlcv"]["symbols"]
|
||||
N_ROWS_PER_SYMBOL = scale_cfg["ohlcv"]["rows_per_symbol"]
|
||||
N_TICKS_TRADES = scale_cfg["tick"]["trades"]
|
||||
N_TICKS_QUOTES = scale_cfg["tick"]["quotes"]
|
||||
TIMING_RUNS = 3
|
||||
|
||||
# Database connection configuration (environment variable overrides)
|
||||
DB_CONFIG = {
|
||||
"clickhouse": {
|
||||
"host": os.environ.get("CLICKHOUSE_HOST", "localhost"),
|
||||
"port": int(os.environ.get("CLICKHOUSE_PORT", "8123")),
|
||||
},
|
||||
"questdb": {
|
||||
"host": os.environ.get("QUESTDB_HOST", "localhost"),
|
||||
"http_port": int(os.environ.get("QUESTDB_HTTP_PORT", "9000")),
|
||||
"ilp_port": int(os.environ.get("QUESTDB_ILP_PORT", "9009")),
|
||||
"pg_port": int(os.environ.get("QUESTDB_PG_PORT", "8812")),
|
||||
},
|
||||
"timescaledb": {
|
||||
"host": os.environ.get("TIMESCALE_HOST", "localhost"),
|
||||
"port": int(os.environ.get("TIMESCALE_PORT", "5437")),
|
||||
"user": os.environ.get("TIMESCALE_USER", "postgres"),
|
||||
"password": os.environ.get("TIMESCALE_PASSWORD", "benchmark"),
|
||||
"database": os.environ.get("TIMESCALE_DB", "postgres"),
|
||||
},
|
||||
"influxdb": {
|
||||
"host": os.environ.get("INFLUXDB_HOST", "localhost"),
|
||||
"port": int(os.environ.get("INFLUXDB_PORT", "8086")),
|
||||
"org": os.environ.get("INFLUXDB_ORG", "benchmark"),
|
||||
"token": os.environ.get("INFLUXDB_TOKEN", "benchmark-token"),
|
||||
"bucket": os.environ.get("INFLUXDB_BUCKET", "benchmark"),
|
||||
},
|
||||
"postgres": {
|
||||
"host": os.environ.get("POSTGRES_HOST", "localhost"),
|
||||
"port": int(os.environ.get("POSTGRES_PORT", "5436")),
|
||||
"user": os.environ.get("POSTGRES_USER", "postgres"),
|
||||
"password": os.environ.get("POSTGRES_PASSWORD", "benchmark"),
|
||||
"database": os.environ.get("POSTGRES_DB", "ml4t"),
|
||||
},
|
||||
}
|
||||
|
||||
# WAL flush timeout (seconds) - adjustable for slower systems
|
||||
WAL_FLUSH_TIMEOUT = int(os.environ.get("WAL_FLUSH_TIMEOUT", "3"))
|
||||
|
||||
# =============================================================================
|
||||
# OUTPUT DIRECTORIES
|
||||
# =============================================================================
|
||||
# Directory structure:
|
||||
# .tmp/ - Transient data (gitignored), regenerated each run
|
||||
# ../output/benchmark/ - Results CSVs for book tables and citation
|
||||
# ../figures/ - Book figures (numbered: figure_3_N_slug.png)
|
||||
|
||||
# Transient benchmark data (synthetic OHLCV, trades, quotes)
|
||||
TMP_DIR = CHAPTER_DIR / ".tmp"
|
||||
TMP_DIR.mkdir(exist_ok=True)
|
||||
BENCHMARK_DIR = TMP_DIR # Alias used by benchmark notebooks
|
||||
|
||||
# Working charts (transient, not book figures)
|
||||
CHARTS_DIR = TMP_DIR / "charts"
|
||||
CHARTS_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# Results for book tables (CSV files for citation in prose)
|
||||
RESULTS_DIR = CHAPTER_DIR / "output" / "benchmark"
|
||||
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# =============================================================================
|
||||
# DATA CLASSES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
"""Container for benchmark results."""
|
||||
|
||||
name: str
|
||||
operation: str
|
||||
time_seconds: float
|
||||
size_bytes: int = 0
|
||||
rows: int = 0
|
||||
|
||||
@property
|
||||
def rows_per_second(self) -> float:
|
||||
return self.rows / self.time_seconds if self.time_seconds > 0 else 0
|
||||
|
||||
@property
|
||||
def mb_per_second(self) -> float:
|
||||
return (self.size_bytes / 1e6) / self.time_seconds if self.time_seconds > 0 else 0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TIMING & VALIDATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def drop_os_caches() -> bool:
|
||||
"""Drop OS page caches for accurate cold-cache benchmarking.
|
||||
|
||||
Requires sudo access. Returns True if successful, False otherwise.
|
||||
On Linux: sync; echo 3 > /proc/sys/vm/drop_caches
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
# Sync first to flush dirty pages
|
||||
subprocess.run(["sync"], check=True, timeout=30)
|
||||
# Drop caches (requires sudo or appropriate permissions)
|
||||
result = subprocess.run(
|
||||
["sudo", "-n", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"],
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def time_operation(func, n_runs: int = TIMING_RUNS, warmup: bool = True) -> tuple[float, Any]:
|
||||
"""Time a function with warm-up and percentile tracking.
|
||||
|
||||
Args:
|
||||
func: Function to time
|
||||
n_runs: Number of timing runs (default: TIMING_RUNS)
|
||||
warmup: Whether to run once before timing to warm up caches/JIT (default: True)
|
||||
|
||||
Returns:
|
||||
(mean_time, result): Mean execution time and last result
|
||||
Note: Timing stats (percentiles) are stored in result.timing_stats if available
|
||||
"""
|
||||
# Warm-up run (exclude from timing)
|
||||
if warmup:
|
||||
from contextlib import suppress
|
||||
|
||||
with suppress(Exception): # Warm-up failure is non-critical
|
||||
func()
|
||||
|
||||
# Timing runs
|
||||
times = []
|
||||
result = None
|
||||
for _ in range(n_runs):
|
||||
start = time.perf_counter()
|
||||
result = func()
|
||||
elapsed = time.perf_counter() - start
|
||||
times.append(elapsed)
|
||||
|
||||
# Calculate statistics
|
||||
times_array = np.array(times)
|
||||
mean_time = float(np.mean(times_array))
|
||||
|
||||
# Store timing stats as metadata (if result supports it)
|
||||
try:
|
||||
if hasattr(result, "__dict__"):
|
||||
result.timing_stats = {
|
||||
"mean": mean_time,
|
||||
"std": float(np.std(times_array)),
|
||||
"min": float(np.min(times_array)),
|
||||
"max": float(np.max(times_array)),
|
||||
"p50": float(np.percentile(times_array, 50)),
|
||||
"p95": float(np.percentile(times_array, 95)),
|
||||
"p99": float(np.percentile(times_array, 99)),
|
||||
}
|
||||
except Exception:
|
||||
pass # Not all result types support metadata
|
||||
|
||||
return mean_time, result
|
||||
|
||||
|
||||
def time_cold_cache(func, drop_caches: bool = True) -> tuple[float, Any]:
|
||||
"""Time a single cold-cache read operation.
|
||||
|
||||
For fair comparison of file formats, drops OS page caches before reading.
|
||||
This ensures memory-mapped formats (Feather) don't benefit from cached pages.
|
||||
|
||||
Args:
|
||||
func: Function to time (should be a file read operation)
|
||||
drop_caches: Whether to drop OS caches first (requires sudo)
|
||||
|
||||
Returns:
|
||||
(time, result): Execution time and result
|
||||
"""
|
||||
if drop_caches:
|
||||
cache_dropped = drop_os_caches()
|
||||
if not cache_dropped:
|
||||
print("Warning: Could not drop OS caches (requires sudo)")
|
||||
|
||||
gc.collect()
|
||||
start = time.perf_counter()
|
||||
result = func()
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
return elapsed, result
|
||||
|
||||
|
||||
def validate_result(
|
||||
result: Any, expected_rows: int, operation: str, tolerance: float = 0.1
|
||||
) -> None:
|
||||
"""Validate benchmark result has reasonable row count.
|
||||
|
||||
Args:
|
||||
result: Result to validate (DataFrame, list, dict with 'dataset', or None)
|
||||
expected_rows: Expected number of rows
|
||||
operation: Operation name for error messages
|
||||
tolerance: Fraction tolerance (0.1 = 10% deviation acceptable)
|
||||
|
||||
Raises:
|
||||
AssertionError: If row count is unreasonable
|
||||
"""
|
||||
if result is None:
|
||||
return # Skip validation for None results (optional databases)
|
||||
|
||||
# Extract row count based on result type
|
||||
if hasattr(result, "shape"): # DataFrame
|
||||
actual_rows = result.shape[0]
|
||||
elif isinstance(result, dict) and "dataset" in result: # QuestDB result
|
||||
actual_rows = len(result["dataset"])
|
||||
elif isinstance(result, list):
|
||||
actual_rows = len(result)
|
||||
else:
|
||||
return # Unknown type, skip validation
|
||||
|
||||
# Check row count is within tolerance
|
||||
min_rows = int(expected_rows * (1 - tolerance))
|
||||
max_rows = int(expected_rows * (1 + tolerance))
|
||||
|
||||
if not (min_rows <= actual_rows <= max_rows):
|
||||
raise AssertionError(
|
||||
f"{operation}: Expected {expected_rows:,} rows (±{tolerance:.0%}), got {actual_rows:,}"
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MATERIALIZATION HELPERS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def force_materialize_polars(df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame:
|
||||
"""Force full materialization of a Polars DataFrame by scanning ALL columns.
|
||||
|
||||
Memory-mapped formats (like Feather/Arrow IPC) may return handles without
|
||||
loading data. This function forces actual data access by touching every column.
|
||||
|
||||
Args:
|
||||
df: Polars DataFrame or LazyFrame
|
||||
|
||||
Returns:
|
||||
Materialized DataFrame with data actually loaded into memory
|
||||
"""
|
||||
if isinstance(df, pl.LazyFrame):
|
||||
df = df.collect()
|
||||
|
||||
# Force materialization by touching EVERY column (not just first 3 numeric)
|
||||
# This ensures all data is actually loaded, not just memory-mapped
|
||||
exprs = []
|
||||
for col_name in df.columns:
|
||||
dtype = df[col_name].dtype
|
||||
if dtype in (pl.Float64, pl.Float32, pl.Int64, pl.Int32, pl.Int16, pl.Int8):
|
||||
# Numeric: compute sum
|
||||
exprs.append(pl.col(col_name).sum().alias(f"{col_name}_sum"))
|
||||
elif dtype == pl.String:
|
||||
# String: compute length sum (forces read of all string data)
|
||||
exprs.append(pl.col(col_name).str.len_bytes().sum().alias(f"{col_name}_len"))
|
||||
elif dtype in (pl.Datetime, pl.Date):
|
||||
# Datetime: compute min/max (forces read)
|
||||
exprs.append(pl.col(col_name).min().alias(f"{col_name}_min"))
|
||||
else:
|
||||
# Other types: count non-null (forces read)
|
||||
exprs.append(pl.col(col_name).count().alias(f"{col_name}_count"))
|
||||
|
||||
if exprs:
|
||||
_ = df.select(exprs).to_dict()
|
||||
|
||||
return df
|
||||
|
||||
|
||||
def force_materialize_pandas(df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Force full materialization of a pandas DataFrame by scanning ALL columns.
|
||||
|
||||
Args:
|
||||
df: pandas DataFrame
|
||||
|
||||
Returns:
|
||||
Materialized DataFrame with all data accessed
|
||||
"""
|
||||
# Force read of ALL columns, not just first 3 numeric
|
||||
result = {}
|
||||
|
||||
# Numeric columns: compute sum
|
||||
numeric_cols = df.select_dtypes(include=[np.number]).columns
|
||||
if len(numeric_cols) > 0:
|
||||
result["numeric_sums"] = df[numeric_cols].sum().to_dict()
|
||||
|
||||
# String/object columns: compute total string length
|
||||
object_cols = df.select_dtypes(include=["object", "string"]).columns
|
||||
if len(object_cols) > 0:
|
||||
result["string_lens"] = {col: df[col].astype(str).str.len().sum() for col in object_cols}
|
||||
|
||||
# Datetime columns: compute min
|
||||
datetime_cols = df.select_dtypes(include=["datetime64"]).columns
|
||||
if len(datetime_cols) > 0:
|
||||
result["datetime_mins"] = {col: df[col].min() for col in datetime_cols}
|
||||
|
||||
# Force evaluation
|
||||
_ = result
|
||||
|
||||
return df
|
||||
|
||||
|
||||
def read_with_materialization(
|
||||
read_func,
|
||||
path: Path,
|
||||
library: str = "polars",
|
||||
) -> tuple[float, Any]:
|
||||
"""Time a read operation with forced materialization.
|
||||
|
||||
Args:
|
||||
read_func: Function to read data (e.g., pl.read_parquet, pd.read_csv)
|
||||
path: Path to file
|
||||
library: "polars" or "pandas"
|
||||
|
||||
Returns:
|
||||
(time_seconds, result): Tuple of read time and DataFrame
|
||||
"""
|
||||
gc.collect()
|
||||
|
||||
start = time.perf_counter()
|
||||
df = read_func(path)
|
||||
|
||||
# Force materialization
|
||||
if library == "polars":
|
||||
df = force_materialize_polars(df)
|
||||
else:
|
||||
df = force_materialize_pandas(df)
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
return elapsed, df
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MEMORY UTILITIES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def estimate_memory_mb(df: pl.DataFrame | pd.DataFrame) -> float:
|
||||
"""Estimate memory usage of a DataFrame in MB.
|
||||
|
||||
Args:
|
||||
df: Polars or pandas DataFrame
|
||||
|
||||
Returns:
|
||||
Estimated memory in megabytes
|
||||
"""
|
||||
if isinstance(df, pl.DataFrame):
|
||||
return df.estimated_size("mb")
|
||||
else:
|
||||
return df.memory_usage(deep=True).sum() / 1_000_000
|
||||
|
||||
|
||||
def get_memory_usage_mb() -> float:
|
||||
"""Get current process memory usage in MB.
|
||||
|
||||
Returns:
|
||||
Memory usage in megabytes
|
||||
"""
|
||||
import psutil
|
||||
|
||||
process = psutil.Process()
|
||||
return process.memory_info().rss / 1_000_000
|
||||
|
||||
|
||||
def run_with_gc(func):
|
||||
"""Run function with garbage collection before and after.
|
||||
|
||||
Args:
|
||||
func: Function to run
|
||||
|
||||
Returns:
|
||||
Function result
|
||||
"""
|
||||
gc.collect()
|
||||
result = func()
|
||||
gc.collect()
|
||||
return result
|
||||
|
||||
|
||||
def save_chart(fig: go.Figure, name: str) -> None:
|
||||
"""Save chart to HTML file instead of opening browser."""
|
||||
path = CHARTS_DIR / f"{name}.html"
|
||||
fig.write_html(str(path), include_plotlyjs="cdn")
|
||||
print(f"Chart saved: {path}")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# DATA GENERATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def generate_ohlcv_data(
|
||||
n_symbols: int = N_SYMBOLS,
|
||||
n_rows: int = N_ROWS_PER_SYMBOL,
|
||||
seed: int = 42,
|
||||
) -> pl.DataFrame:
|
||||
"""Generate realistic synthetic OHLCV panel data (fully vectorized).
|
||||
|
||||
Creates a panel dataset with multiple symbols and time-series bars.
|
||||
OHLCV constraints are enforced: H >= max(O,C), L <= min(O,C).
|
||||
|
||||
This implementation is fully vectorized using numpy arrays, enabling
|
||||
generation of 10M+ rows in seconds (required for XL/XXL scales).
|
||||
|
||||
Args:
|
||||
n_symbols: Number of unique symbols
|
||||
n_rows: Number of rows per symbol
|
||||
seed: Random seed for reproducibility
|
||||
|
||||
Returns:
|
||||
Polars DataFrame with columns: timestamp, symbol, open, high, low, close, volume, vwap, num_trades
|
||||
"""
|
||||
np.random.seed(seed)
|
||||
total_rows = n_symbols * n_rows
|
||||
|
||||
# Generate all timestamps at once (vectorized)
|
||||
base_time = np.datetime64("2024-01-01T00:00:00", "us")
|
||||
minute_offsets = np.arange(n_rows, dtype="timedelta64[m]")
|
||||
single_symbol_times = base_time + minute_offsets
|
||||
|
||||
# Tile timestamps for all symbols
|
||||
timestamps = np.tile(single_symbol_times, n_symbols)
|
||||
|
||||
# Generate symbol array (repeat each symbol n_rows times)
|
||||
symbol_names = np.array([f"SYM_{i:03d}" for i in range(n_symbols)])
|
||||
symbols = np.repeat(symbol_names, n_rows)
|
||||
|
||||
# Generate base prices per symbol, then broadcast to all rows
|
||||
base_prices = 100 + np.random.randn(n_symbols) * 10 # (n_symbols,)
|
||||
|
||||
# Generate returns for all rows at once: (n_symbols, n_rows)
|
||||
returns = np.random.randn(n_symbols, n_rows) * 0.001
|
||||
|
||||
# Cumulative sum along rows axis, then flatten
|
||||
cumret = np.cumsum(returns, axis=1) # (n_symbols, n_rows)
|
||||
prices = (base_prices[:, np.newaxis] * np.exp(cumret)).flatten() # (total_rows,)
|
||||
|
||||
# Generate intrabar noise for all rows
|
||||
noise = np.abs(np.random.randn(total_rows)) * 0.002
|
||||
|
||||
# Generate OHLC
|
||||
opens = prices * (1 - noise * 0.5)
|
||||
closes = prices * (1 + noise * 0.5)
|
||||
highs_raw = prices * (1 + noise)
|
||||
lows_raw = prices * (1 - noise)
|
||||
|
||||
# Enforce OHLC constraints: H >= max(O,C), L <= min(O,C)
|
||||
highs = np.maximum(np.maximum(opens, closes), highs_raw)
|
||||
lows = np.minimum(np.minimum(opens, closes), lows_raw)
|
||||
|
||||
# Generate volume (lognormal distribution)
|
||||
volumes = np.exp(np.random.randn(total_rows) * 0.5 + 10).astype(np.int64)
|
||||
|
||||
# VWAP approximation: typical price (H+L+C)/3
|
||||
vwap = (highs + lows + closes) / 3
|
||||
|
||||
# Number of trades (proportional to volume with noise)
|
||||
num_trades = (volumes / 1000 + np.random.randint(10, 100, total_rows)).astype(np.int32)
|
||||
|
||||
# Build DataFrame directly from numpy arrays (zero-copy where possible)
|
||||
df = pl.DataFrame(
|
||||
{
|
||||
"timestamp": timestamps,
|
||||
"symbol": symbols,
|
||||
"open": opens,
|
||||
"high": highs,
|
||||
"low": lows,
|
||||
"close": closes,
|
||||
"volume": volumes,
|
||||
"vwap": vwap,
|
||||
"num_trades": num_trades,
|
||||
},
|
||||
schema={
|
||||
"timestamp": pl.Datetime("us"),
|
||||
"symbol": pl.String,
|
||||
"open": pl.Float64,
|
||||
"high": pl.Float64,
|
||||
"low": pl.Float64,
|
||||
"close": pl.Float64,
|
||||
"volume": pl.Int64,
|
||||
"vwap": pl.Float64,
|
||||
"num_trades": pl.Int32,
|
||||
},
|
||||
)
|
||||
|
||||
return df
|
||||
|
||||
|
||||
def generate_tick_data(
|
||||
n_trades: int = N_TICKS_TRADES,
|
||||
n_quotes: int = N_TICKS_QUOTES,
|
||||
n_symbols: int = N_SYMBOLS,
|
||||
seed: int = 42,
|
||||
) -> tuple[pl.DataFrame, pl.DataFrame]:
|
||||
"""Generate realistic synthetic tick data (fully vectorized).
|
||||
|
||||
Trades are correlated with quotes (occur near bid/ask) for realistic simulation.
|
||||
|
||||
This implementation is fully vectorized using numpy arrays, enabling
|
||||
generation of millions of ticks in seconds (required for XL/XXL scales).
|
||||
|
||||
Args:
|
||||
n_trades: Total number of trade ticks
|
||||
n_quotes: Total number of quote ticks
|
||||
n_symbols: Number of unique symbols
|
||||
seed: Random seed for reproducibility
|
||||
|
||||
Returns:
|
||||
(trades_df, quotes_df): Tuple of Polars DataFrames
|
||||
"""
|
||||
np.random.seed(seed)
|
||||
|
||||
symbol_names = np.array([f"SYM_{i:03d}" for i in range(n_symbols)])
|
||||
base_time = np.datetime64("2024-01-01T09:30:00", "us")
|
||||
|
||||
# =========================================================================
|
||||
# QUOTES: Fully vectorized generation
|
||||
# =========================================================================
|
||||
|
||||
# Assign quotes to symbols (weighted: some symbols more active)
|
||||
# Use exponential weights so first symbols get more quotes
|
||||
weights = np.exp(-np.arange(n_symbols) * 0.1)
|
||||
weights = weights / weights.sum()
|
||||
quote_symbol_indices = np.random.choice(n_symbols, size=n_quotes, p=weights)
|
||||
quote_symbols = symbol_names[quote_symbol_indices]
|
||||
|
||||
# Generate timestamps (100 microseconds apart)
|
||||
quote_timestamps = base_time + (np.arange(n_quotes) * 100).astype("timedelta64[us]")
|
||||
|
||||
# Generate mid prices per symbol using random walk
|
||||
# Strategy: for each symbol, generate cumulative random walk, then gather
|
||||
quotes_per_symbol = np.bincount(quote_symbol_indices, minlength=n_symbols)
|
||||
max_quotes_per_symbol = quotes_per_symbol.max()
|
||||
|
||||
# Generate random walk increments for all symbols at once
|
||||
base_mids = 150 + np.random.randn(n_symbols) * 20 # Starting mid price per symbol
|
||||
walk_increments = np.random.randn(n_symbols, max_quotes_per_symbol) * 0.01
|
||||
walk_paths = base_mids[:, np.newaxis] + np.cumsum(walk_increments, axis=1)
|
||||
|
||||
# For each quote, look up the appropriate price from that symbol's walk
|
||||
# Track position within each symbol's sequence
|
||||
symbol_counters = np.zeros(n_symbols, dtype=np.int64)
|
||||
mid_prices = np.empty(n_quotes, dtype=np.float64)
|
||||
|
||||
# Vectorized lookup using cumulative counts
|
||||
for sym_idx in range(n_symbols):
|
||||
mask = quote_symbol_indices == sym_idx
|
||||
count = mask.sum()
|
||||
if count > 0:
|
||||
mid_prices[mask] = walk_paths[sym_idx, :count]
|
||||
|
||||
# Generate spreads (0.01-0.05% of mid)
|
||||
spread_pct = 0.0001 + np.abs(np.random.randn(n_quotes)) * 0.0002
|
||||
spreads = mid_prices * spread_pct
|
||||
bids = mid_prices - spreads / 2
|
||||
asks = mid_prices + spreads / 2
|
||||
|
||||
# Generate sizes (lognormal)
|
||||
bid_sizes = np.exp(np.random.randn(n_quotes) * 0.3 + 5).astype(np.int64)
|
||||
ask_sizes = np.exp(np.random.randn(n_quotes) * 0.3 + 5).astype(np.int64)
|
||||
|
||||
quotes_df = pl.DataFrame(
|
||||
{
|
||||
"timestamp": quote_timestamps,
|
||||
"symbol": quote_symbols,
|
||||
"bid": bids,
|
||||
"ask": asks,
|
||||
"bid_size": bid_sizes,
|
||||
"ask_size": ask_sizes,
|
||||
},
|
||||
schema={
|
||||
"timestamp": pl.Datetime("us"),
|
||||
"symbol": pl.String,
|
||||
"bid": pl.Float64,
|
||||
"ask": pl.Float64,
|
||||
"bid_size": pl.Int64,
|
||||
"ask_size": pl.Int64,
|
||||
},
|
||||
)
|
||||
|
||||
# =========================================================================
|
||||
# TRADES: Fully vectorized generation
|
||||
# =========================================================================
|
||||
|
||||
# Assign trades to symbols (same distribution as quotes)
|
||||
trade_symbol_indices = np.random.choice(n_symbols, size=n_trades, p=weights)
|
||||
trade_symbols = symbol_names[trade_symbol_indices]
|
||||
|
||||
# Generate timestamps (100 microseconds apart, offset by 50us from quotes)
|
||||
trade_timestamps = base_time + (np.arange(n_trades) * 100 + 50).astype("timedelta64[us]")
|
||||
|
||||
# Generate sizes
|
||||
trade_sizes = np.exp(np.random.randn(n_trades) * 0.3 + 5).astype(np.int64)
|
||||
|
||||
# Build preliminary trades DataFrame
|
||||
trades_prelim = pl.DataFrame(
|
||||
{
|
||||
"timestamp": trade_timestamps,
|
||||
"symbol": trade_symbols,
|
||||
"price": np.zeros(n_trades), # Placeholder
|
||||
"size": trade_sizes,
|
||||
},
|
||||
schema={
|
||||
"timestamp": pl.Datetime("us"),
|
||||
"symbol": pl.String,
|
||||
"price": pl.Float64,
|
||||
"size": pl.Int64,
|
||||
},
|
||||
)
|
||||
|
||||
# Use ASOF join to match each trade with most recent quote
|
||||
trades_with_quotes = trades_prelim.join_asof(
|
||||
quotes_df.sort(["symbol", "timestamp"]),
|
||||
on="timestamp",
|
||||
by="symbol",
|
||||
strategy="backward",
|
||||
)
|
||||
|
||||
# Calculate trade prices: 25% at ask, 25% at bid, 50% at mid
|
||||
trade_sides = np.random.rand(n_trades)
|
||||
bid_arr = trades_with_quotes["bid"].to_numpy()
|
||||
ask_arr = trades_with_quotes["ask"].to_numpy()
|
||||
|
||||
trade_prices = np.where(
|
||||
trade_sides < 0.25,
|
||||
ask_arr, # Buy at ask (25%)
|
||||
np.where(
|
||||
trade_sides < 0.5,
|
||||
bid_arr, # Sell at bid (25%)
|
||||
(bid_arr + ask_arr) / 2, # Mid (50%)
|
||||
),
|
||||
)
|
||||
|
||||
# Handle NaN prices (trades before first quote)
|
||||
trade_prices = np.where(np.isnan(trade_prices), 150.0, trade_prices)
|
||||
|
||||
# Create final trades DataFrame
|
||||
trades_df = pl.DataFrame(
|
||||
{
|
||||
"timestamp": trades_with_quotes["timestamp"],
|
||||
"symbol": trades_with_quotes["symbol"],
|
||||
"price": trade_prices,
|
||||
"size": trades_with_quotes["size"],
|
||||
}
|
||||
)
|
||||
|
||||
return trades_df, quotes_df
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# RESULTS STORAGE
|
||||
# =============================================================================
|
||||
|
||||
# Map benchmark_type to clean filename prefixes
|
||||
BENCHMARK_TYPE_TO_FILENAME = {
|
||||
"formats": "file_formats",
|
||||
"embedded": "embedded_dbs",
|
||||
"pandas_polars": "pandas_polars",
|
||||
"servers": "server_dbs",
|
||||
# Legacy mappings
|
||||
"local": "local",
|
||||
"server": "server",
|
||||
}
|
||||
|
||||
|
||||
def save_benchmark_results(
|
||||
results: list[BenchmarkResult], benchmark_type: str, scale: str | None = None
|
||||
) -> Path:
|
||||
"""Save benchmark results to CSV for book publication.
|
||||
|
||||
Results are saved to output/benchmark/ with clean filenames for citation in prose.
|
||||
Example: file_formats_l.csv, embedded_dbs_xl.csv
|
||||
|
||||
Args:
|
||||
results: List of BenchmarkResult objects
|
||||
benchmark_type: "formats", "embedded", "pandas_polars", "servers"
|
||||
scale: Scale level (S, L, XL, etc.) or None for auto-detect
|
||||
|
||||
Returns:
|
||||
Path to saved CSV file
|
||||
"""
|
||||
if scale is None:
|
||||
scale = ACTIVE_SCALE
|
||||
|
||||
# Clean filename prefix
|
||||
filename_prefix = BENCHMARK_TYPE_TO_FILENAME.get(benchmark_type, benchmark_type)
|
||||
|
||||
# Create results DataFrame
|
||||
df = pl.DataFrame(
|
||||
[
|
||||
{
|
||||
"benchmark_type": benchmark_type,
|
||||
"scale": scale,
|
||||
"technology": r.name,
|
||||
"operation": r.operation,
|
||||
"time_seconds": r.time_seconds,
|
||||
"size_mb": r.size_bytes / 1_000_000 if r.size_bytes else None,
|
||||
"rows": r.rows if r.rows else None,
|
||||
"rows_per_second": r.rows_per_second if r.rows else None,
|
||||
"mb_per_second": r.mb_per_second if r.size_bytes else None,
|
||||
"timestamp": pd.Timestamp.now().isoformat(),
|
||||
"n_symbols": N_SYMBOLS,
|
||||
"n_rows_per_symbol": N_ROWS_PER_SYMBOL,
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
)
|
||||
|
||||
# Save as CSV (for book tables and prose citation)
|
||||
csv_path = RESULTS_DIR / f"{filename_prefix}_{scale.lower()}.csv"
|
||||
df.write_csv(csv_path)
|
||||
|
||||
print(f"\n📁 Results saved to: {csv_path}")
|
||||
return csv_path
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PRINT CONFIGURATION (only when BENCHMARK_VERBOSE=1 or running as main)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def print_config() -> None:
|
||||
"""Print current benchmark configuration."""
|
||||
print(f"📊 Scale: {ACTIVE_SCALE} ({scale_cfg['target_memory']} target)")
|
||||
print(f" OHLCV: {N_SYMBOLS} symbols × {N_ROWS_PER_SYMBOL:,} rows/symbol")
|
||||
print(f" Ticks: {N_TICKS_TRADES:,} trades, {N_TICKS_QUOTES:,} quotes")
|
||||
print(f" Timing runs: {TIMING_RUNS}")
|
||||
|
||||
|
||||
# Only print on import if explicitly requested
|
||||
if BENCHMARK_VERBOSE:
|
||||
print_config()
|
||||
+794
@@ -0,0 +1,794 @@
|
||||
"""ML4T Visualization Style.
|
||||
|
||||
Canonical color palette, matplotlib rcParams, Plotly template, and chart
|
||||
helpers for all book visualizations.
|
||||
|
||||
## Automatic Styling (Matplotlib)
|
||||
|
||||
ML4T style is applied automatically when running from repo root.
|
||||
The ``matplotlibrc`` file in the repo root is loaded by matplotlib
|
||||
before any other config. No imports or function calls needed.
|
||||
|
||||
## Explicit Color References
|
||||
|
||||
from utils.style import COLORS
|
||||
ax.plot(x, y, color=COLORS['blue'])
|
||||
ax.axhline(0, color=COLORS['amber'], linestyle='--')
|
||||
|
||||
## Plotly
|
||||
|
||||
import plotly.io as pio
|
||||
pio.templates.default = "ml4t" # Auto-registered on import
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from matplotlib.axes import Axes
|
||||
|
||||
# =============================================================================
|
||||
# ML4T COLOR PALETTE
|
||||
# =============================================================================
|
||||
# Aligned with ml4t.io website identity
|
||||
|
||||
COLORS = {
|
||||
# Primary blues (core identity)
|
||||
"blue": "#0a1628", # Deep blue - primary emphasis, main data
|
||||
"blue_light": "#152238", # Lighter blue - secondary elements
|
||||
"slate": "#1a2d4a", # Mid-blue - tertiary, gridlines
|
||||
# Silver tones (backgrounds, text)
|
||||
"silver": "#F8F8F6", # Light silver - text on dark, highlights
|
||||
"silver_muted": "#e8e8e6", # Muted silver - borders, subtle elements
|
||||
# Warm accents (highlights, emphasis)
|
||||
"amber": "#D4A84B", # Warm amber - CTAs, important highlights
|
||||
"amber_light": "#E4B85B", # Lighter amber - hover states
|
||||
"copper": "#C87533", # Copper - secondary accent
|
||||
# Semantic (for data meaning)
|
||||
"positive": "#10b981", # Success green - profits, gains
|
||||
"negative": "#ef4444", # Error red - losses (use sparingly!)
|
||||
"neutral": "#334155", # Slate gray - neutral elements
|
||||
# Backgrounds
|
||||
"bg_light": "#FAFAF9", # Warm off-white (light mode)
|
||||
"bg_dark": "#0a1628", # Deep blue (dark mode)
|
||||
}
|
||||
|
||||
# Grayscale equivalents for print
|
||||
GRAYSCALE = {
|
||||
"blue": 0.10, # ~10% gray (very dark)
|
||||
"slate": 0.25, # ~25% gray
|
||||
"amber": 0.65, # ~65% gray
|
||||
"silver": 0.97, # ~97% gray (nearly white)
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# MATPLOTLIB STYLE CONFIGURATIONS
|
||||
# =============================================================================
|
||||
|
||||
_BASE_STYLE = {
|
||||
# Figure
|
||||
"figure.dpi": 100,
|
||||
"figure.figsize": (10, 6),
|
||||
"savefig.dpi": 150,
|
||||
"savefig.bbox": "tight",
|
||||
"savefig.pad_inches": 0.1,
|
||||
# Axes
|
||||
"axes.spines.top": False,
|
||||
"axes.spines.right": False,
|
||||
"axes.titlesize": 14,
|
||||
"axes.titleweight": "semibold",
|
||||
"axes.titlepad": 12,
|
||||
"axes.labelsize": 11,
|
||||
"axes.labelpad": 8,
|
||||
# Grid
|
||||
"axes.grid": True,
|
||||
"grid.alpha": 0.4,
|
||||
"grid.linewidth": 0.5,
|
||||
# Ticks
|
||||
"xtick.labelsize": 10,
|
||||
"ytick.labelsize": 10,
|
||||
"xtick.major.pad": 4,
|
||||
"ytick.major.pad": 4,
|
||||
# Lines
|
||||
"lines.linewidth": 2,
|
||||
"lines.markersize": 6,
|
||||
# Legend
|
||||
"legend.frameon": False,
|
||||
"legend.fontsize": 10,
|
||||
# Font (prefer DM Sans, fallback to system sans)
|
||||
"font.family": ["sans-serif"],
|
||||
"font.sans-serif": ["DM Sans", "DejaVu Sans", "Helvetica", "Arial"],
|
||||
"font.size": 10,
|
||||
}
|
||||
|
||||
ML4T_LIGHT_STYLE = {
|
||||
**_BASE_STYLE,
|
||||
"figure.facecolor": COLORS["bg_light"],
|
||||
"axes.facecolor": "white",
|
||||
"axes.edgecolor": COLORS["silver_muted"],
|
||||
"axes.labelcolor": COLORS["neutral"],
|
||||
"axes.titlecolor": COLORS["blue"],
|
||||
"xtick.color": COLORS["neutral"],
|
||||
"ytick.color": COLORS["neutral"],
|
||||
"grid.color": COLORS["silver_muted"],
|
||||
"text.color": COLORS["neutral"],
|
||||
}
|
||||
|
||||
ML4T_DARK_STYLE = {
|
||||
**_BASE_STYLE,
|
||||
"figure.facecolor": COLORS["bg_dark"],
|
||||
"axes.facecolor": COLORS["blue_light"],
|
||||
"axes.edgecolor": COLORS["slate"],
|
||||
"axes.labelcolor": COLORS["silver"],
|
||||
"axes.titlecolor": COLORS["silver"],
|
||||
"xtick.color": COLORS["silver_muted"],
|
||||
"ytick.color": COLORS["silver_muted"],
|
||||
"grid.color": COLORS["slate"],
|
||||
"text.color": COLORS["silver"],
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# STYLE APPLICATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def apply_ml4t_style(mode: Literal["light", "dark"] = "light") -> None:
|
||||
"""Apply ML4T style to both Matplotlib and Plotly.
|
||||
|
||||
Args:
|
||||
mode: 'light' (default) for white backgrounds, 'dark' for blue backgrounds
|
||||
"""
|
||||
if mode == "light":
|
||||
plt.rcParams.update(ML4T_LIGHT_STYLE)
|
||||
else:
|
||||
plt.rcParams.update(ML4T_DARK_STYLE)
|
||||
|
||||
# Apply Plotly template if available
|
||||
import contextlib
|
||||
|
||||
with contextlib.suppress(ImportError):
|
||||
_register_plotly_template()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PALETTE HELPERS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def ml4t_palette(n: int = 5, categorical: bool = False) -> list[str]:
|
||||
"""Return colors from the ML4T palette.
|
||||
|
||||
Args:
|
||||
n: Number of colors to return (max 5)
|
||||
categorical: If True, returns distinct colors for categories.
|
||||
If False, returns blue gradient for sequential data.
|
||||
|
||||
Returns:
|
||||
List of hex color strings
|
||||
"""
|
||||
if categorical:
|
||||
colors = [
|
||||
COLORS["blue"],
|
||||
COLORS["amber"],
|
||||
COLORS["slate"],
|
||||
COLORS["copper"],
|
||||
COLORS["silver_muted"],
|
||||
]
|
||||
else:
|
||||
colors = [
|
||||
COLORS["blue"],
|
||||
COLORS["slate"],
|
||||
COLORS["blue_light"],
|
||||
COLORS["silver_muted"],
|
||||
COLORS["silver"],
|
||||
]
|
||||
return colors[:n]
|
||||
|
||||
|
||||
def ml4t_diverging() -> list[str]:
|
||||
"""Return diverging palette (negative to positive).
|
||||
|
||||
Use for data with meaningful zero point (e.g., returns, correlations).
|
||||
|
||||
Returns:
|
||||
List of 3 colors: [negative, neutral, positive]
|
||||
"""
|
||||
return [COLORS["negative"], COLORS["silver_muted"], COLORS["positive"]]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CHART HELPERS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def annotate_peak(ax: Axes, x: object, y: object, label: str, offset: tuple = (10, 10)) -> None:
|
||||
"""Annotate a peak/trough with ML4T styling.
|
||||
|
||||
Args:
|
||||
ax: matplotlib axes
|
||||
x, y: Coordinates of the point
|
||||
label: Text label
|
||||
offset: (x, y) offset in points
|
||||
"""
|
||||
ax.annotate(
|
||||
label,
|
||||
xy=(x, y),
|
||||
xytext=offset,
|
||||
textcoords="offset points",
|
||||
fontsize=9,
|
||||
color=COLORS["neutral"],
|
||||
arrowprops={
|
||||
"arrowstyle": "->",
|
||||
"color": COLORS["amber"],
|
||||
"connectionstyle": "arc3,rad=0.2",
|
||||
},
|
||||
bbox={
|
||||
"boxstyle": "round,pad=0.3",
|
||||
"facecolor": COLORS["silver"],
|
||||
"edgecolor": COLORS["silver_muted"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def add_regime_shading(ax: Axes, periods: list[tuple], label: str = "Crisis") -> None:
|
||||
"""Add regime shading to a time series plot.
|
||||
|
||||
Args:
|
||||
ax: matplotlib axes
|
||||
periods: List of (start, end) tuples defining regime periods
|
||||
label: Label for legend
|
||||
"""
|
||||
for i, (start, end) in enumerate(periods):
|
||||
ax.axvspan(
|
||||
start,
|
||||
end,
|
||||
alpha=0.15,
|
||||
color=COLORS["amber"],
|
||||
label=label if i == 0 else None,
|
||||
)
|
||||
|
||||
|
||||
def format_pct_axis(ax: Axes, axis: Literal["x", "y", "both"] = "y") -> None:
|
||||
"""Format axis as percentage with ML4T styling.
|
||||
|
||||
Args:
|
||||
ax: matplotlib axes
|
||||
axis: Which axis to format ('x', 'y', or 'both')
|
||||
"""
|
||||
from matplotlib.ticker import PercentFormatter
|
||||
|
||||
formatter = PercentFormatter(xmax=1, decimals=0)
|
||||
if axis in ("y", "both"):
|
||||
ax.yaxis.set_major_formatter(formatter)
|
||||
if axis in ("x", "both"):
|
||||
ax.xaxis.set_major_formatter(formatter)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PLOTLY TEMPLATE (optional — only used if Plotly is installed)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _register_plotly_template() -> None:
|
||||
"""Register the ML4T template with Plotly."""
|
||||
import plotly.graph_objects as go
|
||||
import plotly.io as pio
|
||||
|
||||
template = go.layout.Template(
|
||||
layout=go.Layout(
|
||||
font=dict(
|
||||
family="DM Sans, DejaVu Sans, sans-serif",
|
||||
size=11,
|
||||
color=COLORS["neutral"],
|
||||
),
|
||||
paper_bgcolor=COLORS["bg_light"],
|
||||
plot_bgcolor="white",
|
||||
title=dict(
|
||||
font=dict(size=14, color=COLORS["blue"]),
|
||||
x=0.5,
|
||||
xanchor="center",
|
||||
),
|
||||
xaxis=dict(
|
||||
gridcolor=COLORS["silver_muted"],
|
||||
linecolor=COLORS["silver_muted"],
|
||||
tickfont=dict(size=10),
|
||||
title=dict(font=dict(size=11)),
|
||||
showgrid=True,
|
||||
gridwidth=0.5,
|
||||
),
|
||||
yaxis=dict(
|
||||
gridcolor=COLORS["silver_muted"],
|
||||
linecolor=COLORS["silver_muted"],
|
||||
tickfont=dict(size=10),
|
||||
title=dict(font=dict(size=11)),
|
||||
showgrid=True,
|
||||
gridwidth=0.5,
|
||||
),
|
||||
colorway=[
|
||||
COLORS["blue"],
|
||||
COLORS["amber"],
|
||||
COLORS["slate"],
|
||||
COLORS["copper"],
|
||||
COLORS["positive"],
|
||||
COLORS["negative"],
|
||||
],
|
||||
legend=dict(
|
||||
bgcolor="rgba(255,255,255,0.8)",
|
||||
bordercolor=COLORS["silver_muted"],
|
||||
borderwidth=1,
|
||||
font=dict(size=10),
|
||||
),
|
||||
hoverlabel=dict(
|
||||
bgcolor="white",
|
||||
font_size=11,
|
||||
font_family="DM Sans, DejaVu Sans, sans-serif",
|
||||
),
|
||||
)
|
||||
)
|
||||
pio.templates["ml4t"] = template
|
||||
|
||||
|
||||
# Auto-register Plotly template on import
|
||||
_register_plotly_template()
|
||||
HAS_PLOTLY = True
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PUBLICATION (BOOK) STYLE — MIT Press, dual-track (grayscale print + color web)
|
||||
# =============================================================================
|
||||
# Used by `~/ml4t/book/<ch>/figures/scripts/generate_figure_*.py`.
|
||||
# Notebooks may also call `apply_book_style()` to render the same look.
|
||||
#
|
||||
# Two tracks, same data, same script:
|
||||
# - "print": grayscale-first, semantic fills, varied linestyles. Top-level PNG.
|
||||
# - "color": ML4T palette overlay. `color/` subdir.
|
||||
# The grayscale track is the source of truth: data must be legible without color.
|
||||
|
||||
# Semantic grayscale fills — vocabulary mirrors `visualization-style/SKILL.md`.
|
||||
# Use these by ROLE, not by hex. The print track resolves them to grays;
|
||||
# the color track resolves them to the ML4T palette.
|
||||
GRAY_FILLS = {
|
||||
"primary": "#000000", # titles, lead data series, key emphasis
|
||||
"secondary": "#808080", # second series — widened from #404040 for print contrast
|
||||
"tertiary": "#c8c8c8", # third series, supporting elements
|
||||
"quaternary": "#e8e8e8", # fourth series only — keep grayscale separable
|
||||
"muted": "#a8a8a8", # de-emphasized, comparison baselines
|
||||
"border": "#666666", # connectors, axis lines (data side)
|
||||
"highlight": "#d9d9d9", # ~85% white — emphasis band fill
|
||||
"container": "#f2f2f2", # ~95% white — phase container fill
|
||||
"foundation": "#b3b3b3", # ~70% white — foundation layer fill
|
||||
"canvas": "#ffffff", # page background
|
||||
}
|
||||
|
||||
COLOR_FILLS = {
|
||||
"primary": COLORS["blue"], # #0a1628 navy — primary series
|
||||
"secondary": COLORS["amber"], # #D4A84B amber — secondary series
|
||||
"tertiary": COLORS["copper"], # #C87533 copper — tertiary (kept distinct from navy)
|
||||
"quaternary": COLORS["slate"], # #1a2d4a mid-blue — fourth series only
|
||||
"muted": COLORS["silver_muted"], # #e8e8e6
|
||||
"border": COLORS["neutral"], # #334155
|
||||
"highlight": COLORS["amber_light"],
|
||||
"container": COLORS["bg_light"],
|
||||
"foundation": COLORS["silver"],
|
||||
"canvas": "#ffffff",
|
||||
}
|
||||
|
||||
# Categorical cyclers for `axes.prop_cycle`. The print track pairs GRAY_CYCLER
|
||||
# with LINESTYLE_CYCLER so a B&W readout stays legible; the color track relies
|
||||
# on hue alone (no linestyle pairing — see apply_book_style). Color order
|
||||
# prioritizes perceptual separation for the first 4 entries (most figures use
|
||||
# ≤4 series); slate is positioned last because it reads as a second navy next
|
||||
# to blue. GRAY_CYCLER mirrors the GRAY_FILLS weight order (secondary widened
|
||||
# to #808080 for print contrast) while keeping every entry dark enough to read
|
||||
# as a line on white.
|
||||
COLOR_CYCLER = [
|
||||
COLORS["blue"], # navy — primary
|
||||
COLORS["amber"], # gold — secondary
|
||||
COLORS["copper"], # orange — tertiary
|
||||
COLORS["positive"], # green — fourth
|
||||
COLORS["negative"], # red — fifth (semantic, use sparingly)
|
||||
COLORS["slate"], # navy — sixth (only when ≥6 series; reads close to blue)
|
||||
]
|
||||
GRAY_CYCLER = ["#000000", "#808080", "#404040", "#a8a8a8", "#666666", "#c8c8c8"]
|
||||
LINESTYLE_CYCLER = ["-", "--", ":", "-.", "-", "--"]
|
||||
MARKER_CYCLER = ["o", "s", "^", "D", "v", "P"]
|
||||
|
||||
# =============================================================================
|
||||
# CANONICAL FIGURE SIZES (Packt embed width = 5.833")
|
||||
# =============================================================================
|
||||
# Width is fixed at 5.833" — the typeset width Packt uses in the manuscript
|
||||
# template. Heights are picked per layout so panels render at proportions
|
||||
# that don't dominate page vertical space. Use these in generate scripts;
|
||||
# do NOT introduce ad-hoc figsize tuples per figure.
|
||||
PAGE_WIDTH = 5.833 # Packt typeset embed width in inches
|
||||
|
||||
FIGSIZE = {
|
||||
"single_wide": (PAGE_WIDTH, 2.6), # short time series, comparisons
|
||||
"single": (PAGE_WIDTH, 3.4), # default single panel (~1.7:1)
|
||||
"single_tall": (PAGE_WIDTH, 4.0), # detail-heavy single panel
|
||||
"dual_h": (PAGE_WIDTH, 2.6), # two side-by-side panels
|
||||
"dual_h_tall": (PAGE_WIDTH, 3.2), # two side-by-side, taller panels
|
||||
"dual_v": (PAGE_WIDTH, 5.0), # two stacked panels
|
||||
"triple_h": (PAGE_WIDTH, 2.2), # three side-by-side panels, short
|
||||
"triple_h_tall": (PAGE_WIDTH, 3.0), # three side-by-side, detail
|
||||
"grid_2x2": (PAGE_WIDTH, 4.0), # 2 rows × 2 cols, simple axes
|
||||
"grid_2x3": (PAGE_WIDTH, 3.5), # 2 rows × 3 cols
|
||||
"grid_3x2": (PAGE_WIDTH, 5.5), # 3 rows × 2 cols (square-ish grid)
|
||||
"dashboard_2x2": (PAGE_WIDTH, 5.5), # 2×2 with date axes / rotated labels
|
||||
"dashboard_2x3": (PAGE_WIDTH, 4.5), # 2×3 with date axes / rotated labels
|
||||
}
|
||||
|
||||
|
||||
_BOOK_BASE_STYLE = {
|
||||
# Kept in sync with matplotlibrc at repo root. The auto-applied
|
||||
# matplotlibrc covers all default runs; this dict is the explicit-apply
|
||||
# override for book-figure scripts that swap between print and color
|
||||
# tracks via ``apply_book_style()``.
|
||||
"figure.dpi": 100,
|
||||
"figure.figsize": FIGSIZE["single"],
|
||||
"figure.facecolor": COLORS["bg_light"],
|
||||
"figure.constrained_layout.use": True,
|
||||
"savefig.dpi": 300,
|
||||
"savefig.bbox": "tight",
|
||||
"savefig.pad_inches": 0.05,
|
||||
"savefig.facecolor": COLORS["bg_light"],
|
||||
"axes.facecolor": COLORS["bg_light"],
|
||||
"axes.spines.top": False,
|
||||
"axes.spines.right": False,
|
||||
"axes.titlesize": 10,
|
||||
"axes.titleweight": "normal",
|
||||
"axes.titlelocation": "left",
|
||||
"axes.titlepad": 6,
|
||||
"axes.labelsize": 9,
|
||||
"axes.labelpad": 4,
|
||||
"axes.linewidth": 0.75,
|
||||
"axes.grid": False,
|
||||
"axes.axisbelow": True,
|
||||
"grid.linewidth": 0.5,
|
||||
"grid.alpha": 0.6,
|
||||
"grid.linestyle": "--",
|
||||
"xtick.labelsize": 8,
|
||||
"ytick.labelsize": 8,
|
||||
"xtick.major.size": 3,
|
||||
"ytick.major.size": 3,
|
||||
"xtick.major.width": 0.6,
|
||||
"ytick.major.width": 0.6,
|
||||
"xtick.direction": "out",
|
||||
"ytick.direction": "out",
|
||||
"lines.linewidth": 1.4,
|
||||
"lines.markersize": 4,
|
||||
"lines.markeredgewidth": 0,
|
||||
"legend.frameon": False,
|
||||
"legend.fontsize": 8,
|
||||
"legend.handlelength": 2.0,
|
||||
"font.family": ["sans-serif"],
|
||||
"font.sans-serif": ["Source Sans 3", "DejaVu Sans", "Helvetica", "Arial"],
|
||||
"font.size": 9,
|
||||
"image.cmap": "cividis",
|
||||
}
|
||||
|
||||
|
||||
def _cycler(colors: list[str], linestyles: list[str] | None = None):
|
||||
"""Build a prop_cycle from colors + optional linestyles. Local import keeps
|
||||
the module top-level cheap."""
|
||||
from cycler import cycler as cy
|
||||
|
||||
cyc = cy(color=colors)
|
||||
if linestyles is not None:
|
||||
cyc = cyc + cy(linestyle=linestyles[: len(colors)])
|
||||
return cyc
|
||||
|
||||
|
||||
BOOK_PRINT_STYLE = {
|
||||
# PRINT track is for the printed book — on white paper, so revert
|
||||
# the warm-cream backgrounds back to plain white.
|
||||
**_BOOK_BASE_STYLE,
|
||||
"figure.facecolor": "white",
|
||||
"savefig.facecolor": "white",
|
||||
"axes.facecolor": "white",
|
||||
"axes.edgecolor": "#333333",
|
||||
"axes.labelcolor": "#000000",
|
||||
"axes.titlecolor": "#000000",
|
||||
"xtick.color": "#000000",
|
||||
"ytick.color": "#000000",
|
||||
"grid.color": "#cccccc",
|
||||
"text.color": "#000000",
|
||||
}
|
||||
|
||||
BOOK_COLOR_STYLE = {
|
||||
# COLOR track is for web/README/Google Drive — matches the website's
|
||||
# warm-cream bg_light surface.
|
||||
**_BOOK_BASE_STYLE,
|
||||
"axes.edgecolor": COLORS["neutral"],
|
||||
"axes.labelcolor": COLORS["neutral"],
|
||||
"axes.titlecolor": COLORS["neutral"],
|
||||
"xtick.color": COLORS["neutral"],
|
||||
"ytick.color": COLORS["neutral"],
|
||||
"grid.color": COLORS["silver_muted"],
|
||||
"text.color": COLORS["neutral"],
|
||||
}
|
||||
|
||||
|
||||
def apply_book_style(mode: Literal["print", "color"] = "print") -> None:
|
||||
"""Set rcParams for a book-figure generation script.
|
||||
|
||||
Call once at script start (or before each render in a dual-track loop).
|
||||
Resolves the prop_cycle to grayscale (with linestyle variation) for
|
||||
``print`` and to the ML4T color palette for ``color``.
|
||||
"""
|
||||
style = BOOK_PRINT_STYLE if mode == "print" else BOOK_COLOR_STYLE
|
||||
plt.rcParams.update(style)
|
||||
if mode == "print":
|
||||
plt.rcParams["axes.prop_cycle"] = _cycler(GRAY_CYCLER, LINESTYLE_CYCLER)
|
||||
else:
|
||||
plt.rcParams["axes.prop_cycle"] = _cycler(COLOR_CYCLER)
|
||||
|
||||
|
||||
def save_dual(
|
||||
make_fig,
|
||||
output_basename: str,
|
||||
output_dir: str | Path,
|
||||
dpi: int = 300,
|
||||
) -> tuple[Path, Path]:
|
||||
"""Render and save both tracks of a publication figure.
|
||||
|
||||
``make_fig(palette, mode)`` is called twice — once with ``GRAY_FILLS`` /
|
||||
``"print"`` and once with ``COLOR_FILLS`` / ``"color"``. The print PNG
|
||||
lands at ``output_dir/{basename}.png`` (top-level grayscale default).
|
||||
The color PNG lands at ``output_dir/color/{basename}_color.png``.
|
||||
|
||||
Args:
|
||||
make_fig: Callable ``(palette: dict, mode: str) -> matplotlib.Figure``.
|
||||
Must build the figure from scratch each call — the caller closes it.
|
||||
output_basename: e.g. ``"figure_2_2_survivorship_bias"``. No extension.
|
||||
output_dir: chapter ``figures/`` directory.
|
||||
dpi: PNG resolution (default 300).
|
||||
|
||||
Returns:
|
||||
(print_path, color_path) — both absolute.
|
||||
"""
|
||||
output_dir = Path(output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
color_dir = output_dir / "color"
|
||||
color_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Print track first — that's the canonical artifact.
|
||||
apply_book_style("print")
|
||||
fig = make_fig(GRAY_FILLS, "print")
|
||||
print_path = output_dir / f"{output_basename}.png"
|
||||
fig.savefig(print_path, dpi=dpi, bbox_inches="tight", facecolor="white")
|
||||
plt.close(fig)
|
||||
|
||||
# Color track.
|
||||
apply_book_style("color")
|
||||
fig = make_fig(COLOR_FILLS, "color")
|
||||
color_path = color_dir / f"{output_basename}_color.png"
|
||||
fig.savefig(color_path, dpi=dpi, bbox_inches="tight", facecolor="white")
|
||||
plt.close(fig)
|
||||
|
||||
return print_path, color_path
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# BOOK-SPECIFIC FUNCTIONS (LEGACY)
|
||||
# =============================================================================
|
||||
|
||||
# DEPRECATED: Style is now applied automatically via matplotlibrc in repo root.
|
||||
ML4T_STYLE = Path(__file__).parent.parent / "matplotlibrc"
|
||||
|
||||
|
||||
def save_figure(
|
||||
fig,
|
||||
name: str,
|
||||
chapter: str | None = None,
|
||||
formats: list[str] | None = None,
|
||||
dpi: int = 150,
|
||||
) -> None:
|
||||
"""Save figure with ML4T conventions.
|
||||
|
||||
Args:
|
||||
fig: matplotlib figure object
|
||||
name: Base filename (without extension)
|
||||
chapter: Optional chapter directory (e.g., "06_alpha_factor_engineering")
|
||||
formats: List of formats to save (default: ['png', 'pdf'])
|
||||
dpi: Resolution for raster formats (default: 150)
|
||||
"""
|
||||
formats = formats or ["png", "pdf"]
|
||||
|
||||
if chapter:
|
||||
repo_root = Path(__file__).parent.parent
|
||||
output_dir = repo_root / chapter / "visualizations"
|
||||
else:
|
||||
output_dir = Path(".")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for fmt in formats:
|
||||
output_path = output_dir / f"{name}.{fmt}"
|
||||
fig.savefig(output_path, format=fmt, dpi=dpi, bbox_inches="tight")
|
||||
print(f"Saved: {output_path}")
|
||||
|
||||
|
||||
def plot_fidelity_comparison(
|
||||
real_data: np.ndarray,
|
||||
synthetic_data: np.ndarray,
|
||||
title: str = "Real vs Synthetic Distribution",
|
||||
n_samples: int = 1000,
|
||||
figsize: tuple = (12, 5),
|
||||
flatten_method: str = "mean",
|
||||
random_state: int = 42,
|
||||
) -> plt.Figure:
|
||||
"""Create standardized fidelity comparison plot using PCA and t-SNE.
|
||||
|
||||
Designed for grayscale compatibility:
|
||||
- Real data: dark circles (filled)
|
||||
- Synthetic data: amber X markers (open)
|
||||
|
||||
Args:
|
||||
real_data: Real sequences. Shape can be:
|
||||
- (n_samples, seq_len, n_features): 3D time series
|
||||
- (n_samples, n_features): 2D tabular
|
||||
synthetic_data: Synthetic sequences, same shape as real_data
|
||||
title: Plot title
|
||||
n_samples: Number of samples to visualize (subsampled if larger)
|
||||
figsize: Figure size (width, height)
|
||||
flatten_method: How to flatten 3D data to 2D:
|
||||
- "mean": Average across time dimension (default)
|
||||
- "last": Use last timestep only
|
||||
- "flatten": Concatenate all timesteps (high-dim)
|
||||
random_state: Random seed for reproducibility
|
||||
|
||||
Returns:
|
||||
matplotlib Figure object
|
||||
"""
|
||||
from sklearn.decomposition import PCA
|
||||
from sklearn.manifold import TSNE
|
||||
|
||||
np.random.seed(random_state)
|
||||
|
||||
# Handle 3D (time series) vs 2D (tabular) data
|
||||
if real_data.ndim == 3:
|
||||
if flatten_method == "mean":
|
||||
real_flat = real_data.mean(axis=1)
|
||||
synth_flat = synthetic_data.mean(axis=1)
|
||||
elif flatten_method == "last":
|
||||
real_flat = real_data[:, -1, :]
|
||||
synth_flat = synthetic_data[:, -1, :]
|
||||
elif flatten_method == "flatten":
|
||||
real_flat = real_data.reshape(real_data.shape[0], -1)
|
||||
synth_flat = synthetic_data.reshape(synthetic_data.shape[0], -1)
|
||||
else:
|
||||
raise ValueError(f"Unknown flatten_method: {flatten_method}")
|
||||
else:
|
||||
real_flat = real_data
|
||||
synth_flat = synthetic_data
|
||||
|
||||
# Subsample for visualization
|
||||
n_viz = min(n_samples, len(real_flat), len(synth_flat))
|
||||
idx_real = np.random.choice(len(real_flat), n_viz, replace=False)
|
||||
idx_synth = np.random.choice(len(synth_flat), n_viz, replace=False)
|
||||
|
||||
real_sample = real_flat[idx_real]
|
||||
synth_sample = synth_flat[idx_synth]
|
||||
|
||||
# PCA - fit on real, transform both
|
||||
n_features = real_sample.shape[1] if real_sample.ndim > 1 else 1
|
||||
n_pca = min(2, n_features, n_viz)
|
||||
pca = PCA(n_components=n_pca)
|
||||
pca.fit(real_sample)
|
||||
real_pca = pca.transform(real_sample)
|
||||
synth_pca = pca.transform(synth_sample)
|
||||
|
||||
# t-SNE - fit jointly for proper comparison
|
||||
combined = np.vstack([real_sample, synth_sample])
|
||||
perplexity = min(40, max(2, n_viz // 4))
|
||||
n_tsne = min(2, n_features)
|
||||
tsne = TSNE(
|
||||
n_components=n_tsne, perplexity=perplexity, max_iter=1000, random_state=random_state
|
||||
)
|
||||
combined_tsne = tsne.fit_transform(combined)
|
||||
real_tsne = combined_tsne[:n_viz]
|
||||
synth_tsne = combined_tsne[n_viz:]
|
||||
|
||||
# Create figure with aligned axes
|
||||
fig, axes = plt.subplots(1, 2, figsize=figsize)
|
||||
|
||||
# Style constants for grayscale compatibility
|
||||
real_color = COLORS["blue"]
|
||||
synth_color = COLORS["amber"]
|
||||
marker_size = 25
|
||||
alpha = 0.6
|
||||
|
||||
# PCA plot (handle 1D case when n_features < 2)
|
||||
pca_y_real = real_pca[:, 1] if n_pca >= 2 else np.zeros(len(real_pca))
|
||||
pca_y_synth = synth_pca[:, 1] if n_pca >= 2 else np.zeros(len(synth_pca))
|
||||
axes[0].scatter(
|
||||
real_pca[:, 0],
|
||||
pca_y_real,
|
||||
c=real_color,
|
||||
marker="o",
|
||||
s=marker_size,
|
||||
alpha=alpha,
|
||||
label="Real",
|
||||
edgecolors="none",
|
||||
)
|
||||
axes[0].scatter(
|
||||
synth_pca[:, 0],
|
||||
pca_y_synth,
|
||||
c=synth_color,
|
||||
marker="x",
|
||||
s=marker_size,
|
||||
alpha=alpha,
|
||||
label="Synthetic",
|
||||
linewidths=1.5,
|
||||
)
|
||||
axes[0].set_xlabel("PC1")
|
||||
axes[0].set_ylabel("PC2" if n_pca >= 2 else "")
|
||||
axes[0].set_title("PCA Projection")
|
||||
axes[0].legend(loc="upper right", framealpha=0.9)
|
||||
|
||||
# t-SNE plot (handle 1D case when n_features < 2)
|
||||
tsne_y_real = real_tsne[:, 1] if n_tsne >= 2 else np.zeros(len(real_tsne))
|
||||
tsne_y_synth = synth_tsne[:, 1] if n_tsne >= 2 else np.zeros(len(synth_tsne))
|
||||
axes[1].scatter(
|
||||
real_tsne[:, 0],
|
||||
tsne_y_real,
|
||||
c=real_color,
|
||||
marker="o",
|
||||
s=marker_size,
|
||||
alpha=alpha,
|
||||
label="Real",
|
||||
edgecolors="none",
|
||||
)
|
||||
axes[1].scatter(
|
||||
synth_tsne[:, 0],
|
||||
tsne_y_synth,
|
||||
c=synth_color,
|
||||
marker="x",
|
||||
s=marker_size,
|
||||
alpha=alpha,
|
||||
label="Synthetic",
|
||||
linewidths=1.5,
|
||||
)
|
||||
axes[1].set_xlabel("t-SNE 1")
|
||||
axes[1].set_ylabel("t-SNE 2" if n_tsne >= 2 else "")
|
||||
axes[1].set_title("t-SNE Projection")
|
||||
axes[1].legend(loc="upper right", framealpha=0.9)
|
||||
|
||||
fig.suptitle(title, fontsize=14, fontweight="semibold", y=1.02)
|
||||
plt.tight_layout()
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MODULE EXPORTS
|
||||
# =============================================================================
|
||||
|
||||
__all__ = [
|
||||
# Palette
|
||||
"COLORS",
|
||||
"GRAYSCALE",
|
||||
# Matplotlib styles
|
||||
"ML4T_LIGHT_STYLE",
|
||||
"ML4T_DARK_STYLE",
|
||||
# Style application
|
||||
"apply_ml4t_style",
|
||||
# Palette helpers
|
||||
"ml4t_palette",
|
||||
"ml4t_diverging",
|
||||
# Chart helpers
|
||||
"annotate_peak",
|
||||
"add_regime_shading",
|
||||
"format_pct_axis",
|
||||
# Book-specific
|
||||
"ML4T_STYLE",
|
||||
"HAS_PLOTLY",
|
||||
"save_figure",
|
||||
"plot_fidelity_comparison",
|
||||
]
|
||||
Reference in New Issue
Block a user