wehub/stefan-jansen--machine-learning-for-trading
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:26:28 +08:00
commit 0249a267e1
495 changed files with 1277094 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/README.md
+237
View File
@@ -0,0 +1,237 @@
# Notebook Test Suite
Every notebook in this repository is tested via [Papermill](https://papermill.readthedocs.io/) parameter injection. The same code path always runs — only the data scale differs between production and test.
---
## Quick Start
```bash
# Run all environments (ml4t, gpu, py312, benchmark, neo4j)
./scripts/run_all_tests.sh
# Run one environment
./scripts/run_all_tests.sh ml4t
# Rerun everything (ignore already-passed)
./scripts/run_all_tests.sh --force
# Run a specific notebook via pytest
docker compose run --rm ml4t pytest tests/test_chapter_notebooks.py -v -k "01_timegan"
# Run locally (with uv)
uv run pytest tests/test_chapter_notebooks.py -v -k "11_ml_pipeline"
```
---
## How It Works
### Papermill Parameter Injection
Every notebook has a `# %% tags=["parameters"]` cell with **production defaults** — the values readers see in the book:
```python
# %% tags=["parameters"]
MAX_SYMBOLS = 0 # 0 = all symbols (production)
N_EPOCHS = 500
START_DATE = "2006-01-01"
```
During testing, Papermill creates an **injected cell** after the tagged cell that overrides selected values:
```python
# Injected by Papermill
MAX_SYMBOLS = 15 # Reduced for fast execution
N_EPOCHS = 2
```
The notebook code sees only the final (overridden) values. **There are no `if TEST:` branches** — the same code path always runs, just with less data or fewer iterations.
### Output Isolation
Tests set `ML4T_OUTPUT_DIR` to a temp directory. All notebook writes (`get_output_dir()`, `get_case_study_dir()`) redirect there, so production artifacts (trained models, backtest results) are never overwritten.
### Seeded Fixtures
Downstream notebooks (Ch11+) need upstream pipeline outputs (features, labels, model predictions). Rather than running the full pipeline during every test, `tests/fixtures/seed_results.py` creates minimal but schema-correct fixtures:
- Registry databases (`run_log/registry.db`) with realistic training_runs, prediction_sets, and backtest entries
- Feature parquets (350 rows × 15 symbols)
- Label parquets (5 label variants per case study)
- Prediction parquets (200 rows per model)
Fixtures are deterministic and only written if the file doesn't already exist — real upstream results take priority.
---
## Test Data
### Architecture
Tests require two things: **raw market data** (what loaders read) and **pipeline intermediates** (what downstream notebooks consume). Both live in a private repo that CI pulls automatically.
```
~/ml4t/test-data/ # Local clone of ml4t/third-edition-test-data (~2 GB private GitHub repo)
├── data/ # Pre-subsampled raw data (553 MB)
│ ├── etfs/ # 15 most liquid ETFs (full date range)
│ ├── crypto/ # 5 largest perps
│ ├── futures/ # 8 CME products
│ ├── fx/ # 8 major pairs
│ ├── equities/ # 50 US stocks, 3 NASDAQ-100 (minute bars)
│ │ ├── microstructure/ # Synthetic ITCH/LOB/MBO for Ch03
│ │ └── firm_characteristics/ # 200 most-observed per month
│ ├── factors/ # Fama-French + AQR
│ └── manifest.json # Symbol counts per dataset
│
└── intermediates/ # Pre-computed pipeline outputs (301 MB)
├── etfs/ # registry.db, features, labels, predictions
├── crypto_perps_funding/
├── ... # All 9 case studies
└── _metadata.json # Generation timestamp and parameters
```
**Key design decisions:**
- **Same schema, fewer symbols**: Loaders work without code changes. Full date ranges are preserved so cross-validation folds remain valid.
- **Pre-computed intermediates**: Running all 9 case study pipelines from scratch takes ~25 minutes. Pre-computing once and shipping the results lets CI focus on testing individual notebooks.
- **Fixtures as safety net**: If intermediates are missing or a new notebook needs data that wasn't pre-computed, `seed_results.py` generates minimal placeholders at test time.
### Running Tests Locally
**With the test-data repo** (full test coverage):
```bash
# Clone test data once
git clone git@github.com:ml4t/third-edition-test-data.git ~/ml4t/test-data
# Point tests at it
export ML4T_DATA_PATH=~/ml4t/test-data/data
export ML4T_OUTPUT_DIR=/tmp/ml4t-test-output
# Seed intermediates
mkdir -p $ML4T_OUTPUT_DIR
cp -r ~/ml4t/test-data/intermediates/* $ML4T_OUTPUT_DIR/
# Run tests
uv run pytest tests/test_chapter_notebooks.py -v -k "11_ml_pipeline"
```
**Without the test-data repo** (limited coverage):
```bash
# Point at your production data
export ML4T_DATA_PATH=/path/to/your/data
export ML4T_OUTPUT_DIR=/tmp/ml4t-test-output
# Fixtures will be auto-generated for missing intermediates
uv run pytest tests/test_chapter_notebooks.py -v -k "01_ols_inference"
```
### Regenerating Test Data
If the data schema or pipeline logic changes, regenerate the test data:
```bash
# 1. Subsample raw data
uv run python tests/create_test_data.py --output ~/ml4t/test-data
# 2. Generate pipeline intermediates (runs all 9 case studies, ~25 min)
ML4T_DATA_PATH=~/ml4t/test-data/data \
ML4T_OUTPUT_DIR=~/ml4t/test-data/intermediates \
uv run python tests/_internal/generate_intermediates.py
# 3. Commit and push to test-data repo
cd ~/ml4t/test-data && git add -A && git commit -m "regenerate test data"
```
---
## Environments
Each notebook is assigned to exactly one Docker environment via `docker_env` in `overrides.yaml`.
| Environment | Docker Service | Notebooks | What It Provides |
|-------------|---------------|-----------|-----------------|
| `ml4t` | `ml4t` | ~410 | CPU, all Python packages |
| `gpu` | `ml4t-gpu` | ~31 | NVIDIA GPU (PyTorch CUDA) |
| `py312` | `py312` | ~10 | gensim, signatory, esig, pfhedge, tfcausalimpact (Python 3.12) |
| `benchmark` | `benchmark` + database services | 2 | TimescaleDB, ClickHouse, QuestDB, InfluxDB |
| `neo4j` | `ml4t` + Neo4j service | 7 | Neo4j graph database |
Notebooks default to `ml4t` unless tagged otherwise. Multi-environment tags (e.g., `docker_env: ml4t+neo4j+gpu`) require all listed services.
The test runner (`scripts/run_all_tests.sh`) iterates over environments, running only notebooks tagged for each one.
---
## Override Format
Per-notebook configuration lives in `tests/overrides.yaml`:
```yaml
05_synthetic_data/01_timegan:
docker_env: gpu # Runs only in GPU environment
timeout: 600 # Max seconds before test is killed
parameters: # Papermill parameter overrides
TRAIN_STEPS: 100
BATCH_SIZE: 32
case_studies/etfs/07_gbm:
timeout: 300
parameters:
MAX_SYMBOLS: 15
START_DATE: "2020-01-01"
26_mlops_governance/05b_feast_live:
skip: true # Never runs
skip_reason: "Requires Feast feature server"
```
**Override fields:**
| Field | Default | Purpose |
|-------|---------|---------|
| `timeout` | 300 | Max seconds per notebook |
| `docker_env` | `ml4t` | Which Docker environment to use |
| `skip` | false | Skip this notebook entirely |
| `skip_reason` | — | Reason displayed in test output |
| `parameters` | {} | Papermill parameter overrides |
---
## Files
| File | Purpose |
|------|---------|
| `test_chapter_notebooks.py` | Parametrized tests for Ch01-Ch27 teaching notebooks |
| `test_case_studies.py` | Parametrized tests for all 9 case study pipelines |
| `test_backtest_schedule.py` | Backtest-specific integration tests |
| `conftest.py` | Session fixtures: data dirs, output seeding, config patching |
| `pm_helpers.py` | Papermill execution, override loading, Docker env detection |
| `overrides.yaml` | Per-notebook parameter overrides, timeouts, skip/env tags |
| `fixtures/seed_results.py` | Registry DB and parquet fixture generation |
| `_internal/` | Scripts for generating test data and intermediates (not run during tests) |
---
## CI Pipeline
GitHub Actions runs tests on every push and PR:
1. **Checkout** code + test data repo (via deploy key)
2. **Seed** intermediates into `ML4T_OUTPUT_DIR`
3. **Run** pytest inside Docker containers (one job per environment)
4. **Upload** JUnit XML results as artifacts
See `.github/workflows/test.yml` for the full configuration.
---
## Adding a New Notebook
1. Add a `# %% tags=["parameters"]` cell with production defaults
2. Add an entry in `overrides.yaml` with timeout and parameter overrides
3. If GPU-required, add `docker_env: gpu`
4. Run: `./scripts/run_all_tests.sh ml4t` (or the relevant environment)
5. If the notebook depends on upstream pipeline outputs, ensure `seed_results.py` generates the necessary fixtures
tests/__init__.py
View File
tests/add_missing_parameters_cells.py
+118
View File
@@ -0,0 +1,118 @@
"""Add missing Papermill parameters cells to notebooks.
Notebooks that already have a `# %% tags=["parameters"]` cell are skipped.
For notebooks without one, this inserts an empty parameters cell after the
first code cell (or after imports if detectable).
Usage:
uv run python tests/add_missing_parameters_cells.py [--dry-run]
"""
import re
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).parent.parent
DRY_RUN = "--dry-run" in sys.argv
# Parameters cell template
PARAMS_CELL = """
# %% tags=["parameters"]
# Production defaults — Papermill injects overrides for CI
"""
def find_notebooks() -> list[Path]:
"""Find all Jupytext notebooks missing a parameters cell."""
notebooks = []
for d in sorted(REPO_ROOT.glob("[0-9]*_*")):
notebooks.extend(sorted(d.glob("**/*.py")))
for d in sorted((REPO_ROOT / "case_studies").glob("*")):
if d.is_dir() and d.name not in ("utils", "__pycache__"):
notebooks.extend(sorted(d.glob("**/*.py")))
missing = []
for nb in notebooks:
content = nb.read_text()
if "# %%" not in content:
continue
if 'tags=["parameters"]' in content or "tags=['parameters']" in content:
continue
missing.append(nb)
return missing
def add_parameters_cell(nb_path: Path) -> bool:
"""Insert a parameters cell after the first code cell.
Strategy: find the first `# %%` code cell (not markdown, not the header)
and insert the parameters cell after it. If the first code cell has imports,
the parameters cell goes after the import block.
"""
content = nb_path.read_text()
lines = content.split("\n")
# Find insertion point: after the first code cell
# Skip: header (---/jupyter metadata), markdown cells
in_header = False
found_first_code = False
insert_after = None
for i, line in enumerate(lines):
stripped = line.strip()
# Track Jupytext header
if i == 0 and stripped == "# ---":
in_header = True
continue
if in_header:
if stripped == "# ---":
in_header = False
continue
# Found a code cell
if stripped == "# %%" and not found_first_code:
found_first_code = True
# Look ahead to find end of this cell (next # %% or end)
for j in range(i + 1, len(lines)):
if lines[j].strip().startswith("# %%"):
insert_after = j
break
else:
insert_after = len(lines)
break
if insert_after is None:
return False
# Insert the parameters cell
new_lines = (
lines[:insert_after] + PARAMS_CELL.rstrip().split("\n") + [""] + lines[insert_after:]
)
new_content = "\n".join(new_lines)
if not DRY_RUN:
nb_path.write_text(new_content)
return True
def main():
missing = find_notebooks()
print(f"Found {len(missing)} notebooks missing parameters cell")
if DRY_RUN:
print("(dry run — no changes)")
modified = 0
for nb in missing:
rel = nb.relative_to(REPO_ROOT)
ok = add_parameters_cell(nb)
status = "ADDED" if ok else "SKIPPED"
if ok:
modified += 1
print(f" {status}: {rel}")
print(f"\nModified: {modified}/{len(missing)}")
if __name__ == "__main__":
main()
tests/conftest.py
+490
View File
@@ -0,0 +1,490 @@
"""Pytest fixtures for ML4T test suite.
Two modes of operation:
1. CI (GHA): ML4T_DATA_PATH points to pre-subsampled real data (from private repo).
populated_data_dir just returns that path — no synthetic data needed.
2. Local dev: ML4T_DATA_PATH points to full production data or test data.
"""
import json
import os
import shutil
import time
from pathlib import Path
import pytest
import yaml
REPO_ROOT = Path(__file__).parent.parent
# Case study IDs whose config/setup.yaml should be seeded into test output dirs
CASE_STUDY_IDS = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
@pytest.fixture(scope="session", autouse=True)
def ci_env_setup():
"""Create .env file if running in CI (where ML4T_DATA_PATH is set externally).
utils/config.py requires a .env file to exist.
In CI, environment variables are set by the workflow, but the .env
file still needs to exist to avoid FileNotFoundError on import.
"""
env_file = REPO_ROOT / ".env"
created = False
if not env_file.exists():
# Create minimal .env for CI
env_file.write_text(
f"ML4T_PATH={REPO_ROOT}\n"
f"ML4T_DATA_PATH={os.environ.get('ML4T_DATA_PATH', REPO_ROOT / 'data')}\n"
)
created = True
yield
# Clean up CI-created .env (don't leave artifacts)
if created and env_file.exists():
env_file.unlink()
def _resolve_data_path() -> Path | None:
"""Find ML4T_DATA_PATH from env var, .env file, or default location.
pytest-xdist workers may not inherit env vars set by the parent process,
so we also check the .env file and well-known test-data locations.
"""
# 1. Explicit env var (works in single-process pytest and CI)
env_path = os.environ.get("ML4T_DATA_PATH")
if env_path:
p = Path(env_path).expanduser().resolve()
if p.exists() and any(p.iterdir()):
return p
# 2. Read from .env file (works in xdist workers)
env_file = REPO_ROOT / ".env"
if env_file.exists():
for line in env_file.read_text().splitlines():
line = line.strip()
if line.startswith("ML4T_DATA_PATH") and "=" in line:
val = line.split("=", 1)[1].strip().strip('"').strip("'")
if val and not val.startswith("#"):
p = Path(val).expanduser().resolve()
if p.exists() and any(p.iterdir()):
return p
# 3. Well-known test-data repo location
test_data = Path.home() / "ml4t" / "test-data" / "data"
if test_data.exists() and (test_data / "etfs").exists():
return test_data
# 4. Default: repo's own data/ directory
repo_data = REPO_ROOT / "data"
if repo_data.exists() and any(repo_data.glob("*/*.parquet")):
return repo_data
return None
@pytest.fixture(scope="session")
def test_data_dir(tmp_path_factory):
"""Return the data directory for tests.
Resolves ML4T_DATA_PATH from env var, .env file, well-known test-data
repo location, or repo's data/ directory. Works with pytest-xdist.
"""
resolved = _resolve_data_path()
if resolved:
os.environ["ML4T_DATA_PATH"] = str(resolved)
return resolved
# Fallback: create temp directory
data_dir = tmp_path_factory.mktemp("ml4t_data")
os.environ["ML4T_DATA_PATH"] = str(data_dir)
return data_dir
@pytest.fixture(scope="session")
def populated_data_dir(test_data_dir):
"""Return a data directory populated with test data.
If ML4T_DATA_PATH points to pre-populated data (e.g., from GHA checkout
of ml4t/third-edition-test-data), returns it directly.
"""
if (test_data_dir / "etfs" / "market" / "etf_universe.parquet").exists():
return test_data_dir
pytest.skip("No test data available. Set ML4T_DATA_PATH or run in CI.")
@pytest.fixture(scope="session")
def intermediates_dir(test_data_dir):
"""Return directory with pre-computed pipeline intermediates.
When running downstream chapters (Ch11+), they need labels/features
from pipeline stages. These are pre-computed and stored in test-data repo.
"""
idir = test_data_dir.parent / "intermediates"
if idir.exists() and any(idir.iterdir()):
return idir
return None
@pytest.fixture(scope="session")
def seeded_output_dir(tmp_path_factory):
"""Session-scoped output dir seeded with case study config files.
Chapter notebooks that read case study setup.yaml (via get_case_study_dir())
need these configs to exist even when ML4T_OUTPUT_DIR redirects writes to
a temp directory. This fixture copies the real config files into the test
output dir so notebooks can find them.
With pytest-xdist, each worker gets its own subdirectory to avoid races
on shutil.rmtree/copytree when multiple workers seed simultaneously.
"""
base_dir = os.environ.get("ML4T_OUTPUT_DIR")
if base_dir:
# With xdist, append worker id to avoid races
worker_id = os.environ.get("PYTEST_XDIST_WORKER", "")
if worker_id:
output_dir = Path(base_dir) / f"worker_{worker_id}"
else:
output_dir = Path(base_dir)
else:
output_dir = tmp_path_factory.mktemp("ml4t_output")
# Set the env var so notebooks see this worker's output dir
os.environ["ML4T_OUTPUT_DIR"] = str(output_dir)
cs_root = REPO_ROOT / "case_studies"
# Copy per-case-study config files (setup.yaml, training menus, backtest presets, etc.)
for cs_id in CASE_STUDY_IDS:
src_config_dir = cs_root / cs_id / "config"
if not src_config_dir.exists():
continue
dst_config_dir = output_dir / cs_id / "config"
if dst_config_dir.exists():
shutil.rmtree(dst_config_dir)
shutil.copytree(src_config_dir, dst_config_dir)
_trim_label_configs(dst_config_dir)
# Copy global model presets (case_studies/config/) so load_configs() can find them.
# load_configs() resolves presets at {case_dir.parent}/config/{model_type}/*.yaml
# We copy (not symlink) so we can patch presets for fast testing.
global_config_src = cs_root / "config"
global_config_dst = output_dir / "config"
if global_config_src.exists():
if global_config_dst.exists():
shutil.rmtree(global_config_dst)
shutil.copytree(global_config_src, global_config_dst)
_patch_presets_for_testing(global_config_dst)
# Copy pipeline intermediates (features, labels, run_log) from test-data repo.
# These are pre-computed so downstream notebooks (Ch11+) can run without
# executing the full pipeline first.
# Look for intermediates next to data (test-data repo layout) or at well-known path.
data_path = _resolve_data_path()
intermediates_root = None
if data_path:
candidate = Path(data_path).parent / "intermediates"
if candidate.exists():
intermediates_root = candidate
if intermediates_root is None:
# Well-known test-data repo location
candidate = Path.home() / "ml4t" / "test-data" / "intermediates"
if candidate.exists():
intermediates_root = candidate
if intermediates_root and intermediates_root.exists():
for cs_id in CASE_STUDY_IDS:
src = intermediates_root / cs_id
if not src.exists():
continue
dst = output_dir / cs_id
# Copy features, labels, evaluation, run_log, results, benchmark —
# anything that downstream notebooks look for in get_case_study_dir()
for subdir in ["features", "labels", "evaluation", "run_log", "results", "benchmark"]:
src_sub = src / subdir
dst_sub = dst / subdir
if src_sub.exists() and not dst_sub.exists():
shutil.copytree(src_sub, dst_sub)
# Copy top-level intermediate files (e.g. etfs/eligibility.csv,
# protocol.yaml, baseline_checkpoint.yaml) that sit directly in
# intermediates/{cs_id}/ rather than in a subdir. Downstream
# notebooks (etfs 02_labels, 03_financial_features) read these via
# get_case_study_dir(); without this they fail with FileNotFoundError.
for item in src.iterdir():
if item.is_file():
dst_file = dst / item.name
if not dst_file.exists():
dst.mkdir(parents=True, exist_ok=True)
shutil.copy2(item, dst_file)
# Schema reconciliation: test-data predictions parquets were
# generated with an older column convention (y_score / y_true /
# fold_id). Production registry uses (prediction / actual / fold).
# Rename in place so notebooks reading via get_case_study_dir()
# see the canonical names without per-notebook compat shims.
preds_root = dst / "run_log" / "predictions"
if preds_root.exists():
_migrate_predictions_schema(preds_root)
# Copy non-case-study intermediates (chapter-scoped fixtures).
# These are intermediates that downstream teaching notebooks need but aren't
# part of the per-case-study pipeline (e.g., Ch16 signal comparison, Ch20 synthesis).
if intermediates_root and intermediates_root.exists():
for extra_id in ["ch16_signal_method_comparison", "ch20_synthesis"]:
src = intermediates_root / extra_id
if not src.exists():
continue
dst = output_dir / extra_id
if not dst.exists():
shutil.copytree(src, dst)
# Seed minimal results fixtures so downstream notebooks (latent factors, DL,
# backtest) can find baseline results without depending on upstream execution.
# These fill gaps where intermediates don't provide enough (e.g., Ch25 demo
# predictions, Ch15 causal JSON, synthetic registry entries).
from tests.fixtures.seed_results import seed_results
seed_results(output_dir, CASE_STUDY_IDS)
# Symlink AQR factor data so AQRFactorProvider finds it at ~/ml4t/data/aqr_factors
aqr_src = data_path.parent / "data" / "factors" / "aqr" if data_path else None
if aqr_src is None:
aqr_src = Path.home() / "ml4t" / "test-data" / "data" / "factors" / "aqr"
aqr_dst = Path.home() / "ml4t" / "data" / "aqr_factors"
if aqr_src.exists() and not aqr_dst.exists():
aqr_dst.parent.mkdir(parents=True, exist_ok=True)
aqr_dst.symlink_to(aqr_src)
return output_dir
# ---------------------------------------------------------------------------
# Preset patching — reduce workload for CI/test runs
# ---------------------------------------------------------------------------
# Per-model-type overrides applied to copied preset YAMLs.
# Goal: minimal workload that still exercises the training loop + registry.
_TEST_PRESET_PATCHES: dict[str, dict] = {
"lgb": {"max_iterations": 2, "checkpoint_interval": 1},
# DL families: 2 epochs, checkpoint every epoch
"lstm": {"n_epochs": 2, "checkpoint_interval": 1},
"tsmixer": {"n_epochs": 2, "checkpoint_interval": 1},
"tcn": {"n_epochs": 2, "checkpoint_interval": 1},
"nlinear": {"n_epochs": 2, "checkpoint_interval": 1},
"patchtst": {"n_epochs": 2, "checkpoint_interval": 1},
# TabDL: 2 epochs
"tabm": {"n_epochs": 2, "checkpoint_interval": 1},
# Latent factors: 2 epochs
"cae": {"n_epochs": 2, "checkpoint_interval": 1},
"sdf": {"n_epochs": 2, "checkpoint_interval": 1},
"sae": {"n_epochs": 2, "checkpoint_interval": 1},
"ipca": {"n_epochs": 2, "checkpoint_interval": 1},
}
_PREDICTION_COL_RENAMES = {
"y_score": "prediction",
"y_true": "actual",
"fold_id": "fold",
}
def _migrate_predictions_schema(preds_root: Path) -> None:
"""Rename test-data prediction columns to canonical production schema.
Test-data parquets were generated with an older convention
(y_score / y_true / fold_id). Production registry uses
(prediction / actual / fold). Walking the seeded predictions tree once
avoids per-notebook compat shims while keeping test-data immutable on
its own repo schedule.
"""
import polars as pl
for parquet in preds_root.rglob("predictions.parquet"):
cols = pl.read_parquet(parquet, n_rows=0).columns
renames = {old: new for old, new in _PREDICTION_COL_RENAMES.items() if old in cols}
if not renames:
continue
df = pl.read_parquet(parquet).rename(renames)
df.write_parquet(parquet)
def _patch_presets_for_testing(config_dir: Path) -> None:
"""Patch copied preset YAMLs with reduced-workload values for testing."""
for model_type, overrides in _TEST_PRESET_PATCHES.items():
model_dir = config_dir / model_type
if not model_dir.exists():
continue
for preset_path in model_dir.glob("*.yaml"):
preset = yaml.safe_load(preset_path.read_text())
if preset is None:
continue
preset.update(overrides)
with open(preset_path, "w") as f:
yaml.dump(preset, f, default_flow_style=False)
# Max configs per family in label config files (keep tests fast but comprehensive).
# Only applied to families with homogeneous sweep configs (linear, gbm).
# DL/TabDL/latent/causal families are NOT trimmed because each config often
# maps to a dedicated notebook (e.g., 09_dl_lstm, 10_dl_tsmixer).
_MAX_CONFIGS_PER_FAMILY = 2
_TRIM_FAMILIES = {"linear", "gbm"}
def _trim_label_configs(cs_config_dir: Path) -> None:
"""Trim training menu YAMLs to at most _MAX_CONFIGS_PER_FAMILY for sweep families."""
training_dir = cs_config_dir / "training"
label_root = training_dir if training_dir.exists() else cs_config_dir
for label_yaml in label_root.glob("fwd_*.yaml"):
data = yaml.safe_load(label_yaml.read_text())
if data is None or not isinstance(data, dict):
continue
trimmed = False
for family, configs in data.items():
if (
family in _TRIM_FAMILIES
and isinstance(configs, list)
and len(configs) > _MAX_CONFIGS_PER_FAMILY
):
data[family] = configs[:_MAX_CONFIGS_PER_FAMILY]
trimmed = True
if trimmed:
with open(label_yaml, "w") as f:
yaml.dump(data, f, default_flow_style=False)
# ---------------------------------------------------------------------------
# GPU marker — apply `@pytest.mark.gpu` at collection time based on overrides.
# Usage: pytest -m gpu (GPU only) | pytest -m "not gpu" (CPU only)
# ---------------------------------------------------------------------------
def pytest_configure(config):
config.addinivalue_line("markers", "gpu: notebook requires GPU (from overrides.yaml gpu: true)")
config.addinivalue_line(
"markers",
"long_running: notebook takes >10min even with reduced params (from overrides.yaml long_running: true)",
)
config.addinivalue_line(
"markers",
"weekly: notebook tier=weekly — runs only in scheduled weekly-external workflow. "
"To execute locally, set ML4T_TEST_TIER=weekly alongside `pytest -m weekly`; "
"without the env var, matching items are collected but skipped.",
)
config.addinivalue_line(
"markers",
"on_demand: notebook tier=on_demand — runs only on manual dispatch (e.g., GPU Tier 3). "
"To execute locally, set ML4T_TEST_TIER=on_demand alongside `pytest -m on_demand`; "
"without the env var, matching items are collected but skipped.",
)
config.addinivalue_line(
"markers",
"drift: live external drift smoke (Phase 3) — one tiny real pull per free "
"data source. Network-bound; runs only in the scheduled weekly-external "
"workflow's `drift` job, never per-PR.",
)
def pytest_collection_modifyitems(items):
"""Add markers to test items based on overrides.yaml flags."""
from tests.pm_helpers import (
TIER_ON_DEMAND,
TIER_WEEKLY,
get_overrides,
get_reruns,
get_tier,
)
# pytest-rerunfailures provides @pytest.mark.flaky(reruns=N). Detect once
# so per-NB reruns kick in automatically when the dep lands in Step 2.
try:
import pytest_rerunfailures # noqa: F401
has_rerunfailures = True
except ImportError:
has_rerunfailures = False
for item in items:
if hasattr(item, "callspec") and "notebook_path" in item.callspec.params:
nb_path = item.callspec.params["notebook_path"]
rel = (
nb_path.relative_to(REPO_ROOT).with_suffix("")
if hasattr(nb_path, "relative_to")
else nb_path
)
overrides = get_overrides(str(rel)) or {}
if overrides.get("gpu"):
item.add_marker(pytest.mark.gpu)
if overrides.get("long_running"):
item.add_marker(pytest.mark.long_running)
tier = get_tier(overrides)
if tier == TIER_WEEKLY:
item.add_marker(pytest.mark.weekly)
elif tier == TIER_ON_DEMAND:
item.add_marker(pytest.mark.on_demand)
reruns = get_reruns(overrides)
if reruns > 0 and has_rerunfailures:
item.add_marker(pytest.mark.flaky(reruns=reruns, reruns_delay=30))
# ---------------------------------------------------------------------------
# Incremental result saving — write JSONL after each test so results survive
# process kills. Results file: /tmp/ml4t-test-results.jsonl
# ---------------------------------------------------------------------------
_RESULTS_PATH = Path(os.environ.get("ML4T_RESULTS_FILE", "/tmp/ml4t-test-results.jsonl"))
_test_start_times: dict[str, float] = {}
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item):
"""Record test start time."""
_test_start_times[item.nodeid] = time.time()
@pytest.hookimpl(trylast=True)
def pytest_runtest_logreport(report):
"""Write each test result to JSONL as it completes."""
if report.when != "call" and not (report.when == "setup" and report.skipped):
return
start = _test_start_times.pop(report.nodeid, 0)
duration = report.duration if hasattr(report, "duration") else 0
outcome = report.outcome # "passed", "failed", or "skipped"
record = {
"nodeid": report.nodeid,
"outcome": outcome,
"duration_s": round(duration, 2),
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
}
if outcome == "failed" and report.longreprtext:
record["error"] = report.longreprtext[:500]
with open(_RESULTS_PATH, "a") as f:
f.write(json.dumps(record) + "\n")
f.flush()
@pytest.fixture
def clean_env():
"""Fixture that provides a clean environment and restores it after."""
saved_env = os.environ.copy()
yield os.environ
os.environ.clear()
os.environ.update(saved_env)
tests/fixtures/README.md
Vendored
+10
View File
@@ -0,0 +1,10 @@
# Test Fixtures
Minimal results JSON files seeded into test output directories so downstream
notebooks (latent factors, DL, backtest, etc.) can find baseline results
without depending on upstream notebook execution order.
These are NOT real results — they contain plausible placeholder values that
allow the comparison/selection code paths to execute without crashing.
Seeded by `conftest.py::seeded_output_dir` fixture.
tests/fixtures/__init__.py
Vendored
View File
tests/fixtures/seed_results.py
Vendored
+1183
View File
File diff suppressed because it is too large Load Diff
tests/generate_intermediates.py
+327
View File
@@ -0,0 +1,327 @@
#!/usr/bin/env python3
"""Generate pipeline intermediates for the test-data repo.
Runs all 9 case study pipelines through specified stages
via Papermill with test overrides, capturing outputs to the specified directory.
The outputs are committed to ml4t/third-edition-test-data so that downstream
chapters (Ch11+) can read pre-computed labels/features/predictions without
re-running the full pipeline.
Usage:
cd ~/ml4t/third_edition/code
ML4T_DATA_PATH=~/ml4t/test-data/data \
uv run python tests/generate_intermediates.py \
--output ~/ml4t/test-data/intermediates
# Run only through features (stages 01-03)
uv run python tests/generate_intermediates.py \
--output ~/ml4t/test-data/intermediates \
--through-stage 3
# Include DL stages (slow)
uv run python tests/generate_intermediates.py \
--output ~/ml4t/test-data/intermediates \
--through-stage 12 --no-skip-dl
"""
import argparse
import json
import os
import re
import shutil
import time
from datetime import UTC, datetime
from pathlib import Path
import yaml
try:
from tests.pm_helpers import get_overrides, run_notebook
except ModuleNotFoundError:
from pm_helpers import get_overrides, run_notebook
REPO_ROOT = Path(__file__).parent.parent
CASE_STUDIES = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
# Stage patterns to skip when --skip-dl is active (DL/latent/causal are heavy)
DL_STAGE_PATTERNS = re.compile(
r"\d{2}_("
r"dl_|deep_learning|tabular_dl|latent_factors|pca|ipca|cae|sdf|sae|"
r"autoencoder|term_structure_pca|causal_dml"
r")"
)
# ---------------------------------------------------------------------------
# Config seeding — replicate conftest.py seeded_output_dir logic
# ---------------------------------------------------------------------------
# Per-model-type overrides applied to copied preset YAMLs.
# Goal: minimal workload that still exercises the training loop + registry.
_TEST_PRESET_PATCHES: dict[str, dict] = {
"lgb": {"max_iterations": 2, "checkpoint_interval": 1},
"lstm": {"n_epochs": 2, "checkpoint_interval": 1},
"tsmixer": {"n_epochs": 2, "checkpoint_interval": 1},
"tcn": {"n_epochs": 2, "checkpoint_interval": 1},
"nlinear": {"n_epochs": 2, "checkpoint_interval": 1},
"patchtst": {"n_epochs": 2, "checkpoint_interval": 1},
"tabm": {"n_epochs": 2, "checkpoint_interval": 1},
"cae": {"n_epochs": 2, "checkpoint_interval": 1},
"sdf": {"n_epochs": 2, "checkpoint_interval": 1},
"sae": {"n_epochs": 2, "checkpoint_interval": 1},
"ipca": {"n_epochs": 2, "checkpoint_interval": 1},
}
_MAX_CONFIGS_PER_FAMILY = 2
_TRIM_FAMILIES = {"linear", "gbm"}
def _patch_presets_for_testing(config_dir: Path) -> None:
"""Patch copied preset YAMLs with reduced-workload values for testing."""
for model_type, overrides in _TEST_PRESET_PATCHES.items():
model_dir = config_dir / model_type
if not model_dir.exists():
continue
for preset_path in model_dir.glob("*.yaml"):
preset = yaml.safe_load(preset_path.read_text())
if preset is None:
continue
preset.update(overrides)
with open(preset_path, "w") as f:
yaml.dump(preset, f, default_flow_style=False)
def _trim_label_configs(cs_config_dir: Path) -> None:
"""Trim label config YAMLs to at most _MAX_CONFIGS_PER_FAMILY for sweep families."""
for label_yaml in cs_config_dir.glob("fwd_*.yaml"):
data = yaml.safe_load(label_yaml.read_text())
if data is None or not isinstance(data, dict):
continue
trimmed = False
for family, configs in data.items():
if (
family in _TRIM_FAMILIES
and isinstance(configs, list)
and len(configs) > _MAX_CONFIGS_PER_FAMILY
):
data[family] = configs[:_MAX_CONFIGS_PER_FAMILY]
trimmed = True
if trimmed:
with open(label_yaml, "w") as f:
yaml.dump(data, f, default_flow_style=False)
def seed_configs(output_dir: Path) -> None:
"""Copy case study configs and global model presets into output_dir.
Replicates the logic of conftest.py's seeded_output_dir fixture so that
notebooks executed via generate_intermediates.py find patched configs.
"""
cs_root = REPO_ROOT / "case_studies"
# Copy per-case-study config files (setup.yaml, training menus, backtest presets, etc.)
for cs_id in CASE_STUDIES:
src_config_dir = cs_root / cs_id / "config"
if not src_config_dir.exists():
continue
dst_config_dir = output_dir / cs_id / "config"
if dst_config_dir.exists():
shutil.rmtree(dst_config_dir)
shutil.copytree(src_config_dir, dst_config_dir)
_trim_label_configs(dst_config_dir)
# Copy global model presets so load_configs() can find them.
# load_configs() resolves presets at {case_dir.parent}/config/{model_type}/*.yaml
global_config_src = cs_root / "config"
global_config_dst = output_dir / "config"
if global_config_src.exists() and not global_config_dst.exists():
shutil.copytree(global_config_src, global_config_dst)
_patch_presets_for_testing(global_config_dst)
print(f"Seeded configs into {output_dir}")
def discover_stages(cs_dir: Path, through_stage: int, skip_dl: bool) -> list[Path]:
"""Auto-discover pipeline stages in a case study directory.
Returns sorted list of .py notebook paths up through the specified stage number.
Skips DL/latent/causal stages when skip_dl is True.
"""
stages = []
for notebook in sorted(cs_dir.glob("[0-9][0-9]_*.py")):
if notebook.name.startswith("_"):
continue
stage_num = int(notebook.stem[:2])
if stage_num > through_stage:
continue
if skip_dl and DL_STAGE_PATTERNS.match(notebook.stem):
continue
stages.append(notebook)
return stages
def main():
parser = argparse.ArgumentParser(description="Generate pipeline intermediates for CI")
parser.add_argument(
"--output",
type=Path,
required=True,
help="Output directory for intermediates",
)
parser.add_argument(
"--case-studies",
nargs="+",
default=CASE_STUDIES,
help="Case studies to run (default: all)",
)
parser.add_argument(
"--through-stage",
type=int,
default=8,
help="Run stages up to this number (default: 8 = through GBM for all case studies including sp500_options/08_gbm)",
)
parser.add_argument(
"--skip-dl",
action="store_true",
default=True,
help="Skip DL/latent/causal stages (default: True)",
)
parser.add_argument(
"--no-skip-dl",
action="store_false",
dest="skip_dl",
help="Include DL/latent/causal stages",
)
args = parser.parse_args()
output_dir = args.output.expanduser().resolve()
output_dir.mkdir(parents=True, exist_ok=True)
# Seed configs (setup.yaml, label configs, model presets) into output dir
# so notebooks find patched configs when ML4T_OUTPUT_DIR is set.
seed_configs(output_dir)
# Set ML4T_OUTPUT_DIR so all pipeline writes go to our output directory
os.environ["ML4T_OUTPUT_DIR"] = str(output_dir)
os.environ["MPLBACKEND"] = "Agg"
os.environ["PLOTLY_RENDERER"] = "json"
results = {}
total_start = time.time()
for cs in args.case_studies:
cs_dir = REPO_ROOT / "case_studies" / cs
if not cs_dir.exists():
print(f"\nSKIP {cs}: directory not found")
continue
stages = discover_stages(cs_dir, args.through_stage, args.skip_dl)
if not stages:
print(f"\nSKIP {cs}: no stages found")
continue
print(f"\n{'=' * 60}")
print(f"Case study: {cs} ({len(stages)} stages)")
print(f"{'=' * 60}")
cs_failed = False
for notebook in stages:
stage = notebook.stem
if cs_failed:
print(f" {stage}: SKIP (earlier stage failed)")
results[f"{cs}::{stage}"] = "skipped"
continue
rel_path = notebook.relative_to(REPO_ROOT).with_suffix("")
overrides = get_overrides(str(rel_path))
# Skip if overrides say so
if overrides.get("skip"):
reason = overrides.get("skip_reason", "marked skip")
print(f" {stage}: SKIP ({reason})")
results[f"{cs}::{stage}"] = "skipped"
# Pipeline stages (01-05) cascade their skip
stage_num = int(stage[:2])
if stage_num <= 5:
cs_failed = True
continue
timeout = overrides.get("timeout", 300)
parameters = overrides.get("parameters", {})
print(f" {stage}: running...", end="", flush=True)
start = time.time()
result = run_notebook(
py_path=notebook,
parameters=parameters,
timeout=timeout,
output_dir=output_dir,
)
elapsed = time.time() - start
if result["status"] == "ok":
print(f" OK ({elapsed:.0f}s)")
results[f"{cs}::{stage}"] = "ok"
else:
print(f" FAILED ({elapsed:.0f}s)")
print(f" Error: {result['error']}")
results[f"{cs}::{stage}"] = "failed"
cs_failed = True
total_elapsed = time.time() - total_start
# Summary
print(f"\n{'=' * 60}")
print(f"Summary ({total_elapsed:.0f}s total)")
print(f"{'=' * 60}")
ok = sum(1 for v in results.values() if v == "ok")
failed = sum(1 for v in results.values() if v == "failed")
skipped = sum(1 for v in results.values() if v == "skipped")
print(f" OK: {ok} Failed: {failed} Skipped: {skipped}")
if failed:
print("\nFailed stages:")
for k, v in results.items():
if v == "failed":
print(f" - {k}")
# Show output size
total_bytes = sum(f.stat().st_size for f in output_dir.rglob("*") if f.is_file())
print(f"\nOutput: {output_dir} ({total_bytes / 1e6:.1f} MB)")
# Write metadata for staleness tracking
metadata = {
"generated_at": datetime.now(UTC).isoformat(),
"through_stage": args.through_stage,
"skip_dl": args.skip_dl,
"results": results,
"total_seconds": round(total_elapsed),
"size_mb": round(total_bytes / 1e6, 1),
}
metadata_path = output_dir / "_metadata.json"
with open(metadata_path, "w") as f:
json.dump(metadata, f, indent=2)
print(f"Metadata: {metadata_path}")
if __name__ == "__main__":
main()
tests/generate_skip_data.py
+392
View File
@@ -0,0 +1,392 @@
"""Generate synthetic test data for currently-skipped notebooks.
Run once to enrich the test-data repo with minimal synthetic datasets
that allow the remaining skipped notebooks to execute their code paths.
Usage:
uv run python tests/generate_skip_data.py --output ~/ml4t/test-data
This generates data for:
1. FNSPID news dataset (Ch10/07, Ch10/08)
2. SEC 10-Q MD&A text (Ch10/09)
3. ADV columns for Kyle lambda (Ch18/03)
4. Engine divergence predictions (Ch16/07)
5. Signal quality synthesis data (Ch20/02)
6. MLOps drift detection features (Ch26/02)
7. MLOps safe model rollout (Ch26/03)
8. MLOps MLflow registry (Ch26/06)
"""
import argparse
import json
import sqlite3
from datetime import date, timedelta
from pathlib import Path
import numpy as np
import polars as pl
np.random.seed(42)
SYMBOLS_ETF = ["SPY", "QQQ", "IWM", "TLT", "GLD", "XLF", "XLK", "XLE", "EFA", "VWO"]
SYMBOLS_EQ = ["AAPL", "MSFT", "GOOGL", "AMZN", "NVDA", "META", "TSLA", "JPM", "V", "JNJ"]
def generate_fnspid_news(data_dir: Path):
"""Generate synthetic FNSPID financial news data."""
out = data_dir / "alternative" / "news" / "fnspid"
out.mkdir(parents=True, exist_ok=True)
headlines = [
"{sym} reports strong quarterly earnings, beats estimates",
"{sym} shares drop on weaker-than-expected revenue guidance",
"{sym} announces major acquisition worth $2.5B",
"Analysts upgrade {sym} citing improving margins",
"{sym} CEO discusses expansion plans in earnings call",
"Market volatility hits {sym} as sector rotates",
"{sym} launches new product line targeting enterprise customers",
"Institutional investors increase {sym} holdings in Q3",
"{sym} faces regulatory scrutiny over data practices",
"{sym} dividend increase signals management confidence",
]
rows = []
dates = pl.date_range(date(2022, 1, 3), date(2024, 12, 31), "1d", eager=True)
for d in dates:
# 2-5 news items per day
n_items = np.random.randint(2, 6)
for _ in range(n_items):
sym = np.random.choice(SYMBOLS_EQ)
headline = np.random.choice(headlines).format(sym=sym)
rows.append(
{
"ticker": sym,
"timestamp": d,
"title": headline,
"source": np.random.choice(["Reuters", "Bloomberg", "CNBC", "WSJ"]),
}
)
df = pl.DataFrame(rows)
df.write_parquet(out / "fnspid_sample.parquet")
print(f" FNSPID: {len(df)} news items -> {out / 'fnspid_sample.parquet'}")
def generate_sec_10q_mda(data_dir: Path):
"""Generate synthetic SEC 10-Q MD&A text data."""
out = data_dir / "alternative" / "text"
out.mkdir(parents=True, exist_ok=True)
rows = []
for sym in SYMBOLS_EQ[:6]:
for year in range(2019, 2024):
for quarter in range(1, 5):
month = quarter * 3 + 1
if month > 12:
month = 1
year_f = year + 1
else:
year_f = year
filing_date = date(year_f, min(month, 12), 15)
period_end = date(year, quarter * 3, 28)
mda_text = (
f"Management's Discussion and Analysis for {sym}. "
f"During Q{quarter} {year}, revenue increased by {np.random.uniform(2, 15):.1f}% "
f"year-over-year. Operating margins improved to {np.random.uniform(15, 35):.1f}%. "
f"We continue to invest in R&D and expect continued growth. "
f"Key risks include market volatility and regulatory changes."
)
rows.append(
{
"symbol": sym,
"cik": str(np.random.randint(100000, 999999)),
"accession_no": f"0001234567-{year_f:04d}-{np.random.randint(10000, 99999):05d}",
"filing_date": filing_date,
"period_end": period_end,
"mda_text": mda_text,
"mda_word_count": len(mda_text.split()),
"mda_char_count": len(mda_text),
}
)
df = pl.DataFrame(rows)
df.write_parquet(out / "sp500_10q_mda.parquet")
print(f" SEC 10-Q: {len(df)} filings -> {out / 'sp500_10q_mda.parquet'}")
def enrich_adv_columns(data_dir: Path):
"""Add adv_21d (21-day average daily volume) to datasets that need it.
The Kyle lambda market impact calibration notebook (Ch18/03) reads
adv_21d from equity price data. The test data doesn't have this computed.
"""
datasets = [
("etfs", "etf_universe.parquet"),
("equities", "us_equities.parquet"),
]
for subdir, filename in datasets:
path = data_dir / subdir / filename
if not path.exists():
print(f" ADV: SKIP {path} (not found)")
continue
df = pl.read_parquet(path)
if "adv_21d" in df.columns:
print(f" ADV: SKIP {path} (already has adv_21d)")
continue
if "volume" not in df.columns:
print(f" ADV: SKIP {path} (no volume column)")
continue
# Compute rolling 21-day average volume per symbol
sort_cols = ["symbol", "timestamp"] if "symbol" in df.columns else ["timestamp"]
group_col = "symbol" if "symbol" in df.columns else None
if group_col:
df = df.sort(sort_cols).with_columns(
pl.col("volume")
.rolling_mean(window_size=21, min_samples=1)
.over(group_col)
.alias("adv_21d")
)
else:
df = df.sort("timestamp").with_columns(
pl.col("volume").rolling_mean(window_size=21, min_samples=1).alias("adv_21d")
)
df.write_parquet(path)
print(f" ADV: Added adv_21d to {path} ({len(df)} rows)")
def generate_engine_divergence_predictions(intermediates_dir: Path):
"""Generate predictions with model column for Ch16/07 engine divergence."""
out = intermediates_dir / "ch16_signal_method_comparison"
out.mkdir(parents=True, exist_ok=True)
dates = pl.date_range(date(2022, 1, 3), date(2023, 12, 29), "1d", eager=True)
rows = []
for d in dates:
for sym in SYMBOLS_ETF[:5]:
rows.append(
{
"timestamp": d,
"symbol": sym,
"prediction": np.random.normal(0, 0.02),
"model": "ridge_a1.0",
}
)
df = pl.DataFrame(rows)
df.write_parquet(out / "predictions_with_model.parquet")
print(f" Engine divergence: {len(df)} rows -> {out}")
def generate_signal_quality_data(intermediates_dir: Path):
"""Generate synthesis data for Ch20/02 signal quality notebook."""
# The notebook reads from Ch20/01 aggregate_synthesis outputs
out = intermediates_dir / "ch20_synthesis"
out.mkdir(parents=True, exist_ok=True)
case_studies = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
models = ["linear/ridge", "gbm/leaves_15", "deep_learning/lstm", "tabular_dl/tabm_l"]
# IC comparison data
ic_rows = []
for cs in case_studies:
for model in models:
ic_rows.append(
{
"case_study": cs,
"source": model,
"ic_mean": np.random.uniform(-0.02, 0.06),
"ic_std": np.random.uniform(0.01, 0.04),
"n_folds": 5,
}
)
ic_df = pl.DataFrame(ic_rows)
ic_df.write_parquet(out / "ic_comparison.parquet")
# Synthesis JSON
synthesis = {
"case_studies": {
cs: {
"champion": {
"source": "gbm/leaves_15",
"sharpe": float(np.random.uniform(-0.5, 2.0)),
},
"holdout": {
"ic": float(np.random.uniform(-0.02, 0.1)),
"sharpe": float(np.random.uniform(-1, 3)),
},
}
for cs in case_studies
}
}
(out / "all_synthesis.json").write_text(json.dumps(synthesis, indent=2))
print(f" Signal quality: IC comparison + synthesis -> {out}")
def generate_mlops_data(intermediates_dir: Path, data_dir: Path):
"""Generate data for Ch26 MLOps notebooks (02, 03, 06)."""
# Ch26/02 needs ETFs features with adv_21d — handled by enrich_adv_columns
# Ch26/03 needs a linear/lasso validation run in registry
out = intermediates_dir / "us_equities_panel" / "run_log"
out.mkdir(parents=True, exist_ok=True)
db_path = out / "registry.db"
db = sqlite3.connect(str(db_path))
db.execute("""
CREATE TABLE IF NOT EXISTS training_runs (
run_id TEXT PRIMARY KEY,
entry_point TEXT,
source TEXT,
label TEXT,
config_hash TEXT,
created_at TEXT,
ic_mean REAL,
status TEXT DEFAULT 'completed'
)
""")
db.execute("""
CREATE TABLE IF NOT EXISTS prediction_sets (
pred_id TEXT PRIMARY KEY,
run_id TEXT,
entry_point TEXT,
source TEXT,
label TEXT,
config_hash TEXT,
created_at TEXT,
ic_mean REAL,
n_rows INTEGER,
pred_path TEXT
)
""")
db.execute("""
CREATE TABLE IF NOT EXISTS prediction_metrics (
metric_id INTEGER PRIMARY KEY AUTOINCREMENT,
pred_id TEXT,
fold INTEGER,
ic REAL,
n_rows INTEGER
)
""")
# Insert a few synthetic runs
for i, (source, ic) in enumerate(
[
("linear/ridge_a1.0", 0.025),
("linear/lasso_a0.01", 0.018),
("gbm/leaves_15_mae", 0.042),
]
):
run_id = f"run_{i:03d}"
pred_id = f"pred_{i:03d}"
db.execute(
"INSERT OR REPLACE INTO training_runs VALUES (?,?,?,?,?,?,?,?)",
(
run_id,
"06_linear" if "linear" in source else "07_gbm",
source,
"fwd_ret_1d",
f"hash_{i}",
"2026-01-01T00:00:00",
ic,
"completed",
),
)
db.execute(
"INSERT OR REPLACE INTO prediction_sets VALUES (?,?,?,?,?,?,?,?,?,?)",
(
pred_id,
run_id,
"06_linear" if "linear" in source else "07_gbm",
source,
"fwd_ret_1d",
f"hash_{i}",
"2026-01-01T00:00:00",
ic,
1000,
f"predictions/{pred_id}.parquet",
),
)
for fold in range(5):
db.execute(
"INSERT INTO prediction_metrics (pred_id, fold, ic, n_rows) VALUES (?,?,?,?)",
(pred_id, fold, ic + np.random.normal(0, 0.005), 200),
)
db.commit()
db.close()
print(f" MLOps registry: 3 runs -> {db_path}")
# Generate stub predictions for the registry entries
preds_dir = out.parent / "predictions"
preds_dir.mkdir(parents=True, exist_ok=True)
dates = pl.date_range(date(2023, 1, 2), date(2023, 12, 29), "1d", eager=True)
for i in range(3):
rows = []
for d in dates:
for sym in SYMBOLS_EQ[:5]:
rows.append(
{
"timestamp": d,
"symbol": sym,
"prediction": np.random.normal(0, 0.02),
"fold": np.random.randint(0, 5),
}
)
df = pl.DataFrame(rows)
df.write_parquet(preds_dir / f"pred_{i:03d}.parquet")
print(f" MLOps predictions: 3 files -> {preds_dir}")
def main():
parser = argparse.ArgumentParser(
description="Generate synthetic test data for skipped notebooks"
)
parser.add_argument("--output", required=True, help="Test data repo root")
args = parser.parse_args()
root = Path(args.output)
data_dir = root / "data"
intermediates_dir = root / "intermediates"
print("Generating synthetic test data for skipped notebooks...")
print()
print("[1/6] FNSPID news data (Ch10/07, Ch10/08)...")
generate_fnspid_news(data_dir)
print("[2/6] SEC 10-Q MD&A text (Ch10/09)...")
generate_sec_10q_mda(data_dir)
print("[3/6] ADV columns for Kyle lambda (Ch18/03)...")
enrich_adv_columns(data_dir)
print("[4/6] Engine divergence predictions (Ch16/07)...")
generate_engine_divergence_predictions(intermediates_dir)
print("[5/6] Signal quality synthesis data (Ch20/02)...")
generate_signal_quality_data(intermediates_dir)
print("[6/6] MLOps registry and predictions (Ch26/02-06)...")
generate_mlops_data(intermediates_dir, data_dir)
print()
print("Done! Now commit changes to the test-data repo and update overrides.yaml.")
if __name__ == "__main__":
main()
tests/generate_test_data.sh
Executable
+59
View File
@@ -0,0 +1,59 @@
#!/bin/bash
# Regenerate all test data (raw + intermediates) for CI.
#
# Run from the review repo root. Requires:
# - ML4T_DATA_PATH to point to full production data (default: ~/Dropbox/ml4t/data)
# - The working repo at ~/ml4t/third-edition
# - The review repo at ~/ml4t/technical_review
#
# Usage:
# cd ~/ml4t/technical_review
# bash tests/generate_test_data.sh [TEST_DATA_DIR]
set -euo pipefail
TEST_DATA_DIR="${1:-$HOME/ml4t/test-data}"
WORKING_REPO="$HOME/ml4t/third-edition"
REVIEW_REPO="$HOME/ml4t/technical_review"
echo "=== ML4T Test Data Generator ==="
echo "Output: $TEST_DATA_DIR"
echo "Working repo: $WORKING_REPO"
echo "Review repo: $REVIEW_REPO"
echo ""
# Step 1: Generate subsampled raw data
echo "=== Step 1: Generating subsampled data ==="
cd "$WORKING_REPO"
uv run python tests/create_test_data.py \
--source "${ML4T_DATA_PATH:-$HOME/Dropbox/ml4t/data}" \
--output "$TEST_DATA_DIR/data" \
--clean
echo ""
# Step 2: Deploy latest notebooks from third-edition -> review repo
echo "=== Step 2: Deploying latest notebooks ==="
cd "$WORKING_REPO"
uv run python scripts/deploy_to_review.py --all
echo ""
# Step 3: Generate intermediates (runs pipeline notebooks via Papermill)
echo "=== Step 3: Generating intermediates ==="
cd "$REVIEW_REPO"
ML4T_DATA_PATH="$TEST_DATA_DIR/data" \
MPLBACKEND=Agg \
PLOTLY_RENDERER=json \
uv run python tests/generate_intermediates.py \
--output "$TEST_DATA_DIR/intermediates"
echo ""
# Summary
echo "=== Done ==="
echo "Test data directory: $TEST_DATA_DIR"
du -sh "$TEST_DATA_DIR/data" "$TEST_DATA_DIR/intermediates" 2>/dev/null || true
echo ""
echo "Next steps:"
echo " cd $TEST_DATA_DIR"
echo " git add -A"
echo " git commit -m 'update: regenerate test data'"
echo " git push"
tests/generate_test_microstructure.py
+865
View File
@@ -0,0 +1,865 @@
"""Generate minimal synthetic test data for Ch03 microstructure notebooks.
Creates test-sized datasets that match the schemas expected by Ch03 notebooks
and related notebooks (Ch02 futures individual, Ch04 prediction markets).
Writes to ~/ml4t/test-data/data/ which serves as ML4T_DATA_PATH
in CI.
Usage:
uv run python tests/generate_test_microstructure.py
"""
from datetime import date, datetime, time, timedelta
from pathlib import Path
import numpy as np
import polars as pl
# ── Output root ──────────────────────────────────────────────────────────────
TEST_DATA_ROOT = Path.home() / "ml4t" / "test-data" / "data"
# Seed for reproducibility
RNG = np.random.default_rng(42)
# ═════════════════════════════════════════════════════════════════════════════
# 1. ITCH Parsed Messages (for NB 02-10)
# ═════════════════════════════════════════════════════════════════════════════
# These go into the ITCH messages path that load_nasdaq_itch() resolves:
# ML4T_DATA_PATH / "equities" / "market" / "microstructure" / "nasdaq_itch" / "messages"
# Notebooks 02-10 read via utils.limit_orderbook.load_itch_messages(itch_dir, msg_type, symbol)
def _ns_timestamp(hour: int, minute: int, second: int = 0, micro: int = 0) -> datetime:
"""Create a nanosecond-precision datetime on the ITCH trading day (2020-01-30)."""
return datetime(2020, 1, 30, hour, minute, second, micro)
def generate_itch_messages() -> None:
"""Generate all ITCH message type parquet files."""
itch_dir = (
TEST_DATA_ROOT / "equities" / "market" / "microstructure" / "nasdaq_itch" / "messages"
)
# ── R (Stock Directory) ──────────────────────────────────────────────
r_dir = itch_dir / "R"
r_dir.mkdir(parents=True, exist_ok=True)
r_df = pl.DataFrame(
{
"stock_locate": pl.Series([1, 2, 3], dtype=pl.UInt16),
"tracking_number": pl.Series([0, 0, 0], dtype=pl.UInt16),
"timestamp": [
_ns_timestamp(4, 0, 0),
_ns_timestamp(4, 0, 0),
_ns_timestamp(4, 0, 0),
],
"stock": ["AAPL", "MSFT", "NVDA"],
"market_category": ["Q", "Q", "Q"],
"financial_status": ["N", "N", "N"],
"round_lot_size": pl.Series([100, 100, 100], dtype=pl.UInt32),
"round_lots_only": ["N", "N", "N"],
"issue_classification": ["C", "C", "C"],
"issue_subtype": ["Z", "Z", "Z"],
"authenticity": ["P", "P", "P"],
"short_sale_threshold": ["N", "N", "N"],
"ipo_flag": ["N", "N", "N"],
"luld_reference_price_tier": ["1", "1", "1"],
"etp_flag": ["N", "N", "N"],
"etp_leverage_factor": pl.Series([0, 0, 0], dtype=pl.UInt32),
"inverse_indicator": ["N", "N", "N"],
}
).cast({"timestamp": pl.Datetime("ns")})
r_df.write_parquet(r_dir / "part-000000.parquet")
# ── S (System Event) ─────────────────────────────────────────────────
s_dir = itch_dir / "S"
s_dir.mkdir(parents=True, exist_ok=True)
s_df = pl.DataFrame(
{
"stock_locate": pl.Series([0, 0, 0, 0], dtype=pl.UInt16),
"tracking_number": pl.Series([0, 0, 0, 0], dtype=pl.UInt16),
"timestamp": [
_ns_timestamp(4, 0, 0),
_ns_timestamp(9, 30, 0),
_ns_timestamp(16, 0, 0),
_ns_timestamp(20, 0, 0),
],
"event_code": ["O", "Q", "M", "C"],
}
).cast({"timestamp": pl.Datetime("ns")})
s_df.write_parquet(s_dir / "part-000000.parquet")
# ── A (Add Order) ────────────────────────────────────────────────────
# 20 orders for AAPL (stock_locate=1), spanning 10:00 to 15:00
a_dir = itch_dir / "A"
a_dir.mkdir(parents=True, exist_ok=True)
n_orders = 20
base_price_aapl = 320.0 # AAPL price circa Jan 2020
order_refs = list(range(1001, 1001 + n_orders))
sides = ["B" if i % 2 == 0 else "S" for i in range(n_orders)]
shares = [int(RNG.integers(100, 1001)) for _ in range(n_orders)]
# Prices: bids slightly below base, asks slightly above
prices = []
for i, side in enumerate(sides):
offset = RNG.uniform(0.01, 0.50)
if side == "B":
prices.append(round(base_price_aapl - offset, 4))
else:
prices.append(round(base_price_aapl + offset, 4))
# Timestamps spaced across 10:00-15:00 (300 minutes = 18000 seconds)
a_timestamps = [
_ns_timestamp(10, 0) + timedelta(seconds=int(i * 18000 / n_orders)) for i in range(n_orders)
]
# Prices stored as ITCH price4 integers (multiply by 10000) per spec
a_df = pl.DataFrame(
{
"stock_locate": pl.Series([1] * n_orders, dtype=pl.UInt16),
"tracking_number": pl.Series([0] * n_orders, dtype=pl.UInt16),
"timestamp": a_timestamps,
"order_reference_number": pl.Series(order_refs, dtype=pl.UInt64),
"buy_sell_indicator": sides,
"shares": pl.Series(shares, dtype=pl.UInt32),
"stock": ["AAPL"] * n_orders,
"price": pl.Series([int(p * 10000) for p in prices], dtype=pl.UInt32),
}
).cast({"timestamp": pl.Datetime("ns")})
a_df.write_parquet(a_dir / "part-000000.parquet")
# ── D (Order Delete) ─────────────────────────────────────────────────
d_dir = itch_dir / "D"
d_dir.mkdir(parents=True, exist_ok=True)
delete_refs = [1001, 1003, 1005, 1007, 1009]
d_df = pl.DataFrame(
{
"stock_locate": pl.Series([1] * 5, dtype=pl.UInt16),
"tracking_number": pl.Series([0] * 5, dtype=pl.UInt16),
"timestamp": [
a_timestamps[0] + timedelta(seconds=30),
a_timestamps[2] + timedelta(seconds=30),
a_timestamps[4] + timedelta(seconds=30),
a_timestamps[6] + timedelta(seconds=30),
a_timestamps[8] + timedelta(seconds=30),
],
"order_reference_number": pl.Series(delete_refs, dtype=pl.UInt64),
}
).cast({"timestamp": pl.Datetime("ns")})
d_df.write_parquet(d_dir / "part-000000.parquet")
# ── E (Order Executed) ───────────────────────────────────────────────
e_dir = itch_dir / "E"
e_dir.mkdir(parents=True, exist_ok=True)
exec_refs = [1002, 1004, 1006, 1008, 1010, 1012, 1014, 1016]
exec_shares = [min(shares[r - 1001] // 2, 200) for r in exec_refs]
e_df = pl.DataFrame(
{
"stock_locate": pl.Series([1] * 8, dtype=pl.UInt16),
"tracking_number": pl.Series([0] * 8, dtype=pl.UInt16),
"timestamp": [a_timestamps[r - 1001] + timedelta(seconds=60) for r in exec_refs],
"order_reference_number": pl.Series(exec_refs, dtype=pl.UInt64),
"executed_shares": pl.Series(exec_shares, dtype=pl.UInt32),
"match_number": pl.Series(list(range(5001, 5009)), dtype=pl.UInt64),
}
).cast({"timestamp": pl.Datetime("ns")})
e_df.write_parquet(e_dir / "part-000000.parquet")
# ── X (Order Cancel) ─────────────────────────────────────────────────
x_dir = itch_dir / "X"
x_dir.mkdir(parents=True, exist_ok=True)
cancel_refs = [1011, 1013, 1015]
cancel_shares = [shares[r - 1001] // 3 for r in cancel_refs]
x_df = pl.DataFrame(
{
"stock_locate": pl.Series([1] * 3, dtype=pl.UInt16),
"tracking_number": pl.Series([0] * 3, dtype=pl.UInt16),
"timestamp": [a_timestamps[r - 1001] + timedelta(seconds=45) for r in cancel_refs],
"order_reference_number": pl.Series(cancel_refs, dtype=pl.UInt64),
"cancelled_shares": pl.Series(cancel_shares, dtype=pl.UInt32),
}
).cast({"timestamp": pl.Datetime("ns")})
x_df.write_parquet(x_dir / "part-000000.parquet")
# ── C (Order Executed with Price) ────────────────────────────────────
c_dir = itch_dir / "C"
c_dir.mkdir(parents=True, exist_ok=True)
c_refs = [1017, 1018]
c_df = pl.DataFrame(
{
"stock_locate": pl.Series([1] * 2, dtype=pl.UInt16),
"tracking_number": pl.Series([0] * 2, dtype=pl.UInt16),
"timestamp": [
a_timestamps[16] + timedelta(seconds=90),
a_timestamps[17] + timedelta(seconds=90),
],
"order_reference_number": pl.Series(c_refs, dtype=pl.UInt64),
"executed_shares": pl.Series([shares[16] // 4, shares[17] // 4], dtype=pl.UInt32),
"match_number": pl.Series([6001, 6002], dtype=pl.UInt64),
"printable": ["Y", "Y"],
"execution_price": pl.Series(
[int(prices[16] * 10000), int(prices[17] * 10000)], dtype=pl.UInt32
),
}
).cast({"timestamp": pl.Datetime("ns")})
c_df.write_parquet(c_dir / "part-000000.parquet")
# ── P (Non-Cross Trade) ──────────────────────────────────────────────
p_dir = itch_dir / "P"
p_dir.mkdir(parents=True, exist_ok=True)
p_df = pl.DataFrame(
{
"stock_locate": pl.Series([1, 1, 1], dtype=pl.UInt16),
"tracking_number": pl.Series([0, 0, 0], dtype=pl.UInt16),
"timestamp": [
_ns_timestamp(11, 30, 0),
_ns_timestamp(13, 0, 0),
_ns_timestamp(14, 30, 0),
],
"order_reference_number": pl.Series([2001, 2002, 2003], dtype=pl.UInt64),
"buy_sell_indicator": ["B", "S", "B"],
"shares": pl.Series([200, 150, 300], dtype=pl.UInt32),
"stock": ["AAPL", "AAPL", "AAPL"],
"price": pl.Series([int(base_price_aapl * 10000)] * 3, dtype=pl.UInt32),
"match_number": pl.Series([7001, 7002, 7003], dtype=pl.UInt64),
}
).cast({"timestamp": pl.Datetime("ns")})
p_df.write_parquet(p_dir / "part-000000.parquet")
# ── U (Order Replace) ────────────────────────────────────────────────
u_dir = itch_dir / "U"
u_dir.mkdir(parents=True, exist_ok=True)
u_df = pl.DataFrame(
{
"stock_locate": pl.Series([1, 1], dtype=pl.UInt16),
"tracking_number": pl.Series([0, 0], dtype=pl.UInt16),
"timestamp": [
a_timestamps[18] + timedelta(seconds=20),
a_timestamps[19] + timedelta(seconds=20),
],
"original_order_reference_number": pl.Series([1019, 1020], dtype=pl.UInt64),
"new_order_reference_number": pl.Series([3001, 3002], dtype=pl.UInt64),
"shares": pl.Series([500, 600], dtype=pl.UInt32),
"price": pl.Series(
[int((base_price_aapl - 0.10) * 10000), int((base_price_aapl + 0.10) * 10000)],
dtype=pl.UInt32,
),
}
).cast({"timestamp": pl.Datetime("ns")})
u_df.write_parquet(u_dir / "part-000000.parquet")
print(f" ITCH messages written to {itch_dir}")
for sub in sorted(itch_dir.iterdir()):
if sub.is_dir() and sub.name != "enriched":
n = pl.scan_parquet(sub / "*.parquet").select(pl.len()).collect().item()
print(f" {sub.name}/: {n} rows")
# ═════════════════════════════════════════════════════════════════════════════
# 2. DataBento MBO (for NB 09-13)
# ═════════════════════════════════════════════════════════════════════════════
# Path: ML4T_DATA_PATH / "equities" / "market" / "microstructure" / "market_by_order" / "NVDA"
# File naming: xnas-itch-YYYYMMDD.mbo.dbn.parquet (DataBento convention)
# NB09, NB12 expect "timestamp" column. NB10, NB11, NB13 also need it.
# We provide BOTH ts_event and timestamp (same values) for compatibility.
def _generate_mbo_day(base_date: datetime, base_price_nano: int, start_order_id: int) -> list[dict]:
"""Generate one day of MBO messages with realistic bid/ask structure.
Returns a list of row dicts (not yet a DataFrame).
"""
rows: list[dict] = []
order_id = start_order_id
n_cycles = 50 # 50 cycles spread across 6.5 hours of trading
for cycle in range(n_cycles):
cycle_start_ms = cycle * 468_000 # ~7.8 min per cycle
# Phase 1: Adds (build book) - 15 orders per cycle
for i in range(15):
ts = base_date + timedelta(milliseconds=cycle_start_ms + i * 100)
side = "B" if i % 2 == 0 else "A"
if side == "B":
price_offset = -RNG.integers(1, 51) * 10_000_000
else:
price_offset = RNG.integers(1, 51) * 10_000_000
price = base_price_nano + price_offset
size = int(RNG.integers(1, 501))
rows.append(
{
"ts_event": ts,
"ts_recv": ts + timedelta(microseconds=int(RNG.integers(1, 100))),
"action": "A",
"side": side,
"price": price,
"size": size,
"order_id": order_id,
"flags": 0,
"publisher_id": 39,
}
)
order_id += 1
# Phase 2: Modifications - 3 per cycle
for i in range(3):
ts = base_date + timedelta(milliseconds=cycle_start_ms + 1500 + i * 200)
mod_order = order_id - 15 + i * 5
mod_side = "B" if i % 2 == 0 else "A"
if mod_side == "B":
price_offset = -RNG.integers(1, 31) * 10_000_000
else:
price_offset = RNG.integers(1, 31) * 10_000_000
rows.append(
{
"ts_event": ts,
"ts_recv": ts + timedelta(microseconds=int(RNG.integers(1, 100))),
"action": "M",
"side": mod_side,
"price": base_price_nano + price_offset,
"size": int(RNG.integers(1, 300)),
"order_id": mod_order,
"flags": 0,
"publisher_id": 39,
}
)
# Phase 3: Cancels - 3 per cycle
for i in range(3):
ts = base_date + timedelta(milliseconds=cycle_start_ms + 2100 + i * 200)
cancel_order = order_id - 14 + i * 5
rows.append(
{
"ts_event": ts,
"ts_recv": ts + timedelta(microseconds=int(RNG.integers(1, 100))),
"action": "C",
"side": "A" if i % 2 == 0 else "B",
"price": base_price_nano,
"size": 0,
"order_id": cancel_order,
"flags": 0,
"publisher_id": 39,
}
)
# Phase 4: Fills (F) and Trades (T) - 10 per cycle
for i in range(10):
ts = base_date + timedelta(milliseconds=cycle_start_ms + 2700 + i * 300)
fill_order = order_id - 13 + i
fill_size = int(RNG.integers(1, 200))
# Biased aggressor side for realistic imbalance runs.
# Runs of 10 consecutive cycles (~100 trades) with 95% bias,
# creating sustained imbalance that triggers bar boundaries.
# This mimics real institutional order flow patterns.
run_idx = cycle // 10
if run_idx % 2 == 0:
aggressor = "B" if RNG.random() < 0.95 else "A"
else:
aggressor = "A" if RNG.random() < 0.95 else "B"
trade_price = base_price_nano + RNG.integers(-5, 6) * 10_000_000
fill_side = "A" if aggressor == "B" else "B"
rows.append(
{
"ts_event": ts,
"ts_recv": ts + timedelta(microseconds=int(RNG.integers(1, 100))),
"action": "F",
"side": fill_side,
"price": trade_price,
"size": fill_size,
"order_id": fill_order,
"flags": 128,
"publisher_id": 39,
}
)
rows.append(
{
"ts_event": ts + timedelta(microseconds=1),
"ts_recv": ts + timedelta(microseconds=int(RNG.integers(2, 150))),
"action": "T",
"side": aggressor,
"price": trade_price,
"size": fill_size,
"order_id": fill_order,
"flags": 128,
"publisher_id": 39,
}
)
return rows
def generate_mbo_data() -> None:
"""Generate synthetic DataBento MBO tick data for NVDA.
Key schema requirements from notebooks:
- NB09 (lee_ready): expects "timestamp" column, reads parquet directly
- NB10 (information_bars): expects filename like xnas-itch-YYYYMMDD.mbo.dbn.parquet
- NB11 (lob_reconstruction): expects "ts_event" column, reads parquet directly
- NB12 (mbo_analysis): expects "timestamp" column, reads parquet directly
- NB13 (bar_sampling): expects "timestamp" column, filename like xnas-itch-*
We include both ts_event and timestamp columns, and use DataBento file naming.
We also generate enough data (spread across hours) for meaningful analysis.
"""
mbo_dir = TEST_DATA_ROOT / "equities" / "market" / "microstructure" / "market_by_order" / "NVDA"
mbo_dir.mkdir(parents=True, exist_ok=True)
# Remove old file if it exists (was named 20241104.parquet before)
old_file = mbo_dir / "20241104.parquet"
if old_file.exists():
old_file.unlink()
base_price_nano = 140_000_000_000 # $140 in nanodollars
# Generate 3 days of data. NB13 (bar_sampling) computes day-to-day CV which
# needs >= 2 days. NB10 (information_bars) also benefits from more trades.
trading_days = [
datetime(2024, 11, 4, 14, 30, 0), # Monday 9:30 AM ET in UTC
datetime(2024, 11, 5, 14, 30, 0), # Tuesday
datetime(2024, 11, 6, 14, 30, 0), # Wednesday
]
for day_idx, base_date in enumerate(trading_days):
rows = _generate_mbo_day(base_date, base_price_nano, 100_000 + day_idx * 10_000)
df = (
pl.DataFrame(rows)
.cast(
{
"ts_event": pl.Datetime("ns"),
"ts_recv": pl.Datetime("ns"),
"price": pl.Int64,
"size": pl.Int64,
"order_id": pl.Int64,
"flags": pl.Int64,
"publisher_id": pl.Int64,
}
)
.sort("ts_event")
)
# Add canonical "timestamp" column (same as ts_event) for notebooks that expect it.
# NB09, NB12, NB13 use "timestamp"; NB10, NB11 use "ts_event".
df = df.with_columns(pl.col("ts_event").alias("timestamp"))
# Write with DataBento filename convention: xnas-itch-YYYYMMDD.mbo.dbn.parquet
# NB10 and NB13 parse the filename: file_path.name.split("-")[2].split(".")[0]
date_str = base_date.strftime("%Y%m%d")
out_file = mbo_dir / f"xnas-itch-{date_str}.mbo.dbn.parquet"
df.write_parquet(out_file)
print(f" MBO day {date_str}: {len(df)} rows -> {out_file}")
# ═════════════════════════════════════════════════════════════════════════════
# 3. AlgoSeek TAQ (for NB 15-16)
# ═════════════════════════════════════════════════════════════════════════════
# Path: ML4T_DATA_PATH / "equities" / "market" / "microstructure" / "trade_and_quotes" / "symbol=AAPL" / "data.parquet"
# NB15 expects: TRADE, QUOTE BID, QUOTE ASK, QUOTE BID NB, QUOTE ASK NB event types
# NB15 does spread analysis using NBBO quotes and trade size distribution
def generate_taq_data() -> None:
"""Generate synthetic AlgoSeek TAQ tick data for AAPL on 2020-03-16.
Key schema requirements from notebooks:
- NB15 (taq_eda): Needs TRADE, QUOTE BID NB, QUOTE ASK NB event types
for spread analysis. Needs enough trades for size distribution.
- NB16 (taq_lob): Needs QUOTE BID/ASK for LOB reconstruction.
We generate ~600 events with realistic distributions.
"""
taq_dir = (
TEST_DATA_ROOT
/ "equities"
/ "market"
/ "microstructure"
/ "trade_and_quotes"
/ "symbol=AAPL"
)
taq_dir.mkdir(parents=True, exist_ok=True)
# March 16, 2020: AAPL around $250, huge volatility day
base_date = datetime(2020, 3, 16)
base_price = 250.0
exchanges = ["Q", "N", "Z", "P", "K"]
rows = []
# Generate ~600 rows: mix of trades, exchange quotes, and NBBO quotes
# NB15 needs: TRADE events for trade analysis, QUOTE BID NB / QUOTE ASK NB for spread
for i in range(600):
# Random time between 9:30 and 16:00 (6.5 hours = 23400 seconds)
seconds_offset = int(RNG.integers(0, 23400))
ts = base_date + timedelta(
hours=9, minutes=30, seconds=seconds_offset, microseconds=int(RNG.integers(0, 999999))
)
# Event type distribution:
# ~15% trades, ~20% NBBO bids, ~20% NBBO asks, ~20% exchange bids, ~20% exchange asks
# We need QUOTE BID NB and QUOTE ASK NB for NB15's spread analysis
r = RNG.random()
if r < 0.15:
event_type = "TRADE"
price = round(base_price + RNG.normal(0, 5), 2)
quantity = int(RNG.integers(10, 10001))
elif r < 0.35:
event_type = "QUOTE BID NB"
price = round(base_price - abs(RNG.normal(0.03, 0.10)), 2)
quantity = int(RNG.integers(100, 5001))
elif r < 0.55:
event_type = "QUOTE ASK NB"
price = round(base_price + abs(RNG.normal(0.03, 0.10)), 2)
quantity = int(RNG.integers(100, 5001))
elif r < 0.75:
event_type = "QUOTE BID"
price = round(base_price - abs(RNG.normal(0.05, 0.20)), 2)
quantity = int(RNG.integers(100, 5001))
else:
event_type = "QUOTE ASK"
price = round(base_price + abs(RNG.normal(0.05, 0.20)), 2)
quantity = int(RNG.integers(100, 5001))
rows.append(
{
"timestamp": ts,
"event_type": event_type,
"price": price,
"quantity": quantity,
"exchange": exchanges[int(RNG.integers(0, len(exchanges)))],
"conditions": "00000000",
}
)
df = (
pl.DataFrame(rows)
.cast(
{
"timestamp": pl.Datetime("us"),
"price": pl.Float64,
"quantity": pl.Int64,
}
)
.sort("timestamp")
)
df.write_parquet(taq_dir / "data.parquet")
print(f" TAQ data: {len(df)} rows -> {taq_dir / 'data.parquet'}")
# ═════════════════════════════════════════════════════════════════════════════
# 4. IEX Parsed Data (for NB 14)
# ═════════════════════════════════════════════════════════════════════════════
# Path: ML4T_DATA_PATH / "equities" / "market" / "microstructure" / "iex" / "deep" / "parsed" / {type}/
def generate_iex_data() -> None:
"""Generate synthetic IEX DEEP parsed data."""
parsed_dir = (
TEST_DATA_ROOT / "equities" / "market" / "microstructure" / "iex" / "deep" / "parsed"
)
base_date = datetime(2025, 1, 15, 14, 30, 0) # 9:30 AM ET in UTC
base_price = 240.0 # AAPL-ish
# ── Quotes ───────────────────────────────────────────────────────────
quotes_dir = parsed_dir / "quotes"
quotes_dir.mkdir(parents=True, exist_ok=True)
quote_rows = []
for i in range(30):
ts = base_date + timedelta(seconds=i * 60)
spread = round(abs(RNG.normal(0.02, 0.01)), 4)
mid = base_price + RNG.normal(0, 0.5)
quote_rows.append(
{
"timestamp": ts,
"symbol": "AAPL",
"bid_price": round(mid - spread / 2, 2),
"bid_size": int(RNG.integers(100, 5001)),
"ask_price": round(mid + spread / 2, 2),
"ask_size": int(RNG.integers(100, 5001)),
}
)
quotes_df = pl.DataFrame(quote_rows).cast(
{
"timestamp": pl.Datetime("ns"),
"bid_price": pl.Float64,
"ask_price": pl.Float64,
"bid_size": pl.Int64,
"ask_size": pl.Int64,
}
)
quotes_df.write_parquet(quotes_dir / "data.parquet")
print(f" IEX quotes: {len(quotes_df)} rows -> {quotes_dir / 'data.parquet'}")
# ── Trades ───────────────────────────────────────────────────────────
trades_dir = parsed_dir / "trades"
trades_dir.mkdir(parents=True, exist_ok=True)
trade_rows = []
for i in range(20):
ts = base_date + timedelta(seconds=i * 90 + int(RNG.integers(0, 30)))
trade_rows.append(
{
"timestamp": ts,
"symbol": "AAPL",
"price": round(base_price + RNG.normal(0, 0.3), 2),
"size": int(RNG.integers(1, 501)),
}
)
trades_df = pl.DataFrame(trade_rows).cast(
{
"timestamp": pl.Datetime("ns"),
"price": pl.Float64,
"size": pl.Int64,
}
)
trades_df.write_parquet(trades_dir / "data.parquet")
print(f" IEX trades: {len(trades_df)} rows -> {trades_dir / 'data.parquet'}")
# ── Price Levels ─────────────────────────────────────────────────────
price_levels_dir = parsed_dir / "price_levels"
price_levels_dir.mkdir(parents=True, exist_ok=True)
pl_rows = []
for i in range(40):
ts = base_date + timedelta(seconds=i * 45)
side = "bid" if i % 2 == 0 else "ask"
offset = RNG.uniform(0.01, 0.50)
price = round(base_price - offset if side == "bid" else base_price + offset, 2)
pl_rows.append(
{
"timestamp": ts,
"symbol": "AAPL",
"side": side,
"price": price,
"size": int(RNG.integers(100, 3001)),
}
)
pl_df = pl.DataFrame(pl_rows).cast(
{
"timestamp": pl.Datetime("ns"),
"price": pl.Float64,
"size": pl.Int64,
}
)
pl_df.write_parquet(price_levels_dir / "data.parquet")
print(f" IEX price_levels: {len(pl_df)} rows -> {price_levels_dir / 'data.parquet'}")
# ═════════════════════════════════════════════════════════════════════════════
# 5. CME Individual Contracts (for Ch02 NB 04-06)
# ═════════════════════════════════════════════════════════════════════════════
# Path: ML4T_DATA_PATH / "futures" / "market" / "individual" / "{PRODUCT}" / "data.parquet"
# Schema matches what load_cme_futures(continuous=False) returns:
# timestamp (datetime[ns, UTC]), rtype, publisher_id, instrument_id,
# open, high, low, close, volume, product
#
# NB06 (futures_continuous) needs:
# - Multiple contracts with OVERLAPPING date ranges
# - Volume patterns that make front-month detection possible
# - Enough contracts for roll detection to produce adj_close
def generate_individual_futures() -> None:
"""Generate synthetic CME individual contract data for ES, NQ, CL.
Key requirements from NB06 (continuous construction):
- Contracts must overlap in time (concurrent trading)
- Front month should have highest volume (for volume-based roll detection)
- Need at least 3 contracts with clear roll transitions
- Need enough data points for roll gaps to produce adj_close
"""
individual_dir = TEST_DATA_ROOT / "futures" / "market" / "individual"
products = {
"ES": {"base_price": 4500.0, "tick": 0.25},
"NQ": {"base_price": 15500.0, "tick": 0.25},
"CL": {"base_price": 75.0, "tick": 0.01},
}
# Contract months: H=March, M=June, U=Sep, Z=Dec
# Instrument IDs encode contract month. Simulate 4 quarterly contracts
# overlapping across 2024, with volume-based rolls.
contract_specs = [
# (instrument_id, start_day_offset, end_day_offset, is_front_until_day)
# Contract 1 (H24): front month days 0-29, then rolls to contract 2
(49701, 0, 59, 29),
# Contract 2 (M24): front month days 30-89, then rolls to contract 3
(49702, 15, 119, 89),
# Contract 3 (U24): front month days 90-149, then rolls to contract 4
(49703, 75, 179, 149),
# Contract 4 (Z24): front month from day 150 onward
(49704, 135, 209, 209),
]
for product, cfg in products.items():
prod_dir = individual_dir / product
prod_dir.mkdir(parents=True, exist_ok=True)
rows = []
start = datetime(2024, 1, 2, 0, 0, 0)
for inst_id, start_day, end_day, front_until in contract_specs:
# Adjust instrument_id per product to be unique
if product == "NQ":
inst_id += 1000
elif product == "CL":
inst_id += 2000
for day_offset in range(start_day, end_day + 1):
# Generate one bar per day (24 hours apart for hourly-like data)
ts = start + timedelta(days=day_offset)
# Price drifts slightly
base = cfg["base_price"] + RNG.normal(0, cfg["base_price"] * 0.002)
o = round(base, 2)
h = round(base + abs(RNG.normal(0, cfg["base_price"] * 0.001)), 2)
l = round(base - abs(RNG.normal(0, cfg["base_price"] * 0.001)), 2)
c = round(base + RNG.normal(0, cfg["base_price"] * 0.0005), 2)
# Volume: high when front month, low when back month
if day_offset <= front_until:
vol = int(RNG.integers(10000, 50001)) # Front month: high volume
else:
vol = int(RNG.integers(100, 3001)) # Back month: low volume
rows.append(
{
"timestamp": ts,
"rtype": 35,
"publisher_id": 1,
"instrument_id": inst_id,
"open": o,
"high": h,
"low": l,
"close": c,
"volume": vol,
"product": product,
}
)
df = (
pl.DataFrame(rows)
.cast(
{
"timestamp": pl.Datetime("ns", time_zone="UTC"),
"rtype": pl.UInt8,
"publisher_id": pl.UInt16,
"instrument_id": pl.UInt32,
"open": pl.Float64,
"high": pl.Float64,
"low": pl.Float64,
"close": pl.Float64,
"volume": pl.UInt64,
}
)
.sort("timestamp")
)
df.write_parquet(prod_dir / "data.parquet")
print(f" Futures individual {product}: {len(df)} rows -> {prod_dir / 'data.parquet'}")
# ═════════════════════════════════════════════════════════════════════════════
# 6. Kalshi Events (for Ch04 NB 13)
# ═════════════════════════════════════════════════════════════════════════════
# Path: ML4T_DATA_PATH / "prediction_markets" / "kalshi_events.parquet"
# Schema: timestamp (Date), symbol (str), open/high/low/close (Float64), volume (Int64)
def generate_kalshi_data() -> None:
"""Generate synthetic Kalshi prediction market data."""
pm_dir = TEST_DATA_ROOT / "prediction_markets"
pm_dir.mkdir(parents=True, exist_ok=True)
# 5 contracts, ~10 days each = ~50 rows
contracts = [
"KXFED-27APR-T4.25",
"KXFED-27APR-T4.50",
"KXFED-27JUN-T4.00",
"KXINFL-27MAR-T3.0",
"KXGDP-27Q1-T2.0",
]
rows = []
base_date = date(2027, 3, 1)
for contract in contracts:
# Each contract gets a base probability and drifts
base_prob = RNG.uniform(0.2, 0.8)
for day in range(10):
d = base_date + timedelta(days=day)
# Random walk for probability
base_prob = max(0.01, min(0.99, base_prob + RNG.normal(0, 0.03)))
o = round(base_prob, 2)
h = round(min(0.99, base_prob + abs(RNG.normal(0, 0.02))), 2)
l = round(max(0.01, base_prob - abs(RNG.normal(0, 0.02))), 2)
c = round(max(0.01, min(0.99, base_prob + RNG.normal(0, 0.01))), 2)
vol = int(RNG.integers(50, 5001))
rows.append(
{
"timestamp": d,
"symbol": contract,
"open": o,
"high": h,
"low": l,
"close": c,
"volume": vol,
}
)
df = (
pl.DataFrame(rows)
.cast(
{
"timestamp": pl.Date,
"open": pl.Float64,
"high": pl.Float64,
"low": pl.Float64,
"close": pl.Float64,
"volume": pl.Int64,
}
)
.sort(["symbol", "timestamp"])
)
df.write_parquet(pm_dir / "kalshi_events.parquet")
print(f" Kalshi events: {len(df)} rows -> {pm_dir / 'kalshi_events.parquet'}")
# ═════════════════════════════════════════════════════════════════════════════
# Main
# ═════════════════════════════════════════════════════════════════════════════
def main() -> None:
print(f"Generating test microstructure data in {TEST_DATA_ROOT}\n")
print("1. ITCH Parsed Messages")
generate_itch_messages()
print("\n2. DataBento MBO")
generate_mbo_data()
print("\n3. AlgoSeek TAQ")
generate_taq_data()
print("\n4. IEX Parsed Data")
generate_iex_data()
print("\n5. CME Individual Futures")
generate_individual_futures()
print("\n6. Kalshi Prediction Markets")
generate_kalshi_data()
print("\nDone.")
if __name__ == "__main__":
main()
tests/notebook_catalog.py
+769
View File
@@ -0,0 +1,769 @@
from __future__ import annotations
import json
import os
import sqlite3
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from tests.pm_helpers import collect_chapter_notebooks, get_overrides
REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_DB_PATH = REPO_ROOT / ".claude" / "work" / "notebook_testing" / "catalog.sqlite"
CASE_STUDIES = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
TRACKER_SCHEMA_COMPLETE_CHAPTERS = {1, 2, 8, 9, 10, 11, 12, 13, 14, 15}
TRACKER_SCHEMA_IN_PROGRESS_CHAPTERS = {3, 4, 5, 6, 7, 16, 17, 18, 19, 20}
TRACKER_SCHEMA_CASE_STUDIES = {
"etfs": "complete",
"fx_pairs": "complete",
"crypto_perps_funding": "in_progress",
"cme_futures": "pending",
"nasdaq100_microstructure": "pending",
"sp500_equity_option_analytics": "pending",
"sp500_options": "pending",
"us_equities_panel": "pending",
"us_firm_characteristics": "pending",
}
CRYPTO_REPRO_NOTE = (
"Current Binance public downloads no longer reproduce MATICUSDT OHLCV. "
"Crypto case study requires full refreshed model reruns and explicit old-vs-new "
"comparison against the dev registry."
)
HEAVY_KEYWORDS = {
"timegan",
"tailgan",
"sigcwgan",
"diffusion",
"great",
"dp_gan",
"patchtst",
"transformer",
"lstm",
"autoencoder",
"finbert",
"bert",
"ner",
"xgboost",
"lightgbm",
"catboost",
"rl",
"deepm",
"gan",
"backtest_sweep",
}
GPU_KEYWORDS = {"gpu", "cuda", "torch", "tensorflow", "trainer"}
DATA_HINTS = {
"etfs": ("etf", "etfs", "spy"),
"crypto": ("crypto", "perp", "perps", "funding", "premium", "binance"),
"fx": ("fx", "oanda", "eur_usd"),
"futures": ("future", "futures", "cme", "databento", "glbx"),
"us_equities": ("equities", "crsp", "stocks", "ticker", "secedgar"),
"options": ("option", "options", "greeks", "iv"),
"nasdaq_itch": ("itch", "nasdaq100", "algoseek", "taq", "lob", "iex"),
"macro": ("macro", "fred", "yield", "calendar"),
"text": ("text", "news", "filing", "sentiment", "word2vec"),
"synthetic": ("synthetic", "simulation", "regime", "scenario"),
}
@dataclass(slots=True)
class NotebookEntry:
path: str
notebook_key: str
notebook_type: str
chapter: int | None
case_study_id: str | None
stage: str | None
stage_order: int | None
title: str
resource_profile: str
execution_lane: str
parallel_safe: int
worker_slots: int
gpu_required: int
has_parameters_cell: int
parameter_source: str
default_timeout_seconds: int
inputs_hint: str
override_json: str
def utc_now() -> str:
return datetime.now(UTC).replace(microsecond=0).isoformat()
def connect_catalog(db_path: Path | None = None) -> sqlite3.Connection:
path = Path(db_path) if db_path else DEFAULT_DB_PATH
path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(path)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys=ON")
ensure_schema(conn)
return conn
def ensure_schema(conn: sqlite3.Connection) -> None:
conn.executescript(
"""
CREATE TABLE IF NOT EXISTS notebooks (
path TEXT PRIMARY KEY,
notebook_key TEXT NOT NULL UNIQUE,
notebook_type TEXT NOT NULL,
chapter INTEGER,
case_study_id TEXT,
stage TEXT,
stage_order INTEGER,
title TEXT NOT NULL,
resource_profile TEXT NOT NULL,
execution_lane TEXT NOT NULL,
parallel_safe INTEGER NOT NULL DEFAULT 0,
worker_slots INTEGER NOT NULL DEFAULT 1,
gpu_required INTEGER NOT NULL DEFAULT 0,
has_parameters_cell INTEGER NOT NULL DEFAULT 0,
parameter_source TEXT NOT NULL DEFAULT 'none',
default_timeout_seconds INTEGER NOT NULL DEFAULT 300,
inputs_hint TEXT NOT NULL DEFAULT '',
override_json TEXT NOT NULL DEFAULT '{}',
last_inventory_at TEXT,
last_status TEXT NOT NULL DEFAULT 'pending',
last_execution_mode TEXT,
last_runtime_seconds REAL,
last_peak_memory_mb REAL,
last_error_type TEXT,
last_error_message TEXT,
last_run_at TEXT,
last_batch_id TEXT,
notes TEXT NOT NULL DEFAULT ''
);
CREATE TABLE IF NOT EXISTS notebook_runs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
batch_id TEXT NOT NULL,
path TEXT NOT NULL,
notebook_type TEXT NOT NULL,
chapter INTEGER,
case_study_id TEXT,
stage TEXT,
status TEXT NOT NULL,
execution_mode TEXT NOT NULL,
runtime_seconds REAL,
peak_memory_mb REAL,
timeout_seconds INTEGER NOT NULL,
worker_slots INTEGER NOT NULL,
output_root TEXT,
data_root TEXT,
parameter_source TEXT NOT NULL,
parameters_json TEXT NOT NULL DEFAULT '{}',
error_type TEXT,
error_message TEXT,
log_path TEXT,
started_at TEXT NOT NULL,
finished_at TEXT NOT NULL,
FOREIGN KEY(path) REFERENCES notebooks(path)
);
CREATE INDEX IF NOT EXISTS idx_notebooks_type_chapter
ON notebooks(notebook_type, chapter, case_study_id, stage_order);
CREATE INDEX IF NOT EXISTS idx_notebooks_status
ON notebooks(last_status, notebook_type, chapter);
CREATE INDEX IF NOT EXISTS idx_runs_batch
ON notebook_runs(batch_id, notebook_type, chapter, case_study_id, stage);
CREATE TABLE IF NOT EXISTS program_tracker (
item_key TEXT PRIMARY KEY,
track TEXT NOT NULL,
scope_type TEXT NOT NULL,
scope_id TEXT NOT NULL,
label TEXT NOT NULL,
sort_order INTEGER NOT NULL,
required INTEGER NOT NULL DEFAULT 1,
status TEXT NOT NULL DEFAULT 'pending',
status_source TEXT NOT NULL DEFAULT 'auto',
metrics_json TEXT NOT NULL DEFAULT '{}',
notes TEXT NOT NULL DEFAULT '',
updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_program_tracker_track
ON program_tracker(track, scope_type, sort_order, scope_id);
"""
)
conn.commit()
def _read_text(path: Path) -> str:
return path.read_text(encoding="utf-8", errors="ignore")
def _has_parameters_cell(text: str) -> bool:
tags = ('tags=["parameters"]', "tags=['parameters']", '"parameters"', "'parameters'")
return any(tag in text for tag in tags)
def _detect_inputs(text: str, rel_path: str) -> str:
haystack = f"{rel_path.lower()} {text.lower()}"
found = [name for name, markers in DATA_HINTS.items() if any(m in haystack for m in markers)]
return ",".join(sorted(found))
def _classify_entry(
notebook_type: str,
rel_path: str,
text: str,
chapter: int | None,
stage_order: int | None,
overrides: dict,
) -> tuple[str, str, int, int]:
key = rel_path.lower()
haystack = f"{key} {text.lower()}"
if notebook_type == "case_study":
if overrides.get("gpu") or any(k in haystack for k in GPU_KEYWORDS):
return "pipeline_gpu", "serial_case_study", 0, 4
if stage_order is not None and stage_order >= 11:
return "pipeline_model", "serial_case_study", 0, 3
if stage_order is not None and stage_order >= 6:
return "pipeline_feature", "serial_case_study", 0, 2
return "pipeline_setup", "serial_case_study", 0, 1
if overrides.get("gpu") or any(k in haystack for k in GPU_KEYWORDS):
return "gpu", "serial_heavy", 0, 4
if any(k in haystack for k in HEAVY_KEYWORDS):
return "heavy", "serial_heavy", 0, 3
if chapter is not None and chapter <= 6 and not overrides.get("parameters"):
return "light", "parallel_light", 1, 1
if chapter is not None and chapter <= 10:
return "medium", "parallel_medium", 1, 2
if overrides.get("parameters"):
return "medium", "parallel_medium", 1, 2
return "heavy", "serial_heavy", 0, 3
def _default_timeout(overrides: dict) -> int:
return int(overrides.get("timeout", 300))
def _parameter_source(overrides: dict, has_parameters_cell: bool) -> str:
if overrides.get("parameters"):
return "papermill"
if has_parameters_cell:
return "none"
return "config"
def _chapter_entries(repo_root: Path) -> list[NotebookEntry]:
entries: list[NotebookEntry] = []
for path in collect_chapter_notebooks(repo_root, range(1, 28)):
rel = path.relative_to(repo_root)
notebook_key = str(rel.with_suffix("")).replace(os.sep, "/")
text = _read_text(path)
overrides = get_overrides(notebook_key)
chapter = int(rel.parts[0][:2])
has_parameters_cell = int(_has_parameters_cell(text))
resource_profile, execution_lane, parallel_safe, worker_slots = _classify_entry(
"chapter", rel.as_posix(), text, chapter, None, overrides
)
entries.append(
NotebookEntry(
path=rel.as_posix(),
notebook_key=notebook_key,
notebook_type="chapter",
chapter=chapter,
case_study_id=None,
stage=None,
stage_order=None,
title=path.stem,
resource_profile=resource_profile,
execution_lane=execution_lane,
parallel_safe=parallel_safe,
worker_slots=worker_slots,
gpu_required=int(bool(overrides.get("gpu"))),
has_parameters_cell=has_parameters_cell,
parameter_source=_parameter_source(overrides, bool(has_parameters_cell)),
default_timeout_seconds=_default_timeout(overrides),
inputs_hint=_detect_inputs(text, rel.as_posix()),
override_json=json.dumps(overrides, sort_keys=True),
)
)
return entries
def _case_study_entries(repo_root: Path) -> list[NotebookEntry]:
entries: list[NotebookEntry] = []
for case_study_id in CASE_STUDIES:
cs_dir = repo_root / "case_studies" / case_study_id
if not cs_dir.exists():
continue
for path in sorted(cs_dir.glob("[0-9][0-9]_*.py")):
if path.name.startswith("_"):
continue
rel = path.relative_to(repo_root)
notebook_key = str(rel.with_suffix("")).replace(os.sep, "/")
text = _read_text(path)
overrides = get_overrides(notebook_key)
stage = path.stem
stage_order = int(stage[:2]) if stage[:2].isdigit() else None
has_parameters_cell = int(_has_parameters_cell(text))
resource_profile, execution_lane, parallel_safe, worker_slots = _classify_entry(
"case_study", rel.as_posix(), text, None, stage_order, overrides
)
entries.append(
NotebookEntry(
path=rel.as_posix(),
notebook_key=notebook_key,
notebook_type="case_study",
chapter=None,
case_study_id=case_study_id,
stage=stage,
stage_order=stage_order,
title=path.stem,
resource_profile=resource_profile,
execution_lane=execution_lane,
parallel_safe=parallel_safe,
worker_slots=worker_slots,
gpu_required=int(bool(overrides.get("gpu"))),
has_parameters_cell=has_parameters_cell,
parameter_source=_parameter_source(overrides, bool(has_parameters_cell)),
default_timeout_seconds=_default_timeout(overrides),
inputs_hint=_detect_inputs(text, rel.as_posix()),
override_json=json.dumps(overrides, sort_keys=True),
)
)
return entries
def build_inventory(repo_root: Path | None = None) -> list[NotebookEntry]:
root = repo_root or REPO_ROOT
return _chapter_entries(root) + _case_study_entries(root)
def upsert_inventory(conn: sqlite3.Connection, entries: list[NotebookEntry]) -> None:
now = utc_now()
current_paths = [entry.path for entry in entries]
conn.executemany(
"""
INSERT INTO notebooks (
path,
notebook_key,
notebook_type,
chapter,
case_study_id,
stage,
stage_order,
title,
resource_profile,
execution_lane,
parallel_safe,
worker_slots,
gpu_required,
has_parameters_cell,
parameter_source,
default_timeout_seconds,
inputs_hint,
override_json,
last_inventory_at
) VALUES (
:path,
:notebook_key,
:notebook_type,
:chapter,
:case_study_id,
:stage,
:stage_order,
:title,
:resource_profile,
:execution_lane,
:parallel_safe,
:worker_slots,
:gpu_required,
:has_parameters_cell,
:parameter_source,
:default_timeout_seconds,
:inputs_hint,
:override_json,
:last_inventory_at
)
ON CONFLICT(path) DO UPDATE SET
notebook_key=excluded.notebook_key,
notebook_type=excluded.notebook_type,
chapter=excluded.chapter,
case_study_id=excluded.case_study_id,
stage=excluded.stage,
stage_order=excluded.stage_order,
title=excluded.title,
resource_profile=excluded.resource_profile,
execution_lane=excluded.execution_lane,
parallel_safe=excluded.parallel_safe,
worker_slots=excluded.worker_slots,
gpu_required=excluded.gpu_required,
has_parameters_cell=excluded.has_parameters_cell,
parameter_source=excluded.parameter_source,
default_timeout_seconds=excluded.default_timeout_seconds,
inputs_hint=excluded.inputs_hint,
override_json=excluded.override_json,
last_inventory_at=excluded.last_inventory_at
""",
[
{
**asdict(entry),
"last_inventory_at": now,
}
for entry in entries
],
)
if current_paths:
marks = ",".join("?" for _ in current_paths)
stale_paths = [
row[0]
for row in conn.execute(
f"SELECT path FROM notebooks WHERE path NOT IN ({marks})", current_paths
).fetchall()
]
if stale_paths:
stale_marks = ",".join("?" for _ in stale_paths)
conn.execute(f"DELETE FROM notebook_runs WHERE path IN ({stale_marks})", stale_paths)
conn.execute(f"DELETE FROM notebooks WHERE path NOT IN ({marks})", current_paths)
conn.commit()
def refresh_inventory(
conn: sqlite3.Connection, repo_root: Path | None = None
) -> list[NotebookEntry]:
entries = build_inventory(repo_root)
upsert_inventory(conn, entries)
return entries
def resolve_data_root(repo_root: Path | None = None) -> Path | None:
root = repo_root or REPO_ROOT
env_data = os.environ.get("ML4T_DATA_PATH")
if env_data:
path = Path(env_data).expanduser()
if path.exists():
return path
env_file = root / ".env"
if not env_file.exists():
return None
for line in env_file.read_text(encoding="utf-8", errors="ignore").splitlines():
line = line.strip()
if not line or line.startswith("#") or "=" not in line:
continue
key, value = line.split("=", 1)
if key.strip() == "ML4T_DATA_PATH":
path = Path(value.strip().strip('"').strip("'")).expanduser()
if path.exists():
return path
return None
def latest_status_counts(conn: sqlite3.Connection) -> list[sqlite3.Row]:
return conn.execute(
"""
SELECT notebook_type, last_status, COUNT(*) AS n
FROM notebooks
GROUP BY notebook_type, last_status
ORDER BY notebook_type, last_status
"""
).fetchall()
def _latest_notebook_rows(conn: sqlite3.Connection) -> list[sqlite3.Row]:
return conn.execute(
"""
WITH latest AS (
SELECT
nr.*,
ROW_NUMBER() OVER (PARTITION BY path ORDER BY id DESC) AS rn
FROM notebook_runs nr
)
SELECT
n.path,
n.notebook_type,
n.chapter,
n.case_study_id,
n.stage_order,
COALESCE(l.status, 'pending') AS status,
l.runtime_seconds,
l.peak_memory_mb
FROM notebooks n
LEFT JOIN latest l
ON n.path = l.path
AND l.rn = 1
"""
).fetchall()
def _rollup_status(statuses: list[str]) -> str:
if not statuses:
return "pending"
unique = set(statuses)
if unique == {"ok"}:
return "complete"
if unique == {"pending"}:
return "pending"
if {"error", "blocked", "skipped"} & unique:
return "blocked"
if "ok" in unique and "pending" in unique:
return "in_progress"
if "ok" in unique:
return "in_progress"
return "pending"
def _functional_chapter_metrics(
rows: list[sqlite3.Row], chapter: int
) -> tuple[str, dict[str, int]]:
chapter_rows = [
row for row in rows if row["notebook_type"] == "chapter" and row["chapter"] == chapter
]
counts: dict[str, int] = {}
for row in chapter_rows:
counts[row["status"]] = counts.get(row["status"], 0) + 1
return _rollup_status([row["status"] for row in chapter_rows]), {
"total": len(chapter_rows),
**counts,
}
def _functional_case_study_metrics(
rows: list[sqlite3.Row], case_study_id: str, max_stage: int = 5
) -> tuple[str, dict[str, int]]:
cs_rows = [
row
for row in rows
if row["notebook_type"] == "case_study"
and row["case_study_id"] == case_study_id
and row["stage_order"] is not None
and row["stage_order"] <= max_stage
]
counts: dict[str, int] = {}
for row in cs_rows:
counts[row["status"]] = counts.get(row["status"], 0) + 1
return _rollup_status([row["status"] for row in cs_rows]), {
"total": len(cs_rows),
**counts,
}
def _schema_status_for_chapter(chapter: int) -> str:
if chapter in TRACKER_SCHEMA_COMPLETE_CHAPTERS:
return "complete"
if chapter in TRACKER_SCHEMA_IN_PROGRESS_CHAPTERS:
return "in_progress"
return "pending"
def _schema_status_for_case_study(case_study_id: str) -> str:
return TRACKER_SCHEMA_CASE_STUDIES.get(case_study_id, "pending")
def sync_program_tracker(conn: sqlite3.Connection) -> None:
rows = _latest_notebook_rows(conn)
now = utc_now()
tracker_rows: list[dict[str, object]] = [
{
"item_key": "foundation:catalog",
"track": "foundation",
"scope_type": "foundation",
"scope_id": "catalog",
"label": "Notebook catalog and Docker runner",
"sort_order": 0,
"required": 1,
"status": "complete",
"status_source": "manual",
"metrics_json": json.dumps({}, sort_keys=True),
"notes": "SQLite catalog, Docker runner, and isolated-output execution are in place.",
"updated_at": now,
}
]
for chapter in range(1, 21):
functional_status, functional_metrics = _functional_chapter_metrics(rows, chapter)
tracker_rows.append(
{
"item_key": f"chapter:{chapter:02d}:functional",
"track": "functional",
"scope_type": "chapter",
"scope_id": f"{chapter:02d}",
"label": f"Chapter {chapter:02d} functional correctness",
"sort_order": chapter,
"required": 1,
"status": functional_status,
"status_source": "auto",
"metrics_json": json.dumps(functional_metrics, sort_keys=True),
"notes": "",
"updated_at": now,
}
)
tracker_rows.append(
{
"item_key": f"chapter:{chapter:02d}:schema",
"track": "schema",
"scope_type": "chapter",
"scope_id": f"{chapter:02d}",
"label": f"Chapter {chapter:02d} canonical schema retrofit",
"sort_order": chapter,
"required": 1,
"status": _schema_status_for_chapter(chapter),
"status_source": "manual",
"metrics_json": json.dumps({}, sort_keys=True),
"notes": "",
"updated_at": now,
}
)
repro_required = int(chapter >= 11)
repro_status = "pending" if repro_required else "not_required"
repro_notes = (
"Required for model/results notebooks and any book-facing figures or reported results."
if repro_required
else "Teaching notebooks default to functional-only unless book-facing outputs require parity."
)
tracker_rows.append(
{
"item_key": f"chapter:{chapter:02d}:repro",
"track": "repro",
"scope_type": "chapter",
"scope_id": f"{chapter:02d}",
"label": f"Chapter {chapter:02d} dev-registry reproducibility validation",
"sort_order": chapter,
"required": repro_required,
"status": repro_status,
"status_source": "manual",
"metrics_json": json.dumps({}, sort_keys=True),
"notes": repro_notes,
"updated_at": now,
}
)
for idx, case_study_id in enumerate(CASE_STUDIES, start=1):
functional_status, functional_metrics = _functional_case_study_metrics(rows, case_study_id)
tracker_rows.append(
{
"item_key": f"case_study:{case_study_id}:functional_1_5",
"track": "functional",
"scope_type": "case_study",
"scope_id": case_study_id,
"label": f"{case_study_id} stages 1-5 functional correctness",
"sort_order": idx,
"required": 1,
"status": functional_status,
"status_source": "auto",
"metrics_json": json.dumps(functional_metrics, sort_keys=True),
"notes": "",
"updated_at": now,
}
)
tracker_rows.append(
{
"item_key": f"case_study:{case_study_id}:schema_1_5",
"track": "schema",
"scope_type": "case_study",
"scope_id": case_study_id,
"label": f"{case_study_id} early-stage canonical schema retrofit",
"sort_order": idx,
"required": 1,
"status": _schema_status_for_case_study(case_study_id),
"status_source": "manual",
"metrics_json": json.dumps({}, sort_keys=True),
"notes": "",
"updated_at": now,
}
)
tracker_rows.append(
{
"item_key": f"case_study:{case_study_id}:repro",
"track": "repro",
"scope_type": "case_study",
"scope_id": case_study_id,
"label": f"{case_study_id} dev-registry reproducibility validation",
"sort_order": idx,
"required": 1,
"status": "pending",
"status_source": "manual",
"metrics_json": json.dumps({}, sort_keys=True),
"notes": CRYPTO_REPRO_NOTE if case_study_id == "crypto_perps_funding" else "",
"updated_at": now,
}
)
conn.executemany(
"""
INSERT INTO program_tracker (
item_key,
track,
scope_type,
scope_id,
label,
sort_order,
required,
status,
status_source,
metrics_json,
notes,
updated_at
) VALUES (
:item_key,
:track,
:scope_type,
:scope_id,
:label,
:sort_order,
:required,
:status,
:status_source,
:metrics_json,
:notes,
:updated_at
)
ON CONFLICT(item_key) DO UPDATE SET
track=excluded.track,
scope_type=excluded.scope_type,
scope_id=excluded.scope_id,
label=excluded.label,
sort_order=excluded.sort_order,
required=excluded.required,
status=excluded.status,
status_source=excluded.status_source,
metrics_json=excluded.metrics_json,
notes=excluded.notes,
updated_at=excluded.updated_at
""",
tracker_rows,
)
conn.commit()
def tracker_status_counts(conn: sqlite3.Connection) -> list[sqlite3.Row]:
return conn.execute(
"""
SELECT track, status, COUNT(*) AS n
FROM program_tracker
GROUP BY track, status
ORDER BY track, status
"""
).fetchall()
tests/notebook_worker.py
+189
View File
@@ -0,0 +1,189 @@
from __future__ import annotations
import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path
def _sync_if_needed(py_path: Path, sync_policy: str) -> Path:
ipynb_path = py_path.with_suffix(".ipynb")
should_sync = sync_policy == "always" or (sync_policy == "missing" and not ipynb_path.exists())
if should_sync:
result = subprocess.run(
[
sys.executable,
"-m",
"jupytext",
"--to",
"notebook",
"--set-kernel",
"python3",
"--output",
str(ipynb_path),
str(py_path),
],
capture_output=True,
text=True,
timeout=60,
cwd=str(
py_path.parent.parent
if py_path.parent.name.startswith("case_studies")
else py_path.parent
),
)
if result.returncode != 0:
raise RuntimeError(f"Jupytext sync failed for {py_path}: {result.stderr}")
if not ipynb_path.exists():
raise FileNotFoundError(f"Expected .ipynb not found for {py_path}")
return ipynb_path
def _run_full_notebook(
py_path: Path,
timeout: int,
output_dir: Path | None,
data_dir: Path | None,
extra_env: dict[str, str],
sync_policy: str,
) -> dict:
import papermill as pm
start = time.perf_counter()
ipynb_path = _sync_if_needed(py_path, sync_policy)
tmp_out = Path(tempfile.gettempdir()) / f"ml4t-full-{os.getpid()}-{py_path.stem}.ipynb"
saved_env: dict[str, str | None] = {}
rc_dir = Path(tempfile.mkdtemp(prefix="ml4t-mplrc-"))
rc_file = rc_dir / "matplotlibrc"
rc_file.write_text("figure.constrained_layout.use: False\n", encoding="utf-8")
env_vars = {
"MPLBACKEND": "Agg",
"PLOTLY_RENDERER": "json",
"PYTHONUNBUFFERED": "1",
"MATPLOTLIBRC": str(rc_file),
}
if output_dir:
output_dir.mkdir(parents=True, exist_ok=True)
env_vars["ML4T_OUTPUT_DIR"] = str(output_dir)
if data_dir:
env_vars["ML4T_DATA_PATH"] = str(data_dir)
env_vars.update({key: str(value) for key, value in extra_env.items()})
existing = os.environ.get("PYTHONPATH", "")
repo_root = py_path.parents[2] if "case_studies" in py_path.parts else py_path.parent.parent
nb_dir = str(py_path.parent.resolve())
env_vars["PYTHONPATH"] = (
f"{repo_root}:{nb_dir}:{existing}" if existing else f"{repo_root}:{nb_dir}"
)
try:
import torch
torch_lib = str(Path(torch.__file__).parent / "lib")
nvidia_libs = list((Path(torch.__file__).parent.parent / "nvidia").glob("*/lib"))
cuda_paths = [torch_lib] + [str(p) for p in nvidia_libs]
existing_ld = os.environ.get("LD_LIBRARY_PATH", "")
env_vars["LD_LIBRARY_PATH"] = ":".join(cuda_paths + [existing_ld])
except ImportError:
pass
for key, value in env_vars.items():
saved_env[key] = os.environ.get(key)
os.environ[key] = value
try:
pm.execute_notebook(
str(ipynb_path),
str(tmp_out),
parameters={},
cwd=str(repo_root),
kernel_name="python3",
execution_timeout=timeout,
request_save_on_cell_execute=True,
progress_bar=False,
log_output=True,
)
shutil.copy2(tmp_out, ipynb_path)
return {"status": "ok", "error": None, "runtime_seconds": time.perf_counter() - start}
except pm.PapermillExecutionError as exc:
return {
"status": "error",
"error": f"Cell {exc.cell_index} ({exc.ename}): {exc.evalue}",
"runtime_seconds": time.perf_counter() - start,
}
except Exception as exc:
return {
"status": "error",
"error": str(exc),
"runtime_seconds": time.perf_counter() - start,
}
finally:
for key, value in saved_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
tmp_out.unlink(missing_ok=True)
shutil.rmtree(rc_dir, ignore_errors=True)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--path", required=True)
parser.add_argument("--timeout", type=int, required=True)
parser.add_argument("--output-dir")
parser.add_argument("--result-file", required=True)
parser.add_argument("--data-dir")
parser.add_argument("--parameters-json", default="{}")
parser.add_argument("--env-json", default="{}")
parser.add_argument("--execution-mode", choices=["reduced", "full"], default="reduced")
parser.add_argument("--sync-policy", choices=["always", "missing", "never"], default="always")
args = parser.parse_args()
path = Path(args.path).resolve()
output_dir = Path(args.output_dir).resolve() if args.output_dir else None
result_file = Path(args.result_file).resolve()
params = json.loads(args.parameters_json)
extra_env = json.loads(args.env_json)
data_dir = Path(args.data_dir).resolve() if args.data_dir else None
started = time.perf_counter()
if args.execution_mode == "full":
result = _run_full_notebook(
py_path=path,
timeout=args.timeout,
output_dir=output_dir,
data_dir=data_dir,
extra_env=extra_env,
sync_policy=args.sync_policy,
)
else:
from tests.pm_helpers import run_notebook
result = run_notebook(
py_path=path,
parameters=params,
timeout=args.timeout,
output_dir=output_dir,
data_dir=data_dir,
extra_env=extra_env,
)
elapsed = time.perf_counter() - started
payload = {
"status": result.get("status", "error"),
"error": result.get("error"),
"runtime_seconds": elapsed,
}
result_file.parent.mkdir(parents=True, exist_ok=True)
result_file.write_text(json.dumps(payload), encoding="utf-8")
if __name__ == "__main__":
main()
tests/overrides.yaml
+2261
View File
File diff suppressed because it is too large Load Diff
tests/pm_helpers.py
+415
View File
@@ -0,0 +1,415 @@
"""Papermill-based notebook execution helpers.
Provides:
- run_notebook(): Execute a .py notebook via Papermill with parameter injection
- get_overrides(): Load per-notebook overrides from tests/overrides.yaml
- collect_chapter_notebooks(): Discover notebooks in chapter directories
- get_tier() / current_test_tier(): Test-tier routing (per-commit / weekly / on-demand)
- get_record_mode(): VCR cassette mode (consumed by Step 5)
NOTE: Notebooks live directly in chapter directories
(e.g., 05_synthetic_data/01_timegan.py), NOT in code/ subdirs.
overrides.yaml schema (per-notebook, all optional):
timeout: int (seconds, default 300)
parameters: dict (papermill -p overrides)
skip: bool — hard skip in uv-native run (Docker tests ignore)
skip_reason: str
requires_import: str | list[str]
gpu: bool
long_running: bool
docker_env: str — informational (e.g., "benchmark")
tier: "per_commit" | "weekly" | "on_demand" — default "per_commit"
Per-commit runs the Tests workflow on every PR/push.
Weekly runs the weekly-external scheduled workflow (Step 2).
On_demand runs only on manual dispatch (GPU-only NBs).
reruns: int — flaky-retry count (consumed once pytest-rerunfailures lands,
Step 2). Default 0.
record_mode: "replay" | "rewrite" — VCR cassette mode (consumed by Step 5).
Default "replay".
"""
import os
import subprocess
import sys
import tempfile
from pathlib import Path
import yaml
REPO_ROOT = Path(__file__).parent.parent
OVERRIDES_PATH = REPO_ROOT / "tests" / "overrides.yaml"
# Cache loaded overrides
_overrides_cache: dict | None = None
# ---------------------------------------------------------------------------
# Test tier — controls when a notebook runs in CI.
# Per-commit (default): every PR / push triggers the Tests workflow.
# Weekly: only the scheduled weekly-external workflow (Mon 06:00 UTC).
# On-demand: only manual workflow_dispatch (e.g., GPU-only Tier 3).
# ---------------------------------------------------------------------------
TIER_PER_COMMIT = "per_commit"
TIER_WEEKLY = "weekly"
TIER_ON_DEMAND = "on_demand"
VALID_TIERS = frozenset({TIER_PER_COMMIT, TIER_WEEKLY, TIER_ON_DEMAND})
# VCR cassette modes (Step 5: pytest-recording).
RECORD_REPLAY = "replay"
RECORD_REWRITE = "rewrite"
VALID_RECORD_MODES = frozenset({RECORD_REPLAY, RECORD_REWRITE})
def get_tier(overrides: dict) -> str:
"""Return the test tier declared by overrides (default: per_commit)."""
tier = overrides.get("tier") or TIER_PER_COMMIT
if tier not in VALID_TIERS:
raise ValueError(
f"Invalid tier {tier!r} in overrides — must be one of {sorted(VALID_TIERS)}"
)
return tier
def current_test_tier() -> str:
"""Return the tier the current pytest run is targeting.
Read from ML4T_TEST_TIER env var; defaults to per_commit so existing
workflows that don't set it keep their current behavior (only NBs
without a tier key — i.e., tier=per_commit — execute).
"""
tier = os.environ.get("ML4T_TEST_TIER") or TIER_PER_COMMIT
if tier not in VALID_TIERS:
raise ValueError(f"Invalid ML4T_TEST_TIER={tier!r} — must be one of {sorted(VALID_TIERS)}")
return tier
def get_reruns(overrides: dict) -> int:
"""Return per-notebook flaky retry count (default 0).
Consumed by Step 2 (pytest-rerunfailures dep + collection hook adds
@pytest.mark.flaky(reruns=N) when N > 0). Until that lands the value
is parsed but no retries happen.
"""
val = overrides.get("reruns", 0)
if not isinstance(val, int) or val < 0:
raise ValueError(f"Invalid reruns={val!r} — must be non-negative int")
return val
def get_record_mode(overrides: dict) -> str:
"""Return VCR cassette mode for the notebook (default: replay)."""
mode = overrides.get("record_mode") or RECORD_REPLAY
if mode not in VALID_RECORD_MODES:
raise ValueError(
f"Invalid record_mode {mode!r} — must be one of {sorted(VALID_RECORD_MODES)}"
)
return mode
def get_overrides(notebook_key: str) -> dict:
"""Get parameter overrides for a notebook from tests/overrides.yaml.
Args:
notebook_key: Notebook path relative to repo root, no extension.
e.g., "05_synthetic_data/02_tailgan_tail_risk"
Returns:
Dict with optional keys: timeout, gpu, parameters
"""
global _overrides_cache
if _overrides_cache is None:
if OVERRIDES_PATH.exists():
with open(OVERRIDES_PATH) as f:
_overrides_cache = yaml.safe_load(f) or {}
else:
_overrides_cache = {}
return _overrides_cache.get(notebook_key) or {}
def sync_notebook(py_path: Path) -> Path:
"""Sync a .py notebook to a temporary .ipynb via Jupytext.
Writes to a temp file so the real .ipynb (which may contain pre-executed
outputs) is never overwritten.
Args:
py_path: Path to the .py source file
Returns:
Path to a temporary .ipynb file (caller must clean up)
"""
# Write to a temp file — never touch the real .ipynb
tmp_fd, tmp_path_str = tempfile.mkstemp(suffix=".ipynb", prefix=f"_pm_{py_path.stem}_")
os.close(tmp_fd)
tmp_ipynb = Path(tmp_path_str)
result = subprocess.run(
[
sys.executable,
"-m",
"jupytext",
"--to",
"notebook",
"--set-kernel",
"python3",
"--output",
str(tmp_ipynb),
str(py_path),
],
capture_output=True,
text=True,
timeout=60,
cwd=str(REPO_ROOT),
)
if result.returncode != 0:
tmp_ipynb.unlink(missing_ok=True)
raise RuntimeError(f"Jupytext sync failed for {py_path}: {result.stderr}")
if not tmp_ipynb.exists():
raise FileNotFoundError(f"Expected temp .ipynb not found after sync: {tmp_ipynb}")
return tmp_ipynb
def run_notebook(
py_path: Path,
parameters: dict | None = None,
timeout: int = 300,
output_dir: Path | None = None,
data_dir: Path | None = None,
extra_env: dict[str, str] | None = None,
log_path: Path | None = None,
cwd: Path | None = None,
) -> dict:
"""Execute a notebook via Papermill with parameter injection.
This is the core test helper. It:
1. Syncs .py -> .ipynb via Jupytext
2. Executes via Papermill with parameter overrides
3. Logs per-cell progress to log_path (if provided)
4. Returns status and error info
Args:
py_path: Path to the .py notebook source
parameters: Dict of parameters to inject (overrides defaults in parameters cell)
timeout: Per-cell timeout in seconds
output_dir: Directory for ML4T_OUTPUT_DIR (redirects saves to temp)
data_dir: Directory for ML4T_DATA_PATH (test data location)
extra_env: Additional environment variables for notebook execution
log_path: Path to progress log file (appended to)
Returns:
Dict with keys: status ("ok" or "error"), error (str if failed),
duration_s (float), n_cells (int)
"""
import time
import papermill as pm
start = time.time()
nb_name = py_path.stem
def _log(msg: str) -> None:
if log_path:
with open(log_path, "a") as f:
f.write(f"[{time.strftime('%H:%M:%S')}] {msg}\n")
f.flush()
_log(f"START {nb_name} (timeout={timeout}s)")
# Sync to a temp .ipynb (never overwrites the real .ipynb)
tmp_ipynb: Path | None = None
try:
tmp_ipynb = sync_notebook(py_path)
except (RuntimeError, FileNotFoundError) as e:
_log(f"SYNC_FAIL {nb_name}: {e}")
return {
"status": "error",
"error": f"Jupytext sync failed: {e}",
"duration_s": time.time() - start,
"n_cells": 0,
}
ipynb_path = tmp_ipynb
# Executed notebook output path
executed_path = py_path.parent / f"_executed_{py_path.stem}.ipynb"
# Setup environment - Papermill's execute_notebook inherits os.environ,
# so we temporarily set env vars and restore them after execution.
saved_env = {}
env_vars = {
"MPLBACKEND": "Agg",
"PLOTLY_RENDERER": "json",
"DISABLE_HPO": "1",
}
if output_dir:
output_dir.mkdir(parents=True, exist_ok=True)
env_vars["ML4T_OUTPUT_DIR"] = str(output_dir)
if data_dir:
env_vars["ML4T_DATA_PATH"] = str(data_dir)
if extra_env:
env_vars.update({key: str(value) for key, value in extra_env.items()})
# PYTHONPATH includes repo root for utils imports + notebook dir for sibling imports
existing = os.environ.get("PYTHONPATH", "")
nb_dir = str(py_path.parent.resolve())
env_vars["PYTHONPATH"] = (
f"{REPO_ROOT}:{nb_dir}:{existing}" if existing else f"{REPO_ROOT}:{nb_dir}"
)
# Ensure torch's bundled CUDA libraries are found before system ones.
# The system libcudart.so.12 may be outdated and missing symbols like
# cudaGetDriverEntryPointByVersion that torch's bundled version provides.
try:
import torch
torch_lib = str(Path(torch.__file__).parent / "lib")
nvidia_libs = list((Path(torch.__file__).parent.parent / "nvidia").glob("*/lib"))
cuda_paths = [torch_lib] + [str(p) for p in nvidia_libs]
existing_ld = os.environ.get("LD_LIBRARY_PATH", "")
env_vars["LD_LIBRARY_PATH"] = ":".join(cuda_paths + [existing_ld])
except ImportError:
pass
remove_vars = ["TEST", "QUICK_TEST"]
if extra_env:
for key in extra_env:
if key in remove_vars:
remove_vars.remove(key)
# Apply environment changes
for key, value in env_vars.items():
saved_env[key] = os.environ.get(key)
os.environ[key] = value
for key in remove_vars:
saved_env[key] = os.environ.pop(key, None)
# Cell-level progress — always log to /tmp/ml4t-pm-{name}.log for visibility.
# Since request_save_on_cell_execute=True, the executed notebook is updated
# after each cell. Monitor it with: watch -n5 'python -c "import json; ..."'
progress_log = Path(f"/tmp/ml4t-pm-{nb_name}.log")
n_cells = 0
try:
with open(progress_log, "w") as pf:
pf.write(f"START {nb_name} timeout={timeout}s params={parameters}\n")
pm.execute_notebook(
str(ipynb_path),
str(executed_path),
parameters=parameters or {},
cwd=str(cwd or REPO_ROOT),
kernel_name="python3",
execution_timeout=timeout,
request_save_on_cell_execute=True,
progress_bar=False,
log_output=True,
)
# Count cells in executed notebook
try:
import nbformat
nb = nbformat.read(str(executed_path), as_version=4)
n_cells = len([c for c in nb.cells if c.cell_type == "code"])
except Exception:
pass
elapsed = time.time() - start
msg = f"OK {nb_name} ({elapsed:.1f}s, {n_cells} cells)"
_log(msg)
with open(progress_log, "a") as pf:
pf.write(f"{msg}\n")
return {"status": "ok", "error": None, "duration_s": elapsed, "n_cells": n_cells}
except pm.PapermillExecutionError as e:
elapsed = time.time() - start
msg = f"FAIL {nb_name} cell {e.cell_index} ({e.ename}): {e.evalue} ({elapsed:.1f}s)"
_log(msg)
with open(progress_log, "a") as pf:
pf.write(f"{msg}\n")
return {
"status": "error",
"error": f"Cell {e.cell_index} ({e.ename}): {e.evalue}",
"duration_s": elapsed,
"n_cells": e.cell_index,
}
except Exception as e:
elapsed = time.time() - start
_log(f"FAIL {nb_name}: {e} ({elapsed:.1f}s)")
return {"status": "error", "error": str(e), "duration_s": elapsed, "n_cells": 0}
finally:
# Restore environment
for key, value in saved_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
# Clean up temp input notebook
if tmp_ipynb is not None and tmp_ipynb.exists():
tmp_ipynb.unlink()
# Clean up executed notebook
if executed_path.exists():
executed_path.unlink()
def collect_chapter_notebooks(repo_root: Path, chapter_range: range) -> list[Path]:
"""Collect all teaching notebooks from chapter directories.
NOTE: Review repo has flat layout — notebooks live directly in
chapter directories (e.g., 05_synthetic_data/01_timegan.py),
NOT in code/ subdirectories.
Args:
repo_root: Repository root path
chapter_range: Range of chapter numbers to include
Returns:
Sorted list of .py notebook paths
"""
notebooks = []
for chapter_dir in sorted(repo_root.glob("[0-9][0-9]_*")):
if not chapter_dir.is_dir():
continue
# Extract chapter number from directory name
try:
ch_num = int(chapter_dir.name[:2])
except ValueError:
continue
if ch_num not in chapter_range:
continue
# Review repo: notebooks are directly in the chapter directory
for notebook in sorted(chapter_dir.glob("*.py")):
# Skip non-notebook files — use startswith to avoid false positives
# (e.g., "test_" must not match "backtest_")
if any(
notebook.name.startswith(prefix)
for prefix in [
"test_",
"conftest",
"extract_book_figures",
"export_figures",
"batch_",
]
) or any(x in notebook.name for x in ["__pycache__", "__init__"]):
continue
# Skip archived/draft/reserved directories
if any(
x in str(notebook)
for x in ["_archive", "archived", "drafts", "inventory", "_reserved"]
):
continue
# Skip helper/utility files (start with _)
if notebook.name.startswith("_"):
continue
if not notebook.name[0].isdigit() and not notebook.with_suffix(".ipynb").exists():
continue
notebooks.append(notebook)
return notebooks
tests/sample_registry_for_tests.py
+197
View File
@@ -0,0 +1,197 @@
"""Sample real registry.db data into test intermediates.
Copies a representative subset from each case study's production registry
into the test-data repo. This gives insight/synthesis/strategy_analysis
notebooks real data to work with in CI.
Sampling strategy:
- Model-side tables (training_runs, prediction_sets, prediction_metrics,
fold_metrics): copied in full — small enough.
- Backtest tables: top N per (family × stage) by Sharpe, plus ALL holdout
backtests. Includes corresponding backtest_fold_metrics.
Usage:
uv run python tests/sample_registry_for_tests.py
Writes to: ~/ml4t/test-data/intermediates/{cs}/run_log/registry.db
"""
import contextlib
import sqlite3
from pathlib import Path
REPO_ROOT = Path(__file__).parent.parent
CODE_CS_DIR = REPO_ROOT / "case_studies"
TEST_DATA_ROOT = Path.home() / "ml4t" / "test-data"
INTERMEDIATES_DIR = TEST_DATA_ROOT / "intermediates"
CASE_STUDY_IDS = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
# Keep top N backtests per (family, stage) by absolute Sharpe
TOP_N_PER_GROUP = 3
def _copy_rows(src, dst, table: str, rows: list) -> int:
"""Insert rows into dst table with proper column quoting."""
if not rows:
return 0
cols = [d[0] for d in src.execute(f"SELECT * FROM {table} LIMIT 1").description]
quoted = [f'"{c}"' for c in cols]
ph = ",".join(["?"] * len(cols))
dst.executemany(f"INSERT OR IGNORE INTO {table} ({','.join(quoted)}) VALUES ({ph})", rows)
return len(rows)
def sample_registry(cs_id: str) -> dict:
"""Sample from production registry into test intermediates. Returns stats."""
src_db = CODE_CS_DIR / cs_id / "run_log" / "registry.db"
if not src_db.exists():
return {"status": "SKIP", "reason": "no source registry.db"}
dst_dir = INTERMEDIATES_DIR / cs_id / "run_log"
dst_dir.mkdir(parents=True, exist_ok=True)
dst_db = dst_dir / "registry.db"
# Remove old DB to start fresh
dst_db.unlink(missing_ok=True)
src = sqlite3.connect(str(src_db))
try:
dst = sqlite3.connect(str(dst_db))
try:
return _populate_sample_db(src, dst, dst_db)
finally:
dst.close()
finally:
src.close()
def _populate_sample_db(src, dst, dst_db) -> dict:
stats: dict = {}
# 1. Copy schema from source (dump CREATE statements)
schema_sql = []
for row in src.execute(
"SELECT sql FROM sqlite_master WHERE type='table' AND sql IS NOT NULL"
).fetchall():
schema_sql.append(row[0])
for sql in schema_sql:
dst.execute(sql)
# Also copy indexes
for row in src.execute(
"SELECT sql FROM sqlite_master WHERE type='index' AND sql IS NOT NULL"
).fetchall():
with contextlib.suppress(sqlite3.OperationalError):
dst.execute(row[0])
# 2. Copy model-side tables in full
for table in ["training_runs", "prediction_sets", "prediction_metrics", "fold_metrics"]:
rows = src.execute(f"SELECT * FROM {table}").fetchall()
n = _copy_rows(src, dst, table, rows)
stats[table] = n
# 3. Sample backtests: top N per (family, stage) by |Sharpe|, plus all holdout
# First, get sampled backtest hashes
sampled_bt_hashes = set()
# 3a. Top N per family × stage (validation backtests)
top_n_sql = """
WITH ranked AS (
SELECT
b.backtest_hash,
b.stage,
t.family,
bm.sharpe,
ROW_NUMBER() OVER (
PARTITION BY b.stage, t.family
ORDER BY ABS(bm.sharpe) DESC
) AS rn
FROM backtest_runs b
JOIN backtest_metrics bm ON b.backtest_hash = bm.backtest_hash
JOIN prediction_sets p ON b.prediction_hash = p.prediction_hash
JOIN training_runs t ON p.training_hash = t.training_hash
WHERE p.split != 'holdout'
)
SELECT backtest_hash FROM ranked WHERE rn <= ?
"""
for row in src.execute(top_n_sql, (TOP_N_PER_GROUP,)).fetchall():
sampled_bt_hashes.add(row[0])
# 3b. ALL holdout backtests
holdout_sql = """
SELECT b.backtest_hash
FROM backtest_runs b
JOIN prediction_sets p ON b.prediction_hash = p.prediction_hash
WHERE p.split = 'holdout'
"""
for row in src.execute(holdout_sql).fetchall():
sampled_bt_hashes.add(row[0])
stats["backtest_runs_sampled"] = len(sampled_bt_hashes)
# 3c. Copy sampled backtest data (runs, metrics, fold_metrics)
if sampled_bt_hashes:
hash_list = list(sampled_bt_hashes)
batch_size = 500
for table in ["backtest_runs", "backtest_metrics", "backtest_fold_metrics"]:
count = 0
for i in range(0, len(hash_list), batch_size):
batch = hash_list[i : i + batch_size]
placeholders = ",".join(["?"] * len(batch))
rows = src.execute(
f"SELECT * FROM {table} WHERE backtest_hash IN ({placeholders})",
batch,
).fetchall()
count += _copy_rows(src, dst, table, rows)
stats[table] = count
dst.commit()
stats["file_size_kb"] = dst_db.stat().st_size // 1024
stats["status"] = "OK"
return stats
def main():
print(f"Sampling registries from {CODE_CS_DIR}")
print(f"Writing to {INTERMEDIATES_DIR}")
print(f"Top {TOP_N_PER_GROUP} backtests per (family × stage) + all holdout\n")
total_size = 0
for cs_id in CASE_STUDY_IDS:
print(f"--- {cs_id} ---")
stats = sample_registry(cs_id)
if stats["status"] != "OK":
print(f" {stats['status']}: {stats.get('reason', '')}")
continue
for table in [
"training_runs",
"prediction_sets",
"prediction_metrics",
"fold_metrics",
"backtest_runs",
"backtest_metrics",
"backtest_fold_metrics",
]:
print(f" {table:30s} {stats.get(table, 0):>6}")
print(f" {'file size (KB)':30s} {stats['file_size_kb']:>6}")
total_size += stats["file_size_kb"]
print(f"\nTotal registry size: {total_size} KB ({total_size / 1024:.1f} MB)")
if __name__ == "__main__":
main()
tests/test_allocation.py
+355
View File
@@ -0,0 +1,355 @@
"""Tests for case_studies/utils/allocation.py — portfolio weight contracts.
These allocators sit between model predictions and the backtest engine.
A silent regression here would corrupt every Ch17+ result.
The tests pin *structural* contracts, not exact numerical values:
- Long-only weights sum to 1 per timestamp
- Long-only weights are all non-negative
- Long/short weights are dollar-neutral (net ≈ 0) with gross leverage ≈ 2
(inverse-vol / risk-parity / HRP) or gross ≈ 1 (MVO)
- Exactly ``top_k`` assets per side are selected when enough assets exist
- Output columns are ``[timestamp, symbol, weight]`` in the expected order
Exact MVO values come from SLSQP and may vary across scipy versions, so the
numeric pins are loose (gross, net, count) rather than per-asset weights.
The ``synthetic_panel`` fixture builds 8 symbols × 300 dates of random-walk
prices. MVO needs a full lookback window (126 days); HRP needs ``vol_window``
(63 days). The fixture gives both allocators enough runway before the
rebalance timestamps.
"""
from __future__ import annotations
import numpy as np
import polars as pl
import pytest
from case_studies.utils.allocation import (
compute_hrp_weights,
compute_inverse_vol_weights,
compute_mvo_weights,
compute_risk_parity_weights,
)
# -----------------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------------
@pytest.fixture(scope="module")
def synthetic_panel() -> tuple[pl.DataFrame, pl.DataFrame]:
"""Return (predictions, prices) for 8 symbols × 300 days with 3 rebalance dates.
Prices: geometric random walk with asset-specific vol (so inverse-vol
produces distinguishable weights). Scores are ascending by symbol id
(S0..S7) so top_k picks deterministically.
"""
rng = np.random.default_rng(42)
n_symbols = 8
n_dates = 300
symbols = [f"S{i}" for i in range(n_symbols)]
ts = pl.date_range(pl.date(2023, 1, 1), pl.date(2023, 12, 31), "1d", eager=True)[:n_dates]
vols = 0.005 + 0.005 * np.arange(n_symbols) / n_symbols # 0.5% to ~1%
shocks = rng.normal(0.0, vols[None, :], (n_dates, n_symbols))
prices = 100.0 * np.exp(np.cumsum(shocks, axis=0))
price_rows: list[dict] = []
for i, t in enumerate(ts):
for j, s in enumerate(symbols):
price_rows.append({"timestamp": t, "symbol": s, "close": float(prices[i, j])})
prices_df = pl.DataFrame(price_rows)
pred_dates = ts[-3:]
pred_rows: list[dict] = []
for t in pred_dates:
for j, s in enumerate(symbols):
pred_rows.append({"timestamp": t, "symbol": s, "y_score": float(j)})
predictions = pl.DataFrame(pred_rows)
return predictions, prices_df
# -----------------------------------------------------------------------------
# Shared contract checks
# -----------------------------------------------------------------------------
def _assert_output_shape(out: pl.DataFrame) -> None:
assert set(out.columns) == {"timestamp", "symbol", "weight"}
def _assert_long_only_sums_to_1(out: pl.DataFrame) -> None:
per_date = out.group_by("timestamp").agg(pl.col("weight").sum().alias("s")).sort("timestamp")
for s in per_date["s"].to_list():
assert abs(s - 1.0) < 1e-6, per_date
def _assert_non_negative(out: pl.DataFrame) -> None:
assert (out["weight"] < 0).sum() == 0
def _assert_dollar_neutral(out: pl.DataFrame, gross_target: float) -> None:
per_date = (
out.group_by("timestamp")
.agg(
net=pl.col("weight").sum(),
gross=pl.col("weight").abs().sum(),
)
.sort("timestamp")
)
for net, gross in zip(per_date["net"].to_list(), per_date["gross"].to_list(), strict=True):
assert abs(net) < 1e-6, f"long-short should net to 0, got {net}"
assert abs(gross - gross_target) < 1e-6, f"expected gross={gross_target}, got {gross}"
def _assert_top_k_selected(out: pl.DataFrame, top_k: int) -> None:
per_date = (
out.group_by("timestamp").agg(n=pl.col("symbol").count()).sort("timestamp")["n"].to_list()
)
for n in per_date:
assert n == top_k, f"expected {top_k} selected, got {n}"
# -----------------------------------------------------------------------------
# compute_inverse_vol_weights
# -----------------------------------------------------------------------------
def test_inverse_vol_long_only_contracts(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_inverse_vol_weights(predictions, prices, top_k=4)
_assert_output_shape(out)
_assert_long_only_sums_to_1(out)
_assert_non_negative(out)
_assert_top_k_selected(out, top_k=4)
def test_inverse_vol_picks_top_k_by_score(synthetic_panel) -> None:
"""Scores ascending S0..S7 → top 4 should be S4..S7."""
predictions, prices = synthetic_panel
out = compute_inverse_vol_weights(predictions, prices, top_k=4)
assert set(out["symbol"].unique().to_list()) == {"S4", "S5", "S6", "S7"}
def test_inverse_vol_long_short_is_dollar_neutral(synthetic_panel) -> None:
"""Long/short with top_k=3 → 3 longs @ +w_i, 3 shorts @ -w_j, gross≈2 (two sides of 1)."""
predictions, prices = synthetic_panel
out = compute_inverse_vol_weights(predictions, prices, top_k=3, long_short=True)
_assert_dollar_neutral(out, gross_target=2.0)
def test_inverse_vol_produces_nonuniform_weights(synthetic_panel) -> None:
"""Weights are 1/σ-normalized — selected assets have heterogeneous vols,
so weights must not collapse to equal-weight (0.25 for top_k=4).
"""
predictions, prices = synthetic_panel
out = compute_inverse_vol_weights(predictions, prices, top_k=4)
last_ts = out["timestamp"].max()
slice_ = out.filter(pl.col("timestamp") == last_ts)
weights = np.array(slice_["weight"].to_list())
# Range of weights should be nontrivial (> 1% spread)
assert weights.max() - weights.min() > 0.01
def test_inverse_vol_deterministic(synthetic_panel) -> None:
predictions, prices = synthetic_panel
a = compute_inverse_vol_weights(predictions, prices, top_k=4)
b = compute_inverse_vol_weights(predictions, prices, top_k=4)
assert a.sort("timestamp", "symbol").equals(b.sort("timestamp", "symbol"))
# -----------------------------------------------------------------------------
# compute_risk_parity_weights
# -----------------------------------------------------------------------------
def test_risk_parity_long_only_contracts(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_risk_parity_weights(predictions, prices, top_k=4)
_assert_output_shape(out)
_assert_long_only_sums_to_1(out)
_assert_non_negative(out)
_assert_top_k_selected(out, top_k=4)
def test_risk_parity_long_short_is_dollar_neutral(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_risk_parity_weights(predictions, prices, top_k=3, long_short=True)
_assert_dollar_neutral(out, gross_target=2.0)
def test_risk_parity_assigns_less_to_high_vol_than_inverse_vol(synthetic_panel) -> None:
"""Risk-parity uses 1/σ^1.5 (steeper penalty than inverse-vol's 1/σ).
High-vol assets should be relatively *less* weighted under risk-parity
than under inverse-vol.
"""
predictions, prices = synthetic_panel
iv = compute_inverse_vol_weights(predictions, prices, top_k=4)
rp = compute_risk_parity_weights(predictions, prices, top_k=4)
last_ts = iv["timestamp"].max()
iv_weights = dict(
zip(
iv.filter(pl.col("timestamp") == last_ts)["symbol"].to_list(),
iv.filter(pl.col("timestamp") == last_ts)["weight"].to_list(),
strict=True,
)
)
rp_weights = dict(
zip(
rp.filter(pl.col("timestamp") == last_ts)["symbol"].to_list(),
rp.filter(pl.col("timestamp") == last_ts)["weight"].to_list(),
strict=True,
)
)
# S7 is the highest-vol asset selected — risk-parity should weight it lower than inverse-vol
assert rp_weights["S7"] < iv_weights["S7"]
def test_risk_parity_deterministic(synthetic_panel) -> None:
predictions, prices = synthetic_panel
a = compute_risk_parity_weights(predictions, prices, top_k=4)
b = compute_risk_parity_weights(predictions, prices, top_k=4)
assert a.sort("timestamp", "symbol").equals(b.sort("timestamp", "symbol"))
# -----------------------------------------------------------------------------
# compute_hrp_weights
# -----------------------------------------------------------------------------
def test_hrp_long_only_contracts(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_hrp_weights(predictions, prices, top_k=4)
_assert_output_shape(out)
_assert_long_only_sums_to_1(out)
_assert_non_negative(out)
_assert_top_k_selected(out, top_k=4)
def test_hrp_long_short_is_dollar_neutral(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_hrp_weights(predictions, prices, top_k=3, long_short=True)
_assert_dollar_neutral(out, gross_target=2.0)
def test_hrp_falls_back_to_equal_weight_on_short_history() -> None:
"""With <20 days of history, HRP cannot form a covariance matrix → equal-weight."""
rng = np.random.default_rng(0)
n_dates = 10 # well under the 20-obs floor
ts = pl.date_range(pl.date(2023, 1, 1), pl.date(2023, 1, 10), "1d", eager=True)
price_rows = []
for i, t in enumerate(ts):
for j, s in enumerate(["A", "B", "C", "D"]):
price_rows.append({"timestamp": t, "symbol": s, "close": 100.0 + float(rng.normal())})
prices = pl.DataFrame(price_rows)
predictions = pl.DataFrame(
{
"timestamp": [ts[-1]] * 4,
"symbol": ["A", "B", "C", "D"],
"y_score": [0.0, 1.0, 2.0, 3.0],
}
)
out = compute_hrp_weights(predictions, prices, top_k=4)
# Equal-weight = 1/4 for each
for w in out["weight"].to_list():
assert abs(w - 0.25) < 1e-9
def test_hrp_deterministic(synthetic_panel) -> None:
predictions, prices = synthetic_panel
a = compute_hrp_weights(predictions, prices, top_k=4)
b = compute_hrp_weights(predictions, prices, top_k=4)
assert a.sort("timestamp", "symbol").equals(b.sort("timestamp", "symbol"))
# -----------------------------------------------------------------------------
# compute_mvo_weights
# -----------------------------------------------------------------------------
def test_mvo_long_only_contracts(synthetic_panel) -> None:
predictions, prices = synthetic_panel
out = compute_mvo_weights(predictions, prices, top_k=4, max_weight=0.5)
_assert_output_shape(out)
_assert_non_negative(out)
# Some assets may be dropped from the output if their optimal weight is
# below 1e-6; don't pin the count. Weights should still sum to ~1.
per_date = out.group_by("timestamp").agg(pl.col("weight").sum().alias("s"))
for s in per_date["s"].to_list():
assert abs(s - 1.0) < 1e-6
def test_mvo_long_short_gross_normalized_to_1_and_dollar_neutral(synthetic_panel) -> None:
"""MVO long/short normalizes gross to 1 and uses a dollar-neutral constraint."""
predictions, prices = synthetic_panel
out = compute_mvo_weights(predictions, prices, top_k=3, long_short=True, max_weight=0.5)
_assert_dollar_neutral(out, gross_target=1.0)
def test_mvo_respects_position_cap(synthetic_panel) -> None:
"""No weight should exceed max_weight after normalization (long-only path)."""
predictions, prices = synthetic_panel
out = compute_mvo_weights(predictions, prices, top_k=8, max_weight=0.15)
# Small tolerance for renormalization + float error
assert out["weight"].max() <= 0.15 + 5e-3
def test_mvo_falls_back_to_equal_weight_on_short_history() -> None:
"""With 20 dates of history but lookback=126, MVO falls back to equal weight."""
rng = np.random.default_rng(0)
ts = pl.date_range(pl.date(2023, 1, 1), pl.date(2023, 1, 20), "1d", eager=True)
price_rows = []
for t in ts:
for s in ["A", "B", "C", "D"]:
price_rows.append({"timestamp": t, "symbol": s, "close": 100.0 + float(rng.normal())})
prices = pl.DataFrame(price_rows)
predictions = pl.DataFrame(
{
"timestamp": [ts[-1]] * 4,
"symbol": ["A", "B", "C", "D"],
"y_score": [0.0, 1.0, 2.0, 3.0],
}
)
out = compute_mvo_weights(predictions, prices, top_k=4)
for w in out["weight"].to_list():
assert abs(w - 0.25) < 1e-9
def test_mvo_returns_empty_frame_when_fewer_than_3_assets_selected() -> None:
"""top_k=2 → <3 assets → MVO skips the date and emits an empty frame."""
rng = np.random.default_rng(0)
ts = pl.date_range(pl.date(2023, 1, 1), pl.date(2023, 12, 31), "1d", eager=True)[:200]
price_rows = []
for t in ts:
for s in ["A", "B"]:
price_rows.append({"timestamp": t, "symbol": s, "close": 100.0 + float(rng.normal())})
prices = pl.DataFrame(price_rows)
predictions = pl.DataFrame(
{"timestamp": [ts[-1]] * 2, "symbol": ["A", "B"], "y_score": [0.0, 1.0]}
)
out = compute_mvo_weights(predictions, prices, top_k=2)
assert out.height == 0
assert set(out.columns) == {"timestamp", "symbol", "weight"}
# -----------------------------------------------------------------------------
# Input flexibility: accept either 'close' or 'ret' column in prices
# -----------------------------------------------------------------------------
def test_inverse_vol_accepts_ret_column_directly(synthetic_panel) -> None:
"""If prices already carry 'ret', the allocator uses it instead of pct_change('close')."""
predictions, prices = synthetic_panel
ret_prices = (
prices.sort("timestamp", "symbol")
.with_columns(ret=pl.col("close").pct_change().over("symbol"))
.select("timestamp", "symbol", "ret")
)
out = compute_inverse_vol_weights(predictions, ret_prices, top_k=4)
_assert_long_only_sums_to_1(out)
_assert_top_k_selected(out, top_k=4)
tests/test_artifact_specs.py
+57
View File
@@ -0,0 +1,57 @@
from __future__ import annotations
from case_studies.utils.backtest_loaders import (
get_backtest_config,
load_backtest_prices,
)
from utils.modeling import load_modeling_dataset
def test_us_equities_pilot_helpers_preserve_current_outputs() -> None:
bt = get_backtest_config("us_equities_panel")
prices = load_backtest_prices("us_equities_panel", max_symbols=2)
mds = load_modeling_dataset("us_equities_panel", "fwd_ret_1d", max_symbols=2)
assert bt.primary_label == "fwd_ret_1d"
assert bt.label_buffer == "1D"
assert bt.calendar == "NYSE"
assert bt.cadence == "daily_close"
assert prices.columns == ["symbol", "timestamp", "open", "high", "low", "close", "volume"]
assert prices["symbol"].n_unique() == 2
assert mds.label_col == "fwd_ret_1d"
assert mds.date_col == "timestamp"
assert mds.entity_cols == ["symbol"]
assert mds.join_cols == ["symbol", "timestamp"]
assert len(mds.feature_names) == 72
assert len(mds.splits) == 16
assert mds.label_buffer == "1D"
assert mds.task_type == "regression"
def test_microstructure_pilot_helpers_preserve_current_outputs() -> None:
bt = get_backtest_config("nasdaq100_microstructure")
prices = load_backtest_prices("nasdaq100_microstructure", max_symbols=2)
mds = load_modeling_dataset("nasdaq100_microstructure", "fwd_ret_15m", max_symbols=2)
assert bt.primary_label == "fwd_ret_15m"
assert bt.label_buffer == "15min"
assert bt.calendar == "NYSE"
assert bt.cadence == "15_minute"
# Microstructure carries OHLCV + bid/ask OHLC so the backtest engine can
# cost spread-aware fills.
required_cols = ["symbol", "timestamp", "open", "high", "low", "close", "volume"]
assert all(c in prices.columns for c in required_cols)
assert "bid_close" in prices.columns and "ask_close" in prices.columns
assert prices["symbol"].n_unique() == 2
assert mds.label_col == "fwd_ret_15m"
assert mds.date_col == "timestamp"
assert mds.entity_cols == ["symbol"]
assert mds.join_cols == ["symbol", "timestamp"]
assert len(mds.feature_names) == 88
assert len(mds.splits) == 2
assert mds.label_buffer == "15min"
assert mds.task_type == "regression"
tests/test_backtest_presets.py
+249
View File
@@ -0,0 +1,249 @@
from __future__ import annotations
from datetime import datetime
import polars as pl
from case_studies.utils.backtest_loaders import get_backtest_config, load_backtest_prices
from case_studies.utils.backtest_presets import (
cost_view,
ensure_backtest_spec,
is_backtest_spec,
load_backtest_preset,
strategy_view,
)
from case_studies.utils.backtest_runner import normalize_prediction_columns
from case_studies.utils.registry.specs import backtest_hash_from_parts
from case_studies.utils.registry.store import _infer_stage
def test_etf_backtest_base_preset_exists() -> None:
preset = load_backtest_preset("etfs")
assert preset["calendar"]["calendar"] == "NYSE"
assert preset["commission"]["rate"] == 0.0006
assert preset["slippage"]["rate"] == 0.0004
def test_ensure_backtest_spec_builds_composite_spec() -> None:
bt = get_backtest_config("etfs")
prices = load_backtest_prices("etfs", max_symbols=2)
legacy_spec = {
"chapter": "ch18",
"signal": {"method": "equal_weight_top_k", "top_k": 10, "long_short": False},
"execution": {
"mode": "engine",
"engine_preset": "realistic",
"cadence": bt.cadence,
"fill_timing": bt.execution_delay.upper(),
},
"costs": {"commission_bps": 6.0, "slippage_bps": 4.0},
"allocation": {"method": "risk_parity", "top_k": 10},
}
spec = ensure_backtest_spec(
"etfs",
bt,
legacy_spec,
prices=prices,
prediction_hash="pred123",
initial_cash=1_000_000.0,
)
assert is_backtest_spec(spec)
assert spec["strategy"]["signal"]["method"] == "equal_weight_top_k"
assert spec["strategy"]["allocation"]["method"] == "risk_parity"
assert spec["strategy"]["rebalance"]["cadence"] == bt.cadence
assert spec["backtest_config"]["commission"]["rate"] == 0.0006
assert spec["backtest_config"]["slippage"]["rate"] == 0.0004
assert spec["backtest_config"]["metadata"]["prediction_hash"] == "pred123"
assert cost_view(spec) == {"commission_bps": 6.0, "slippage_bps": 4.0}
def test_backtest_hash_changes_with_resolved_config() -> None:
bt = get_backtest_config("etfs")
prices = load_backtest_prices("etfs", max_symbols=2)
base = {
"signal": {"method": "equal_weight_top_k", "top_k": 10, "long_short": False},
"execution": {
"mode": "engine",
"engine_preset": "realistic",
"cadence": bt.cadence,
"fill_timing": bt.execution_delay.upper(),
},
"costs": {"commission_bps": 6.0, "slippage_bps": 4.0},
}
cheap = ensure_backtest_spec(
"etfs",
bt,
base,
prices=prices,
prediction_hash="pred123",
initial_cash=1_000_000.0,
)
expensive = ensure_backtest_spec(
"etfs",
bt,
{**base, "costs": {"commission_bps": 10.0, "slippage_bps": 4.0}},
prices=prices,
prediction_hash="pred123",
initial_cash=1_000_000.0,
)
assert backtest_hash_from_parts("pred123", cheap) != backtest_hash_from_parts(
"pred123", expensive
)
def test_stage_inference_supports_v2_specs() -> None:
spec = {
"version": 2,
"chapter": "ch19",
"strategy": {
"signal": {"method": "equal_weight_top_k", "top_k": 10},
"rebalance": {"mode": "engine", "cadence": "monthly_month_end"},
"risk": {"name": "trailing", "position_rules": [{"type": "trailing_stop"}]},
},
"backtest_config": {"commission": {"rate": 0.0006}, "slippage": {"rate": 0.0004}},
}
assert _infer_stage(spec) == "risk_overlay"
assert strategy_view(spec)["risk"]["name"] == "trailing"
def test_microstructure_preset_auto_loads_quote_columns() -> None:
prices = load_backtest_prices("nasdaq100_microstructure", max_symbols=2)
assert "bid_open" in prices.columns
assert "ask_open" in prices.columns
def test_sp500_options_short_only_engine_spec_preserves_explicit_feed() -> None:
bt = get_backtest_config("sp500_options")
prices = load_backtest_prices("sp500_options", max_symbols=2)
legacy_spec = {
"chapter": "ch16",
"signal": {
"method": "score_weighted_top_k",
"top_k": 10,
"long_short": False,
"direction": "short_only",
},
"execution": {
"mode": "engine",
"engine_preset": "realistic",
"cadence": bt.cadence,
"fill_timing": bt.execution_delay.upper(),
},
"costs": {"commission_bps": bt.commission_bps, "slippage_bps": bt.slippage_bps},
}
spec = ensure_backtest_spec(
"sp500_options",
bt,
legacy_spec,
prices=prices,
prediction_hash="pred123",
initial_cash=1_000_000.0,
)
cfg = spec["backtest_config"]
assert spec["strategy"]["rebalance"]["mode"] == "engine"
assert cfg["account"]["allow_short_selling"] is True
assert cfg["execution"]["execution_price"] == "quote_side"
assert cfg["execution"]["mark_price"] == "price"
assert cfg["feed"]["price_col"] == "instr_mid"
assert cfg["feed"]["bid_col"] == "instr_bid"
assert cfg["feed"]["ask_col"] == "instr_ask"
def test_ensure_backtest_spec_pins_enforce_sessions_for_cme_calendar() -> None:
"""CME-calendar specs must set enforce_sessions=True at construction time.
Regression: without this, ensure_backtest_spec emits a runtime with
enforce_sessions=False, but _run_engine later mutates it to True, so the
plan-time and post-engine hashes diverge (verify fires "0/N in_registry").
BacktestConfig.to_dict() does NOT serialize enforce_sessions; the hash
picks it up through ``_runtime_backtest_config`` in the spec (canonical_json
uses default=str on the dataclass repr), so the runtime is the load-bearing
surface and what we pin here.
Also pins that the NYSE projection branch does NOT trip the re-serialize
path — its original ``backtest_config`` dict and metadata are preserved.
"""
bt_cme = get_backtest_config("cme_futures")
prices_cme = load_backtest_prices("cme_futures", max_symbols=2)
canonical_cme = {
"version": 2,
"chapter": "ch18",
"strategy": {
"signal": {"method": "equal_weight_top_k", "top_k": 5, "long_short": True},
"rebalance": {"mode": "engine", "cadence": bt_cme.cadence},
},
"backtest_config": {
"commission": {"rate": 0.0},
"slippage": {"rate": 0.0},
"metadata": {"chapter": "ch18", "extra": "preserved"},
},
}
spec_cme = ensure_backtest_spec(
"cme_futures",
bt_cme,
canonical_cme,
prices=prices_cme,
prediction_hash="pred_cme",
initial_cash=1_000_000.0,
)
assert spec_cme["_runtime_backtest_config"].enforce_sessions is True
# Caller-supplied metadata keys must survive the re-serialization.
assert spec_cme["backtest_config"]["metadata"]["extra"] == "preserved"
assert spec_cme["backtest_config"]["metadata"]["prediction_hash"] == "pred_cme"
# NYSE specs (etfs) must NOT trip the projection re-serialize path —
# the original backtest_config dict should be preserved untouched.
bt_etf = get_backtest_config("etfs")
prices_etf = load_backtest_prices("etfs", max_symbols=2)
canonical_etf = {
"version": 2,
"chapter": "ch18",
"strategy": {
"signal": {"method": "equal_weight_top_k", "top_k": 10, "long_short": False},
"rebalance": {"mode": "engine", "cadence": bt_etf.cadence},
},
"backtest_config": {
"commission": {"rate": 0.0006},
"slippage": {"rate": 0.0004},
"metadata": {"chapter": "ch18", "extra": "preserved"},
},
}
spec_etf = ensure_backtest_spec(
"etfs",
bt_etf,
canonical_etf,
prices=prices_etf,
prediction_hash="pred_etf",
initial_cash=1_000_000.0,
)
assert spec_etf["_runtime_backtest_config"].enforce_sessions is False
# Original dict preserved (we did not re-serialize) — keys the caller
# passed in are exactly the keys present.
assert set(spec_etf["backtest_config"].keys()) == {"commission", "slippage", "metadata"}
assert spec_etf["backtest_config"]["metadata"]["extra"] == "preserved"
assert spec_etf["backtest_config"]["metadata"]["prediction_hash"] == "pred_etf"
def test_normalize_prediction_columns_maps_causal_and_legacy_fields() -> None:
df = pl.DataFrame(
{
"timestamp": [datetime(2024, 1, 2)],
"symbol": ["AAPL"],
"fold": [0],
"actual": [0.01],
"prediction": [0.02],
}
)
normalized = normalize_prediction_columns(df)
assert "y_score" in normalized.columns
assert "y_true" in normalized.columns
assert "fold_id" in normalized.columns
assert normalized["y_score"].to_list() == [0.02]
assert normalized["y_true"].to_list() == [0.01]
tests/test_backtest_runner_helpers.py
+227
View File
@@ -0,0 +1,227 @@
"""Regression tests for case_studies/utils/backtest_runner.py helpers.
Pins the P2.4 fixes from roborev jobs #2904, #2501, #2502, #2500:
- ``_align_symbol_dtype`` surfaces case-study context on ticker-vs-id mismatches.
- ``substitute_continuous_return_for_classification`` raises on duplicate
(timestamp, symbol) rows in the continuous-return parquet and on left-join
height changes.
- ``apply_universe_filter`` collapses sub-daily timestamps to the date grain
before computing the within-date rank.
- ``_MAX_NULL_RATE`` constant is wired through ``max_null_rate`` parameter.
"""
from __future__ import annotations
from datetime import datetime
from pathlib import Path
from textwrap import dedent
import polars as pl
import pytest
from case_studies.utils.backtest_runner import (
_MAX_NULL_RATE,
_align_symbol_dtype,
apply_universe_filter,
substitute_continuous_return_for_classification,
)
def test_max_null_rate_constant_default() -> None:
assert _MAX_NULL_RATE == 0.10
def test_align_symbol_dtype_same_dtype_passthrough() -> None:
target = pl.DataFrame({"symbol": ["A", "B"]})
other = pl.DataFrame({"symbol": ["C", "D"]})
out = _align_symbol_dtype(target, other, case_study="x", target_side="t", other_side="o")
assert out.schema["symbol"] == pl.Utf8
# Returned frame is the original when dtypes match.
assert out.equals(other)
def test_align_symbol_dtype_int_target_numeric_string_source() -> None:
target = pl.DataFrame({"symbol": [1, 2]}, schema={"symbol": pl.UInt32})
other = pl.DataFrame({"symbol": ["10", "20"]})
out = _align_symbol_dtype(
target, other, case_study="us_firm", target_side="weights", other_side="prices"
)
assert out.schema["symbol"] == pl.UInt32
assert out["symbol"].to_list() == [10, 20]
def test_align_symbol_dtype_int_target_ticker_source_raises_with_context() -> None:
target = pl.DataFrame({"symbol": [1, 2]}, schema={"symbol": pl.UInt32})
other = pl.DataFrame({"symbol": ["AAPL", "MSFT"]})
with pytest.raises(TypeError, match=r"case_study='broken'"):
_align_symbol_dtype(
target,
other,
case_study="broken",
target_side="weights",
other_side="prices",
)
def test_align_symbol_dtype_int_source_to_string_target() -> None:
target = pl.DataFrame({"symbol": ["A"]})
other = pl.DataFrame({"symbol": [1, 2]}, schema={"symbol": pl.UInt32})
out = _align_symbol_dtype(target, other, case_study="x", target_side="t", other_side="o")
assert out.schema["symbol"] == pl.Utf8
def test_apply_universe_filter_collapses_intraday_to_date_grain(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Sub-daily bars share a date but rank should be within-date, not within-bar.
Without the date-collapse fix, two intraday bars per (date, symbol) would
produce a denominator of 2N instead of N for the daily rank, silently
filtering against a within-bar universe.
"""
cs = "sp500_options_test"
cs_dir = tmp_path / cs / "config"
cs_dir.mkdir(parents=True)
(cs_dir / "setup.yaml").write_text(
dedent(
"""
backtest:
sweep:
htm_cost_cascade:
liquid_quantile: 0.50
"""
).strip()
)
import case_studies.utils.backtest_runner as br
monkeypatch.setattr(br, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
# ``CASE_STUDIES_DIR`` is imported lazily inside the function, so also
# patch the source module ``utils`` so the rebinding wins.
import utils as _utils # type: ignore
monkeypatch.setattr(_utils, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
# Two intraday bars per (date, symbol). Without date-collapse, rank
# denominator would be 4 (two bars × two symbols) and both symbols would
# land at the 0.50 quantile; with date-collapse, denominator is 2 (two
# symbols), and the tighter-spread symbol (A) is the unique survivor.
d1 = datetime(2024, 1, 2)
bar_open = datetime(2024, 1, 2, 9, 30)
bar_close = datetime(2024, 1, 2, 16, 0)
prices = pl.DataFrame(
{
"timestamp": [bar_open, bar_close, bar_open, bar_close],
"symbol": ["A", "A", "B", "B"],
"instr_rel_spread": [0.01, 0.012, 0.05, 0.06],
}
)
predictions = pl.DataFrame(
{
"timestamp": [d1, d1],
"symbol": ["A", "B"],
}
)
out = apply_universe_filter(
predictions, prices, case_study=cs, signal_config={"universe_filter": "liquid"}
)
# Only the tighter-spread symbol (A) survives the 0.50 quantile.
assert out["symbol"].to_list() == ["A"]
def test_substitute_continuous_return_dedupe_assertion(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
cs = "test_cs"
cs_dir = tmp_path / cs
(cs_dir / "config").mkdir(parents=True)
(cs_dir / "labels").mkdir()
(cs_dir / "config" / "setup.yaml").write_text(
dedent(
"""
labels:
classification_eval_label:
fwd_dir_1d: fwd_ret_1d
"""
).strip()
)
# Continuous-return parquet with a duplicate (timestamp, symbol) row.
d1 = datetime(2024, 1, 2)
eval_df = pl.DataFrame(
{
"timestamp": [d1, d1, d1], # 2× (d1, "A") — duplicate!
"symbol": ["A", "A", "B"],
"fwd_ret_1d": [0.01, 0.02, 0.03],
}
)
eval_df.write_parquet(cs_dir / "labels" / "fwd_ret_1d.parquet")
predictions = pl.DataFrame(
{
"timestamp": [d1, d1],
"symbol": ["A", "B"],
"y_score": [0.1, 0.2],
"y_true": [1, 0],
}
)
import case_studies.utils.backtest_runner as br
import utils as _utils # type: ignore
monkeypatch.setattr(_utils, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
monkeypatch.setattr(br, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
with pytest.raises(ValueError, match=r"duplicate \(timestamp, symbol\)"):
substitute_continuous_return_for_classification(
predictions, case_study=cs, label="fwd_dir_1d"
)
def test_substitute_continuous_return_max_null_rate_param(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Passing ``max_null_rate=1.0`` allows callers in a legitimately high-null regime."""
cs = "test_cs_nulls"
cs_dir = tmp_path / cs
(cs_dir / "config").mkdir(parents=True)
(cs_dir / "labels").mkdir()
(cs_dir / "config" / "setup.yaml").write_text(
dedent(
"""
labels:
classification_eval_label:
fwd_dir_1d: fwd_ret_1d
"""
).strip()
)
d1 = datetime(2024, 1, 2)
d2 = datetime(2024, 1, 3)
# Eval parquet only covers d1, not d2 — predictions on d2 will null-match.
eval_df = pl.DataFrame({"timestamp": [d1], "symbol": ["A"], "fwd_ret_1d": [0.01]})
eval_df.write_parquet(cs_dir / "labels" / "fwd_ret_1d.parquet")
predictions = pl.DataFrame(
{
"timestamp": [d1, d2, d2, d2],
"symbol": ["A", "A", "B", "C"],
"y_score": [0.1, 0.2, 0.3, 0.4],
"y_true": [1, 0, 1, 0],
}
)
import case_studies.utils.backtest_runner as br
import utils as _utils # type: ignore
monkeypatch.setattr(_utils, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
monkeypatch.setattr(br, "CASE_STUDIES_DIR", str(tmp_path), raising=False)
# Default cap (10%) raises: 3/4 = 75% null rate.
with pytest.raises(ValueError, match=r"exceeds max_null_rate"):
substitute_continuous_return_for_classification(
predictions, case_study=cs, label="fwd_dir_1d"
)
# Override loosens the cap; missing rows are dropped instead of raised.
out = substitute_continuous_return_for_classification(
predictions, case_study=cs, label="fwd_dir_1d", max_null_rate=1.0
)
assert out.height == 1
assert out["y_true"].to_list() == [0.01]
tests/test_backtest_schedule.py
+345
View File
@@ -0,0 +1,345 @@
"""Tests for calendar-aware backtest schedule resolution and execution delay mapping.
These tests validate the fixes for:
- Finding 1: Calendar-named cadences must use actual calendar dates, not elapsed time
- Finding 2: Execution delay mapping must be explicit, not substring-based
- Finding 3: Session enforcement must be enabled for CME
- Finding 4: Vectorized path must use resolved schedule, not gather_every
"""
from datetime import datetime
import polars as pl
import pytest
# ---------------------------------------------------------------------------
# resolve_rebalance_timestamps tests
# ---------------------------------------------------------------------------
def _make_weekday_series(start: str, end: str) -> pl.Series:
"""Create a daily timestamp series (weekdays only, simulating trading days)."""
dates = pl.date_range(
pl.lit(datetime.strptime(start, "%Y-%m-%d")),
pl.lit(datetime.strptime(end, "%Y-%m-%d")),
interval="1d",
eager=True,
)
# Filter to weekdays (Mon=1..Fri=5 in Polars)
df = pl.DataFrame({"ts": dates}).filter(pl.col("ts").dt.weekday() <= 5)
return df["ts"].sort()
class TestResolveRebalanceTimestamps:
"""Test calendar-aware schedule resolution."""
def test_monthly_month_end_returns_actual_month_ends(self):
"""ETF scenario: monthly_month_end must return last session of each month."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = _make_weekday_series("2016-01-01", "2016-06-30")
result = resolve_rebalance_timestamps(ts, "monthly_month_end")
all_dates = ts.to_list()
for dt in result.to_list():
month, year = dt.month, dt.year
later = [d for d in all_dates if d.year == year and d.month == month and d > dt]
assert len(later) == 0, f"Month-end {dt} is not the last session in {year}-{month:02d}"
def test_monthly_month_end_not_every_21_days(self):
"""Regression: month-end should NOT produce evenly-spaced 21-day intervals."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = _make_weekday_series("2016-01-01", "2016-12-31")
result = resolve_rebalance_timestamps(ts, "monthly_month_end")
assert 11 <= len(result) <= 12
gaps = result.diff().drop_nulls().dt.total_seconds() / 86400
gap_values = set(int(g) for g in gaps.to_list())
assert len(gap_values) > 1, "Month-end gaps should not all be identical"
def test_weekly_friday_returns_end_of_week(self):
"""CME/SP500 scenario: weekly_friday_close returns last session per ISO week."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = _make_weekday_series("2018-08-01", "2018-10-31")
result = resolve_rebalance_timestamps(ts, "weekly_friday_close")
# Exclude last week (may be incomplete if date range ends mid-week)
dates = result.to_list()
# All complete weeks should end on Friday
for dt in dates[:-1]:
assert dt.weekday() == 4, f"Expected Friday, got {dt} (weekday={dt.weekday()})"
def test_weekly_friday_holiday_fallback(self):
"""If Friday is missing (holiday), should take Thursday of that week."""
from datetime import date
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = _make_weekday_series("2018-08-27", "2018-09-14")
# Remove Friday 2018-09-07 (simulate holiday)
# ts contains date objects, so compare with date
friday_to_remove = date(2018, 9, 7)
ts_list = [d for d in ts.to_list() if d != friday_to_remove]
ts_filtered = pl.Series("ts", ts_list)
result = resolve_rebalance_timestamps(ts_filtered, "weekly_friday_close")
# The week of Sep 3-7 should still have a rebalance, but on Thursday Sep 6
week_36_dates = [dt for dt in result.to_list() if dt.isocalendar()[1] == 36]
assert len(week_36_dates) == 1
assert week_36_dates[0] == date(2018, 9, 6), (
f"Expected Thursday fallback, got {week_36_dates[0]}"
)
def test_daily_returns_all_timestamps(self):
"""Daily cadence should return every available timestamp."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = _make_weekday_series("2020-01-01", "2020-01-31")
for cadence in ("daily", "daily_close", "daily_ny_close"):
result = resolve_rebalance_timestamps(ts, cadence)
assert len(result) == len(ts), f"{cadence}: expected all {len(ts)} dates"
def test_eight_hour_returns_all_timestamps(self):
"""8-hour funding cadence should return all timestamps."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
dates = pl.datetime_range(
datetime(2020, 1, 1),
datetime(2020, 1, 31),
interval="8h",
eager=True,
)
result = resolve_rebalance_timestamps(dates, "8_hour_funding_aligned")
assert len(result) == len(dates.unique())
def test_empty_series(self):
"""Empty input should return empty output."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
ts = pl.Series("ts", [], dtype=pl.Datetime("us"))
result = resolve_rebalance_timestamps(ts, "monthly_month_end")
assert len(result) == 0
# ---------------------------------------------------------------------------
# Execution delay mapping tests
# ---------------------------------------------------------------------------
class TestExecutionDelayMapping:
"""Test explicit execution delay -> ExecutionMode mapping."""
def test_next_bar_open_maps_to_next_bar(self):
from ml4t.backtest import ExecutionMode
from case_studies.utils.backtest_presets import (
resolve_execution_mode as _resolve_execution_mode,
)
assert _resolve_execution_mode("NEXT_BAR_OPEN") == ExecutionMode.NEXT_BAR
assert _resolve_execution_mode("next_bar_open") == ExecutionMode.NEXT_BAR
def test_monday_open_maps_to_next_bar(self):
"""Regression: monday_open must NOT fall through to SAME_BAR."""
from ml4t.backtest import ExecutionMode
from case_studies.utils.backtest_presets import (
resolve_execution_mode as _resolve_execution_mode,
)
assert _resolve_execution_mode("MONDAY_OPEN") == ExecutionMode.NEXT_BAR
assert _resolve_execution_mode("monday_open") == ExecutionMode.NEXT_BAR
def test_1_bar_maps_to_next_bar(self):
from ml4t.backtest import ExecutionMode
from case_studies.utils.backtest_presets import (
resolve_execution_mode as _resolve_execution_mode,
)
assert _resolve_execution_mode("1_BAR") == ExecutionMode.NEXT_BAR
assert _resolve_execution_mode("1_bar") == ExecutionMode.NEXT_BAR
def test_at_funding_timestamp_maps_to_same_bar(self):
from ml4t.backtest import ExecutionMode
from case_studies.utils.backtest_presets import (
resolve_execution_mode as _resolve_execution_mode,
)
assert _resolve_execution_mode("AT_FUNDING_TIMESTAMP") == ExecutionMode.SAME_BAR
def test_unknown_token_raises_value_error(self):
"""Unknown execution delay must raise, not silently degrade."""
from case_studies.utils.backtest_presets import (
resolve_execution_mode as _resolve_execution_mode,
)
with pytest.raises(ValueError, match="Unknown execution delay"):
_resolve_execution_mode("SOME_RANDOM_TOKEN")
# ---------------------------------------------------------------------------
# Vectorized thinning tests
# ---------------------------------------------------------------------------
class TestThinToRebalanceDates:
"""Test that vectorized thinning uses calendar-aware schedule."""
def test_weekly_friday_keeps_fridays_not_every_5th(self):
"""Regression: sp500_options weekly_friday must keep actual Fridays."""
from case_studies.utils.backtest_loaders import thin_to_rebalance_dates
dates = _make_weekday_series("2020-01-01", "2020-02-28")
preds = pl.DataFrame(
{
"timestamp": dates,
"symbol": ["SPY"] * len(dates),
"y_score": [0.1] * len(dates),
"y_true": [0.01] * len(dates),
}
)
# fwd_ret_5d on weekly_friday schedule: step=1 (5d horizon <= 7d gap)
result = thin_to_rebalance_dates(preds, cadence="weekly_friday", step=1)
unique_dates = result["timestamp"].unique().sort()
# Exclude last date (may be incomplete week)
for dt in unique_dates.to_list()[:-1]:
assert dt.weekday() in (3, 4), f"Expected Thu/Fri, got {dt} (weekday={dt.weekday()})"
def test_monthly_month_end_keeps_month_ends(self):
"""Regression: ETFs monthly_month_end must keep actual month-end sessions."""
from case_studies.utils.backtest_loaders import thin_to_rebalance_dates
dates = _make_weekday_series("2016-01-01", "2016-06-30")
all_dates = dates.to_list()
preds = pl.DataFrame(
{
"timestamp": dates,
"symbol": ["SPY"] * len(dates),
"y_score": [0.1] * len(dates),
"y_true": [0.01] * len(dates),
}
)
# fwd_ret_21d on monthly_month_end schedule: step=1 (21d horizon <= 30d gap)
result = thin_to_rebalance_dates(preds, cadence="monthly_month_end", step=1)
unique_dates = result["timestamp"].unique().sort()
assert 5 <= len(unique_dates) <= 6
for dt in unique_dates.to_list():
month, year = dt.month, dt.year
later = [d for d in all_dates if d.year == year and d.month == month and d > dt]
assert len(later) == 0, f"{dt} is not the last trading day of {year}-{month:02d}"
# ---------------------------------------------------------------------------
# Rebalance-step lookup (declared in each case study's setup.yaml)
# ---------------------------------------------------------------------------
class TestGetRebalanceStep:
"""Verify that per-label thinning steps are read from setup.yaml.
Replaces the legacy regex-based implementation which silently returned
step=1 for any label without a digit-unit token (e.g., ``ret_to_expiry``),
producing 4-5× inflated Sharpe for overlapping-cohort strategies.
The step is now a design-time constant declared under
``labels.rebalance_step`` in each case study's setup.yaml.
"""
def test_sp500_options_ret_to_expiry_is_5(self):
"""Regression: ret_to_expiry must thin weekly_friday cohorts by 5.
30-day DTE on a 7-day schedule -> ceil(30/7) = 5. Pre-fix this
silently returned 1 (overlapping 5-cohort double-counting).
"""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("sp500_options", "ret_to_expiry") == 5
def test_cme_fwd_ret_21d_is_3(self):
"""cme_futures fwd_ret_21d on weekly_friday -> ceil(21/7) = 3."""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("cme_futures", "fwd_ret_21d") == 3
def test_us_firm_fwd_ret_1m_is_1(self):
"""Monthly label on monthly schedule -> step=1."""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("us_firm_characteristics", "fwd_ret_1m") == 1
def test_nasdaq100_fwd_ret_60m_is_4(self):
"""nasdaq100 fwd_ret_60m on 15-minute schedule -> ceil(60/15) = 4."""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("nasdaq100_microstructure", "fwd_ret_60m") == 4
def test_nasdaq100_fwd_ret_5m_is_1(self):
"""Regression: fwd_ret_5m on 15-minute schedule must stay at 1.
Pre-fix, the regex matched `(5, m)` and the old `n <= 12` branch
mis-read it as 5 MONTHS, computing step ~10,000 and collapsing
backtests to a handful of points.
"""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("nasdaq100_microstructure", "fwd_ret_5m") == 1
def test_crypto_fwd_ret_24h_is_3(self):
"""crypto 24h label on 8h schedule -> ceil(24/8) = 3."""
from case_studies.utils.backtest_loaders import get_rebalance_step
assert get_rebalance_step("crypto_perps_funding", "fwd_ret_24h") == 3
def test_unknown_label_raises(self):
"""Unknown label must raise KeyError pointing at setup.yaml."""
from case_studies.utils.backtest_loaders import get_rebalance_step
with pytest.raises(KeyError, match="rebalance_step"):
get_rebalance_step("sp500_options", "fwd_ret_unknown_label")
# ---------------------------------------------------------------------------
# Integration: engine schedule set membership
# ---------------------------------------------------------------------------
class TestEngineScheduleIntegration:
"""Verify the engine path builds correct schedule sets."""
def test_etf_monthly_schedule_matches_month_ends(self):
"""ETF backtest should rebalance on actual month-end sessions."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
dates = _make_weekday_series("2015-12-01", "2016-04-30")
schedule = resolve_rebalance_timestamps(dates, "monthly_month_end")
all_dates = dates.to_list()
for dt in schedule.to_list():
month, year = dt.month, dt.year
later = [d for d in all_dates if d.month == month and d.year == year and d > dt]
assert len(later) == 0, f"Rebalance {dt} is not month-end: later dates {later[:3]}"
def test_cme_weekly_schedule_matches_fridays(self):
"""CME backtest should rebalance on Friday sessions."""
from case_studies.utils.backtest_loaders import resolve_rebalance_timestamps
dates = _make_weekday_series("2018-08-01", "2018-10-26") # End on a Friday
schedule = resolve_rebalance_timestamps(dates, "weekly_friday_close")
for dt in schedule.to_list():
assert dt.weekday() == 4, (
f"CME rebalance {dt} should be Friday, got weekday={dt.weekday()}"
)
tests/test_case_studies.py
+156
View File
@@ -0,0 +1,156 @@
"""Test case study pipeline notebooks via Papermill parameter injection.
Each case study notebook runs independently against pre-generated intermediates
(labels, features, predictions, registries) stored in the test-data repo.
A failure in one notebook does NOT cascade to skip later notebooks.
Stages are auto-discovered: any [0-9][0-9]_*.py file in a case study
directory is treated as a pipeline stage.
Usage:
# All case studies
pytest tests/test_case_studies.py -v
# Specific case study
pytest tests/test_case_studies.py -v -k "etfs"
# Specific stage
pytest tests/test_case_studies.py -v -k "03_features"
"""
import re
from pathlib import Path
import pytest
from tests.pm_helpers import current_test_tier, get_overrides, get_tier, run_notebook
REPO_ROOT = Path(__file__).parent.parent
# All case studies
CASE_STUDIES = [
"etfs",
"crypto_perps_funding",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"fx_pairs",
"cme_futures",
"sp500_options",
"us_equities_panel",
]
# Pattern for numbered pipeline stages — allows optional single-letter suffix
# (e.g., 10a_pca, 11b_ipca) for per-estimator notebook splits.
_STAGE_RE = re.compile(r"^\d{2}[a-z]?_")
def _collect_case_study_tests():
"""Collect all case study pipeline notebooks as (case_study, stage, path) tuples.
Auto-discovers files matching ^\\d{2}[a-z]?_ in each case study directory,
sorted numerically. Skips helper files (starting with _).
"""
tests = []
for cs in CASE_STUDIES:
cs_dir = REPO_ROOT / "case_studies" / cs
if not cs_dir.exists():
continue
for notebook in sorted(cs_dir.glob("[0-9][0-9]*.py")):
if notebook.name.startswith("_"):
continue
if not _STAGE_RE.match(notebook.name):
continue
stage = notebook.stem # e.g., "06_linear" or "11a_pca"
tests.append((cs, stage, notebook))
return tests
CASE_STUDY_TESTS = _collect_case_study_tests()
print(f"Found {len(CASE_STUDY_TESTS)} case study pipeline notebooks to test")
@pytest.mark.parametrize(
"case_study,stage,notebook_path",
CASE_STUDY_TESTS,
ids=lambda *args: None, # Custom IDs below
)
def test_case_study_pipeline(
case_study, stage, notebook_path, populated_data_dir, seeded_output_dir
):
"""Execute a case study pipeline stage via Papermill.
Each notebook runs independently — intermediates (labels, features,
predictions, registries) are pre-generated in the test-data repo.
"""
# Check case-study-level skip (e.g., "case_studies/nasdaq100_microstructure")
cs_key = f"case_studies/{case_study}"
cs_overrides = get_overrides(cs_key)
if cs_overrides.get("skip"):
pytest.skip(f"Skipped: {cs_overrides.get('skip_reason', 'case study skipped')}")
rel_path = notebook_path.relative_to(REPO_ROOT).with_suffix("")
overrides = get_overrides(str(rel_path))
# Tier routing: skip when NB tier doesn't match the current run tier.
nb_tier = get_tier(overrides)
run_tier = current_test_tier()
if nb_tier != run_tier:
pytest.skip(f"Tier {nb_tier} — current run tier is {run_tier}")
# Skip if overrides say so
if overrides.get("skip"):
reason = overrides.get("skip_reason", "marked skip in overrides")
pytest.skip(f"Skipped: {reason}")
# Check required imports (e.g., gensim, signatory, duckdb)
requires = overrides.get("requires_import")
if requires:
pkg = requires if isinstance(requires, str) else requires[0]
try:
__import__(pkg)
except ImportError:
pytest.skip(f"Requires {pkg} (not installed in this Docker image)")
# Check GPU requirement
if overrides.get("gpu"):
try:
import torch
if not torch.cuda.is_available():
pytest.skip("GPU required but not available")
except ImportError:
pytest.skip("GPU required but torch not installed")
timeout = overrides.get("timeout", 300)
parameters = overrides.get("parameters", {})
result = run_notebook(
py_path=notebook_path,
parameters=parameters,
timeout=timeout,
output_dir=seeded_output_dir,
data_dir=populated_data_dir,
)
if result["status"] == "error":
pytest.fail(
f"\n{'=' * 70}\n"
f"Pipeline failed: {case_study}::{stage}\n"
f"{'=' * 70}\n"
f"Error: {result['error']}\n"
f"{'=' * 70}\n"
)
# Custom test IDs
def pytest_collection_modifyitems(items):
"""Set readable test IDs for case study tests."""
for item in items:
if "test_case_study_pipeline" in item.name and hasattr(item, "callspec"):
cs = item.callspec.params.get("case_study", "")
stage = item.callspec.params.get("stage", "")
item._nodeid = f"{item.parent.nodeid}::{cs}::{stage}"
tests/test_case_study_analytics.py
+495
View File
@@ -0,0 +1,495 @@
"""Tests for case_studies/utils/analytics.py — registry query contracts.
Per memory rule ``feedback_results_json``: "Registry only, never JSONs."
This module is the canonical path Ch6–Ch20 insights notebooks take to pull
IC / AUC / Sharpe numbers into the book. A silent regression in any of
these queries would corrupt every cross-chapter summary table.
The tests pin three layers:
1. **Metadata invariants** — the handful of dicts that declare the 9
case-study IDs must agree on keys, so a new case study can't be added
to one dict and forgotten in another.
2. **Path resolution** — ``_cs_dir`` / ``_registry_path`` honor
``ML4T_OUTPUT_DIR`` for test isolation.
3. **Query contracts** — against a seeded SQLite registry:
- ``load_model_ic`` filters by family, split, case_studies list and
returns the expected rows with a ``case_study`` label column.
- ``load_classification_metrics`` requires ``task_type = 'classification'``.
- ``load_best_ic_per_family`` picks max IC per (case_study, family)
pair and optionally restricts to the primary label.
- ``load_chapter_backtests("ch16")`` maps to ``stage="signal"`` and
joins backtest_runs × backtest_metrics × prediction_sets ×
training_runs.
- Spec helpers: ``extract_cost_bps`` sums commission + slippage
from either v1 or v2 backtest specs; ``extract_allocator`` reads
strategy.allocation.method.
"""
from __future__ import annotations
import json
import sqlite3
from pathlib import Path
import polars as pl
import pytest
from case_studies.utils import analytics
# -----------------------------------------------------------------------------
# Metadata invariants
# -----------------------------------------------------------------------------
def test_case_study_ids_match_metadata_keys() -> None:
assert list(analytics.CASE_STUDY_META.keys()) == analytics.CASE_STUDY_IDS
@pytest.mark.parametrize(
"dict_name",
["PRIMARY_LABELS", "SHORT_NAMES", "DATASET_META", "CADENCE_MAP", "DISPLAY_NAMES"],
)
def test_metadata_dict_keys_match_case_study_ids(dict_name) -> None:
"""Every metadata dict must enumerate the same 9 case studies — otherwise
cross-dict joins in load_best_ic_per_family / load_chapter_backtests drop
rows silently.
"""
d = getattr(analytics, dict_name)
assert set(d.keys()) == set(analytics.CASE_STUDY_IDS), (
f"{dict_name} keys differ: missing {set(analytics.CASE_STUDY_IDS) - set(d.keys())}, "
f"extra {set(d.keys()) - set(analytics.CASE_STUDY_IDS)}"
)
def test_primary_labels_are_non_empty_strings() -> None:
for cs, lbl in analytics.PRIMARY_LABELS.items():
assert isinstance(lbl, str) and lbl, f"{cs}: empty/non-string primary label"
# -----------------------------------------------------------------------------
# Path resolution
# -----------------------------------------------------------------------------
def test_cs_dir_production_path(monkeypatch) -> None:
"""With no ML4T_OUTPUT_DIR, _cs_dir falls back to REPO_ROOT/case_studies."""
monkeypatch.delenv("ML4T_OUTPUT_DIR", raising=False)
from utils.paths import REPO_ROOT
assert analytics._cs_dir() == REPO_ROOT / "case_studies"
def test_cs_dir_redirects_to_output_dir_when_registry_present(tmp_path, monkeypatch) -> None:
"""With ML4T_OUTPUT_DIR set AND a registry.db present under tmp, _cs_dir
returns the tmp root instead of the production case_studies path.
"""
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
(tmp_path / "etfs" / "run_log").mkdir(parents=True)
(tmp_path / "etfs" / "run_log" / "registry.db").touch()
assert analytics._cs_dir("etfs") == tmp_path
def test_cs_dir_falls_back_when_registry_missing_under_output_dir(tmp_path, monkeypatch) -> None:
"""ML4T_OUTPUT_DIR set but no registry.db under it → fall back to production."""
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
from utils.paths import REPO_ROOT
assert analytics._cs_dir("etfs") == REPO_ROOT / "case_studies"
def test_registry_path_is_three_levels_deep(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
(tmp_path / "etfs" / "run_log").mkdir(parents=True)
(tmp_path / "etfs" / "run_log" / "registry.db").touch()
p = analytics._registry_path("etfs")
assert p == tmp_path / "etfs" / "run_log" / "registry.db"
# -----------------------------------------------------------------------------
# _query behavior on empty / missing databases
# -----------------------------------------------------------------------------
def test_query_returns_empty_df_when_path_missing(tmp_path) -> None:
missing = tmp_path / "nope.db"
out = analytics._query(missing, "SELECT 1")
assert isinstance(out, pl.DataFrame)
assert out.is_empty()
def test_query_returns_empty_df_when_no_rows(tmp_path) -> None:
db = tmp_path / "empty.db"
conn = sqlite3.connect(str(db))
conn.execute("CREATE TABLE t (x INTEGER)")
conn.commit()
conn.close()
out = analytics._query(db, "SELECT * FROM t")
assert out.is_empty()
def test_query_returns_populated_df(tmp_path) -> None:
db = tmp_path / "rows.db"
conn = sqlite3.connect(str(db))
conn.execute("CREATE TABLE t (x INTEGER, y TEXT)")
conn.executemany("INSERT INTO t VALUES (?, ?)", [(1, "a"), (2, "b")])
conn.commit()
conn.close()
out = analytics._query(db, "SELECT * FROM t ORDER BY x")
assert out.to_dicts() == [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}]
# -----------------------------------------------------------------------------
# Seeded-registry fixture: builds the schema + minimal rows
# -----------------------------------------------------------------------------
def _create_registry_schema(conn: sqlite3.Connection) -> None:
"""Create the subset of the registry schema the analytics queries touch."""
conn.executescript(
"""
CREATE TABLE training_runs (
training_hash TEXT PRIMARY KEY,
family TEXT NOT NULL,
label TEXT NOT NULL,
config_name TEXT,
spec_json TEXT,
created_at TEXT NOT NULL
);
CREATE TABLE prediction_sets (
prediction_hash TEXT PRIMARY KEY,
training_hash TEXT NOT NULL,
checkpoint_value INTEGER,
checkpoint_kind TEXT,
split TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE TABLE prediction_metrics (
prediction_hash TEXT PRIMARY KEY,
computed_at TEXT NOT NULL,
ic_mean REAL,
ic_std REAL,
ic_t REAL,
n_folds REAL,
pct_positive REAL,
task_type TEXT,
accuracy REAL,
balanced_accuracy REAL,
auc_roc REAL,
auc_pr REAL,
log_loss REAL,
brier_score REAL
);
CREATE TABLE backtest_runs (
backtest_hash TEXT PRIMARY KEY,
prediction_hash TEXT NOT NULL,
spec_json TEXT,
stage TEXT,
created_at TEXT NOT NULL
);
CREATE TABLE backtest_metrics (
backtest_hash TEXT PRIMARY KEY,
computed_at TEXT NOT NULL,
sharpe REAL,
sortino REAL,
total_return REAL,
max_drawdown REAL
);
CREATE TABLE fold_metrics (
prediction_hash TEXT NOT NULL,
fold_id INTEGER NOT NULL,
computed_at TEXT NOT NULL,
ic REAL,
PRIMARY KEY (prediction_hash, fold_id)
);
"""
)
def _seed_registry(db_path: Path) -> None:
"""Populate a registry with 4 training runs + predictions + 2 backtests.
Layout (etfs):
linear_a / fwd_ret_21d / validation / ic=0.05 / regression
linear_b / fwd_ret_21d / validation / ic=0.08 / regression <- best linear
gbm_a / fwd_ret_21d / validation / ic=0.10 / regression <- best gbm (primary)
gbm_a / fwd_ret_21d / holdout / ic=0.03 / regression
linear_c / fwd_dir_5d / validation / ic=0.04 / classification (task_type='classification')
Plus 1 ch16 (signal) backtest and 1 ch17 (allocation) backtest on the
same gbm prediction_hash.
"""
db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(db_path))
_create_registry_schema(conn)
# training_runs
conn.executemany(
"INSERT INTO training_runs VALUES (?, ?, ?, ?, ?, ?)",
[
("th_lin_a", "linear", "fwd_ret_21d", "ridge_a100", None, "2024-01-01T00:00:00"),
("th_lin_b", "linear", "fwd_ret_21d", "ridge_b100", None, "2024-01-01T00:00:00"),
("th_gbm_a", "gbm", "fwd_ret_21d", "default", None, "2024-01-01T00:00:00"),
("th_lin_c", "linear", "fwd_dir_5d", "logistic", None, "2024-01-01T00:00:00"),
],
)
# prediction_sets
conn.executemany(
"INSERT INTO prediction_sets VALUES (?, ?, ?, ?, ?, ?)",
[
("ph_lin_a_val", "th_lin_a", 0, "final", "validation", "2024-01-02T00:00:00"),
("ph_lin_b_val", "th_lin_b", 0, "final", "validation", "2024-01-02T00:00:00"),
("ph_gbm_a_val", "th_gbm_a", 100, "final", "validation", "2024-01-02T00:00:00"),
("ph_gbm_a_hol", "th_gbm_a", 100, "final", "holdout", "2024-01-02T00:00:00"),
("ph_lin_c_val", "th_lin_c", 0, "final", "validation", "2024-01-02T00:00:00"),
],
)
# prediction_metrics — regression and classification
conn.executemany(
"INSERT INTO prediction_metrics (prediction_hash, computed_at, ic_mean, task_type, "
"auc_roc, accuracy) VALUES (?, ?, ?, ?, ?, ?)",
[
("ph_lin_a_val", "2024-01-03", 0.05, "regression", None, None),
("ph_lin_b_val", "2024-01-03", 0.08, "regression", None, None),
("ph_gbm_a_val", "2024-01-03", 0.10, "regression", None, None),
("ph_gbm_a_hol", "2024-01-03", 0.03, "regression", None, None),
("ph_lin_c_val", "2024-01-03", 0.04, "classification", 0.62, 0.55),
],
)
# backtest_runs — one signal (ch16) + one allocation (ch17) on the gbm prediction
spec_v2 = {
"version": 2,
"strategy": {"allocation": {"method": "inverse_vol"}},
"backtest_config": {
"commission": {"rate": 0.0005}, # 5 bps
"slippage": {"rate": 0.0003}, # 3 bps
},
}
conn.executemany(
"INSERT INTO backtest_runs VALUES (?, ?, ?, ?, ?)",
[
("bh_sig", "ph_gbm_a_val", json.dumps(spec_v2), "signal", "2024-01-04"),
("bh_alloc", "ph_gbm_a_val", json.dumps(spec_v2), "allocation", "2024-01-04"),
],
)
conn.executemany(
"INSERT INTO backtest_metrics (backtest_hash, computed_at, sharpe, sortino, "
"total_return, max_drawdown) VALUES (?, ?, ?, ?, ?, ?)",
[
("bh_sig", "2024-01-05", 1.2, 1.8, 0.35, -0.10),
("bh_alloc", "2024-01-05", 1.5, 2.2, 0.45, -0.08),
],
)
conn.commit()
conn.close()
@pytest.fixture
def seeded_registries(tmp_path, monkeypatch) -> Path:
"""Build registries for etfs + crypto_perps_funding under a temp output dir.
The second case study (crypto) is intentionally empty (only schema) so
multi-case-study queries have a no-op partition to merge against.
"""
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
_seed_registry(tmp_path / "etfs" / "run_log" / "registry.db")
# crypto: schema but no rows
crypto_db = tmp_path / "crypto_perps_funding" / "run_log" / "registry.db"
crypto_db.parent.mkdir(parents=True)
conn = sqlite3.connect(str(crypto_db))
_create_registry_schema(conn)
conn.commit()
conn.close()
return tmp_path
# -----------------------------------------------------------------------------
# load_model_ic
# -----------------------------------------------------------------------------
def test_load_model_ic_returns_all_families_by_default(seeded_registries) -> None:
df = analytics.load_model_ic(case_studies=["etfs", "crypto_perps_funding"], split="validation")
# etfs: 3 validation rows (lin_a, lin_b, gbm_a) with regression task_type;
# lin_c is also validation but the query doesn't filter task_type here.
assert df.height == 4
assert set(df["case_study"].unique().to_list()) == {"etfs"}
assert set(df["family"].unique().to_list()) == {"linear", "gbm"}
def test_load_model_ic_filters_by_family(seeded_registries) -> None:
df = analytics.load_model_ic(families="gbm", case_studies=["etfs"], split="validation")
assert df["family"].unique().to_list() == ["gbm"]
assert df.height == 1
def test_load_model_ic_filters_by_split_holdout(seeded_registries) -> None:
df = analytics.load_model_ic(case_studies=["etfs"], split="holdout")
# only gbm_a has a holdout prediction
assert df.height == 1
assert df["split"].to_list() == ["holdout"]
assert df["ic_mean"].to_list() == [0.03]
def test_load_model_ic_returns_empty_when_no_case_study_has_data(seeded_registries) -> None:
df = analytics.load_model_ic(case_studies=["crypto_perps_funding"], split="validation")
assert df.is_empty()
def test_load_model_ic_has_case_study_label_column(seeded_registries) -> None:
df = analytics.load_model_ic(case_studies=["etfs"], split="validation")
assert "case_study" in df.columns
assert df["case_study"].unique().to_list() == ["etfs"]
# -----------------------------------------------------------------------------
# load_classification_metrics
# -----------------------------------------------------------------------------
def test_load_classification_metrics_filters_task_type_eq_1(seeded_registries) -> None:
"""Only the linear_c row (task_type='classification') should come back."""
df = analytics.load_classification_metrics(case_studies=["etfs"], split="validation")
assert df.height == 1
assert df["family"].to_list() == ["linear"]
assert df["auc_roc"].to_list() == [0.62]
assert df["task_type"].to_list() == ["classification"]
def test_load_classification_metrics_excludes_regression_rows(seeded_registries) -> None:
"""Regression rows (task_type='regression') must not leak into the classification view."""
df = analytics.load_classification_metrics(case_studies=["etfs"], split="validation")
# Spec: no rows with null AUC should appear
assert df.filter(pl.col("auc_roc").is_null()).is_empty()
# -----------------------------------------------------------------------------
# load_best_ic_per_family
# -----------------------------------------------------------------------------
def test_load_best_ic_per_family_picks_max_per_pair(seeded_registries) -> None:
"""Primary label for etfs is fwd_ret_21d. Among linear runs on that label,
lin_b (IC=0.08) beats lin_a (IC=0.05). gbm has only one run (IC=0.10).
"""
best = analytics.load_best_ic_per_family(case_studies=["etfs"], split="validation")
rows = {(r["case_study"], r["family"]): r for r in best.to_dicts()}
assert rows[("etfs", "linear")]["ic_mean"] == 0.08
assert rows[("etfs", "linear")]["config_name"] == "ridge_b100"
assert rows[("etfs", "gbm")]["ic_mean"] == 0.10
def test_load_best_ic_per_family_primary_label_excludes_other_labels(
seeded_registries,
) -> None:
"""With use_primary_label=True (default), the linear_c run on fwd_dir_5d must
not appear — only rows with label == primary are kept.
"""
best = analytics.load_best_ic_per_family(case_studies=["etfs"], split="validation")
assert all(row["label"] == "fwd_ret_21d" for row in best.to_dicts())
def test_load_best_ic_per_family_use_primary_label_false_includes_all(
seeded_registries,
) -> None:
"""With use_primary_label=False, the fwd_dir_5d row competes for best linear."""
best = analytics.load_best_ic_per_family(
case_studies=["etfs"], split="validation", use_primary_label=False
)
# linear: best of {lin_a 0.05, lin_b 0.08, lin_c 0.04} is lin_b
linear_row = next(r for r in best.to_dicts() if r["family"] == "linear")
assert linear_row["ic_mean"] == 0.08
def test_load_best_ic_per_family_adds_display_name(seeded_registries) -> None:
best = analytics.load_best_ic_per_family(case_studies=["etfs"], split="validation")
assert "display_name" in best.columns
assert best["display_name"].unique().to_list() == ["ETFs"]
# -----------------------------------------------------------------------------
# load_chapter_backtests
# -----------------------------------------------------------------------------
def test_load_chapter_backtests_ch16_maps_to_signal_stage(seeded_registries) -> None:
"""chapter='ch16' → stage='signal' → returns the bh_sig backtest row."""
df = analytics.load_chapter_backtests("ch16", case_studies=["etfs"])
assert df.height == 1
assert df["backtest_hash"].to_list() == ["bh_sig"]
def test_load_chapter_backtests_explicit_stage_overrides_chapter(seeded_registries) -> None:
df = analytics.load_chapter_backtests("ch16", stage="allocation", case_studies=["etfs"])
assert df["backtest_hash"].to_list() == ["bh_alloc"]
def test_load_chapter_backtests_joins_sharpe_and_training_columns(seeded_registries) -> None:
df = analytics.load_chapter_backtests("ch17", case_studies=["etfs"])
assert df.height == 1
row = df.to_dicts()[0]
assert row["sharpe"] == 1.5
assert row["family"] == "gbm"
assert row["config_name"] == "default"
def test_load_chapter_backtests_metrics_filter_selects_columns(seeded_registries) -> None:
df = analytics.load_chapter_backtests("ch17", case_studies=["etfs"], metrics=["sharpe"])
# Only meta columns + sharpe; sortino/total_return/max_drawdown excluded
assert "sharpe" in df.columns
assert "sortino" not in df.columns
assert "total_return" not in df.columns
def test_load_chapter_backtests_returns_empty_for_unused_stage(seeded_registries) -> None:
df = analytics.load_chapter_backtests("ch18", case_studies=["etfs"])
assert df.is_empty()
# -----------------------------------------------------------------------------
# Spec helpers
# -----------------------------------------------------------------------------
def test_parse_backtest_spec_round_trips_json() -> None:
spec = {"version": 2, "strategy": {"foo": "bar"}, "backtest_config": {}}
assert analytics.parse_backtest_spec(json.dumps(spec)) == spec
def test_extract_cost_bps_sums_commission_and_slippage_from_v2_spec() -> None:
spec_json = json.dumps(
{
"version": 2,
"strategy": {},
"backtest_config": {
"commission": {"rate": 0.0005}, # 5 bps
"slippage": {"rate": 0.0003}, # 3 bps
},
}
)
assert analytics.extract_cost_bps(spec_json) == pytest.approx(8.0)
def test_extract_cost_bps_handles_v1_spec() -> None:
"""v1 specs store costs in a flat ``costs`` dict; cost_view falls back to it."""
spec_json = json.dumps({"costs": {"commission_bps": 2.0, "slippage_bps": 1.5}})
assert analytics.extract_cost_bps(spec_json) == pytest.approx(3.5)
def test_extract_allocator_reads_strategy_allocation_method() -> None:
spec_json = json.dumps(
{
"version": 2,
"strategy": {"allocation": {"method": "risk_parity"}},
"backtest_config": {},
}
)
assert analytics.extract_allocator(spec_json) == "risk_parity"
def test_extract_allocator_defaults_to_unknown_when_missing() -> None:
spec_json = json.dumps({"version": 2, "strategy": {}, "backtest_config": {}})
assert analytics.extract_allocator(spec_json) == "unknown"
tests/test_chapter_imports.py
+50
View File
@@ -0,0 +1,50 @@
"""Guard: per-chapter helper modules import from the repo root.
Chapter directories are number-prefixed (``25_live_trading``), so they are not
Python packages and their helper modules (``async_utils`` etc.) are only
importable when the chapter directory is on ``sys.path``. ``sitecustomize.py``
(declared as a top-level py-module in pyproject) arranges that at interpreter
startup in every environment. This test pins that contract: a bare
``import async_utils`` must succeed in a fresh interpreter started from the
repo root, with no chapter directory injected onto the path.
If this fails, ``sitecustomize.py`` or its ``[tool.setuptools] py-modules``
declaration was likely removed, or the package needs reinstalling
(``uv pip install -e .``).
"""
from __future__ import annotations
import os
import re
import subprocess
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
_CHAPTER_DIR = re.compile(r"/\d\d_[^/]+/?$")
def test_chapter_helper_imports_from_repo_root() -> None:
# Inherit the real environment (so PYTHONPATH=/app survives in Docker CI),
# but strip any pre-injected chapter directory so the test genuinely
# exercises the sitecustomize hook rather than a path the harness added.
env = dict(os.environ)
if pp := env.get("PYTHONPATH"):
kept = [p for p in pp.split(os.pathsep) if not _CHAPTER_DIR.search(p)]
env["PYTHONPATH"] = os.pathsep.join(kept)
# Representative sibling helpers from number-prefixed chapter dirs.
code = "import async_utils, limit_orderbook; print('ok')"
result = subprocess.run(
[sys.executable, "-c", code],
cwd=REPO_ROOT,
capture_output=True,
text=True,
env=env,
)
assert result.returncode == 0, (
"Bare chapter-helper import failed from the repo root — the "
"sitecustomize path hook is not active.\n"
f"stdout: {result.stdout}\nstderr: {result.stderr}"
)
tests/test_chapter_notebooks.py
+122
View File
@@ -0,0 +1,122 @@
"""Test chapter teaching notebooks via Papermill parameter injection.
Instead of the legacy TEST=1 environment variable (which creates divergent code paths),
this module uses Papermill to inject medium-scale parameter overrides into notebooks.
The same code path always runs; only the scale differs.
When ML4T_OUTPUT_DIR is set and contains pre-generated intermediates (from
generate_intermediates.py), chapter notebooks that depend on case study artifacts
(labels, features, predictions) will find them. This is seeded in CI by copying
intermediates from the test-data repo into ML4T_OUTPUT_DIR before running tests.
Usage:
# All chapters
pytest tests/test_chapter_notebooks.py -v
# Specific chapter
pytest tests/test_chapter_notebooks.py -v -k "ch05"
# Specific notebook
pytest tests/test_chapter_notebooks.py -v -k "tailgan"
"""
from pathlib import Path
import pytest
from tests.pm_helpers import (
collect_chapter_notebooks,
current_test_tier,
get_overrides,
get_tier,
run_notebook,
)
REPO_ROOT = Path(__file__).parent.parent
# Collect all chapter teaching notebooks (Ch01-Ch26)
CHAPTER_RANGE = range(1, 27)
CHAPTER_NOTEBOOKS = collect_chapter_notebooks(REPO_ROOT, CHAPTER_RANGE)
# Also collect per-dataset card notebooks (data/*/dataset_card.py, data/*/*/dataset_card.py)
for notebook in sorted(REPO_ROOT.glob("data/**/dataset_card.py")):
CHAPTER_NOTEBOOKS.append(notebook)
print(f"Found {len(CHAPTER_NOTEBOOKS)} chapter notebooks to test")
@pytest.mark.parametrize(
"notebook_path",
CHAPTER_NOTEBOOKS,
ids=lambda p: p.relative_to(REPO_ROOT).as_posix().replace("/", "::"),
)
def test_chapter_notebook(notebook_path, populated_data_dir, seeded_output_dir):
"""Execute a chapter notebook via Papermill with medium-scale overrides.
Each notebook runs with:
- Production defaults (what readers see)
- Papermill-injected overrides from tests/overrides.yaml (medium scale)
- ML4T_OUTPUT_DIR set to seeded output dir (has case study configs)
- MPLBACKEND=Agg, PLOTLY_RENDERER=json (headless rendering)
Markers (applied at collection time via conftest.py):
- ``pytest -m gpu`` — run only GPU-requiring notebooks
- ``pytest -m "not gpu"`` — run only CPU notebooks
"""
rel_path = notebook_path.relative_to(REPO_ROOT).with_suffix("")
overrides = get_overrides(str(rel_path))
# Tier routing: skip when NB tier doesn't match the current run tier.
# Default tier is per_commit; weekly/on_demand NBs require their dedicated
# workflow to set ML4T_TEST_TIER explicitly.
nb_tier = get_tier(overrides)
run_tier = current_test_tier()
if nb_tier != run_tier:
pytest.skip(f"Tier {nb_tier} — current run tier is {run_tier}")
# Skip if overrides say so (e.g., missing test data)
if overrides.get("skip"):
pytest.skip(f"Skipped: {overrides.get('skip_reason', 'marked skip in overrides')}")
# Check required imports (e.g., gensim, signatory, duckdb)
requires = overrides.get("requires_import")
if requires:
pkg = requires if isinstance(requires, str) else requires[0]
try:
__import__(pkg)
except ImportError:
pytest.skip(f"Requires {pkg} (not installed in this Docker image)")
# Check GPU requirement
if overrides.get("gpu"):
try:
import torch
if not torch.cuda.is_available():
pytest.skip("GPU required but not available")
except ImportError:
pytest.skip("GPU required but torch not installed")
timeout = overrides.get("timeout", 300)
parameters = overrides.get("parameters", {})
# Data layer notebooks expect to run from their own directory (for config.yaml)
notebook_cwd = notebook_path.parent if "data/" in str(rel_path) else None
result = run_notebook(
py_path=notebook_path,
parameters=parameters,
timeout=timeout,
output_dir=seeded_output_dir,
data_dir=populated_data_dir,
cwd=notebook_cwd,
)
if result["status"] == "error":
pytest.fail(
f"\n{'=' * 70}\n"
f"Notebook failed: {rel_path}\n"
f"{'=' * 70}\n"
f"Error: {result['error']}\n"
f"{'=' * 70}\n"
)
tests/test_compose_mounts.py
+117
View File
@@ -0,0 +1,117 @@
"""Phase 2a — deterministic compose-mount contract (regression-lock for #361).
The in-container free-data download bug (#361) was that ``docker-compose.yml``
mounted the reader's data directory **read-only** (``:ro``), so every downloader
failed with ``Read-only file system`` the moment it tried to write to
``/data``. The functional container smoke (Phase 2b, ``container-smoke.yml``)
catches that end-to-end, but a full 12 GB image pull is too heavy to run per PR.
This test pins the exact compose contract behind that fix, statically, in
milliseconds: the data volume is writable and ``ML4T_DATA_PATH`` points at the
mount. No Docker, no network — just parse the compose file. If someone flips the
mount back to ``:ro`` or repoints ``ML4T_DATA_PATH``, this fails on the PR that
does it, not weeks later in a reader's terminal.
"""
from __future__ import annotations
import re
from pathlib import Path
import pytest
import yaml
REPO_ROOT = Path(__file__).parent.parent
COMPOSE_FILE = REPO_ROOT / "docker-compose.yml"
# Services a reader actually runs the download workflow in. Benchmark/db-only
# services don't mount /data for writes and are out of scope here.
_WRITE_SERVICES = ("ml4t", "ml4t-gpu")
def _load_compose() -> dict:
# SafeLoader resolves YAML anchors/aliases and ``<<`` merge keys, so each
# service dict already carries the merged ``x-common`` volumes/environment.
return yaml.safe_load(COMPOSE_FILE.read_text())
# Container target ``/data`` with an optional trailing mode. Anchored at the end
# so an interpolated source like ``${ML4T_DATA_PATH:-./data}`` — which itself
# contains ``:`` — can't be mistaken for the target/mode.
_DATA_MOUNT_RE = re.compile(r":(?P<target>/data)(?::(?P<mode>[a-zA-Z]+))?$")
def _data_mounts(service: dict) -> list[tuple[str, str]]:
"""``(volume_string, mode)`` for every mount whose container target is ``/data``.
``mode`` is ``""`` when the entry omits an explicit rw/ro suffix.
"""
mounts = []
for vol in service.get("volumes", []):
if not isinstance(vol, str):
continue # long-form mounts not used for /data here
m = _DATA_MOUNT_RE.search(vol)
if m:
mounts.append((vol, m.group("mode") or ""))
return mounts
def test_compose_file_exists():
assert COMPOSE_FILE.exists(), f"missing {COMPOSE_FILE}"
@pytest.mark.parametrize("service_name", _WRITE_SERVICES)
def test_data_mount_is_writable(service_name):
"""The reader's /data mount must be read-write (the #361 fix)."""
compose = _load_compose()
service = compose["services"][service_name]
mounts = _data_mounts(service)
assert mounts, f"{service_name}: no /data volume mount found"
for vol, mode in mounts:
assert mode != "ro", (
f"{service_name}: /data is mounted read-only ('{vol}') — this is the "
f"#361 bug; the in-container download workflow cannot write to /data"
)
assert mode == "rw", (
f"{service_name}: /data mount '{vol}' must be explicitly ':rw' so the "
f"download workflow can populate it"
)
@pytest.mark.parametrize("service_name", _WRITE_SERVICES)
def test_data_path_env_points_at_mount(service_name):
"""ML4T_DATA_PATH inside the container must be the /data mount target.
Every downloader resolves its output root from ML4T_DATA_PATH; if it doesn't
equal the writable mount, files land somewhere the host mount can't see
(the ETF wrong-dir class of bug).
"""
compose = _load_compose()
env = compose["services"][service_name].get("environment", [])
# environment is a list of "KEY=VALUE" strings in this compose file.
env_map = {}
for item in env:
if isinstance(item, str) and "=" in item:
key, _, value = item.partition("=")
env_map[key] = value
assert env_map.get("ML4T_DATA_PATH") == "/data", (
f"{service_name}: ML4T_DATA_PATH is {env_map.get('ML4T_DATA_PATH')!r}, "
f"expected '/data' (the writable mount target)"
)
def test_no_service_mounts_data_readonly():
"""Defensive: no service anywhere reintroduces a read-only /data mount."""
compose = _load_compose()
offenders = []
for name, service in compose.get("services", {}).items():
if not isinstance(service, dict):
continue
for vol, mode in _data_mounts(service):
if mode == "ro":
offenders.append(f"{name} -> {vol}")
assert not offenders, f"read-only /data mount(s) found: {offenders}"
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_cv_splits.py
+361
View File
@@ -0,0 +1,361 @@
"""Tests for utils/cv_splits.py — walk-forward split generation.
Pins the invariants that every Ch11+ pipeline depends on:
- Pure duration/calendar normalization (regex-based, hermetic).
- load_evaluation_config reads setup.yaml's ``evaluation`` block and merges
the market_data semantics calendar.
- generate_cv_splits produces n_splits folds with the correct chronology,
backward walk-forward direction, embargo gap (label_buffer), and respects
the holdout_start boundary.
- make_walk_forward_config returns int label_horizon for calendar-aware
case studies (trading days) and Timedelta for 24/7 crypto.
Uses the real etfs and crypto_perps_funding setup.yaml files as ground
truth so the tests double as regression guards on those configs — if the
n_splits / train_size / val_size values are reordered, these tests will
flag it before a sweep wastes GPU time.
"""
from __future__ import annotations
import pandas as pd
import polars as pl
import pytest
import yaml
from utils.cv_splits import (
_map_calendar_id,
_normalize_duration,
_normalize_label_buffer,
generate_cv_splits,
load_evaluation_config,
make_walk_forward_config,
make_wf_config,
)
# -----------------------------------------------------------------------------
# Pure: _map_calendar_id
# -----------------------------------------------------------------------------
@pytest.mark.parametrize(
"setup_name, expected",
[
(None, None),
("NYSE", "NYSE"),
("CME", "CME_Equity"),
("FX", "CME_FX"),
("crypto", None), # 24/7 → disable calendar-aware splitting
("LSE", "LSE"), # unknown → pass through
],
)
def test_map_calendar_id(setup_name, expected) -> None:
assert _map_calendar_id(setup_name) == expected
# -----------------------------------------------------------------------------
# Pure: _normalize_duration (ISO 8601 stripping + unit aliasing)
# -----------------------------------------------------------------------------
@pytest.mark.parametrize(
"raw, normalized",
[
("P5Y", "5Y"),
("P1Y", "1Y"),
("1Y", "1Y"),
("PT8H", "8h"),
("8H", "8h"), # H → h for pd.Timedelta compatibility
("21D", "21D"),
("15T", "15min"), # T is a legacy pandas minute alias
],
)
def test_normalize_duration(raw, normalized) -> None:
assert _normalize_duration(raw) == normalized
# -----------------------------------------------------------------------------
# Pure: _normalize_label_buffer (inherits normalization + M → days)
# -----------------------------------------------------------------------------
@pytest.mark.parametrize(
"raw, normalized",
[
("21D", "21D"),
("PT8H", "8h"),
("1M", "30D"), # month → 30 days (pd.Timedelta rejects raw M)
("3M", "90D"),
("P6M", "180D"),
],
)
def test_normalize_label_buffer(raw, normalized) -> None:
assert _normalize_label_buffer(raw) == normalized
# -----------------------------------------------------------------------------
# load_evaluation_config
# -----------------------------------------------------------------------------
def test_load_evaluation_config_etfs_keys_and_values() -> None:
"""etfs is NYSE / 10Y train / 1Y val / 8 splits / backward (ground truth)."""
cfg = load_evaluation_config("etfs")
assert cfg["n_splits"] == 8
assert cfg["train_size"] == "10Y"
assert cfg["val_size"] == "1Y"
assert cfg["holdout_start"] == "2024-01-01"
assert cfg["holdout_end"] == "2025-12-31"
assert cfg["calendar"] == "NYSE"
def test_load_evaluation_config_crypto_keeps_24_7_calendar() -> None:
"""crypto sets calendar: crypto (24/7); preserved in the returned config."""
cfg = load_evaluation_config("crypto_perps_funding")
assert cfg["calendar"] == "crypto"
def test_load_evaluation_config_raises_on_missing_section(tmp_path, monkeypatch) -> None:
"""A setup.yaml without an ``evaluation`` section raises KeyError.
We spoof the case-study dir via ML4T_OUTPUT_DIR. The fallback path
(re-read from source) won't find the fake id either, so the outer
check raises.
"""
cs_id = "_cv_splits_test_missing_evaluation"
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
cfg_dir = tmp_path / cs_id / "config"
cfg_dir.mkdir(parents=True)
(cfg_dir / "setup.yaml").write_text(yaml.safe_dump({"labels": {"primary": "x"}}))
with pytest.raises(KeyError, match="evaluation"):
load_evaluation_config(cs_id)
# -----------------------------------------------------------------------------
# generate_cv_splits — uses real etfs config (NYSE, 10Y/1Y, 8 splits, backward)
# -----------------------------------------------------------------------------
@pytest.fixture(scope="module")
def etfs_daily_frame() -> pl.DataFrame:
"""~24 years of business days — enough for 8 backward folds of 10+1 years."""
ts = pd.date_range("1999-01-01", "2023-12-31", freq="B")
return pl.DataFrame({"timestamp": pl.Series(ts)})
@pytest.fixture(scope="module")
def etfs_splits(etfs_daily_frame) -> list[dict]:
return generate_cv_splits(etfs_daily_frame, case_study_id="etfs", label_buffer="21D")
def test_generate_cv_splits_etfs_returns_n_splits_folds(etfs_splits) -> None:
assert len(etfs_splits) == 8
def test_generate_cv_splits_etfs_fold_ids_are_0_through_n_minus_1(etfs_splits) -> None:
assert [s["fold"] for s in etfs_splits] == list(range(len(etfs_splits)))
def test_generate_cv_splits_etfs_folds_have_required_keys(etfs_splits) -> None:
required = {"fold", "train_start", "train_end", "val_start", "val_end"}
for s in etfs_splits:
assert required <= set(s)
def test_generate_cv_splits_etfs_intra_fold_chronology(etfs_splits) -> None:
"""Within each fold: train_start ≤ train_end < val_start ≤ val_end."""
for s in etfs_splits:
assert s["train_start"] <= s["train_end"]
assert s["train_end"] < s["val_start"]
assert s["val_start"] <= s["val_end"]
def test_generate_cv_splits_etfs_backward_walk_forward(etfs_splits) -> None:
"""fold_direction=backward → fold 0 is the most recent, folds step back."""
for i in range(len(etfs_splits) - 1):
assert etfs_splits[i]["val_start"] > etfs_splits[i + 1]["val_start"]
def test_generate_cv_splits_etfs_embargo_respects_label_buffer(etfs_splits) -> None:
"""The gap between train_end and val_start covers the 21-trading-day label
horizon. On NYSE that is roughly 29-32 calendar days; allow a generous
lower bound to avoid flaking on holiday spacing.
"""
for s in etfs_splits:
gap = s["val_start"] - s["train_end"]
assert gap >= pd.Timedelta(days=21), s # at minimum 21 calendar days
def test_generate_cv_splits_etfs_val_before_holdout(etfs_splits) -> None:
"""All validation windows end strictly before the holdout_start (2024-01-01)."""
holdout_start = pd.Timestamp("2024-01-01")
for s in etfs_splits:
assert s["val_end"] < holdout_start, s
def test_generate_cv_splits_etfs_train_size_10y(etfs_splits) -> None:
"""10Y train_size — span should be ~10 years (±2 months for calendar alignment)."""
for s in etfs_splits:
span = s["train_end"] - s["train_start"]
assert pd.Timedelta(days=365 * 10 - 60) <= span <= pd.Timedelta(days=365 * 10 + 60), s
def test_generate_cv_splits_etfs_val_size_1y(etfs_splits) -> None:
"""1Y val_size — span should be ~1 year."""
for s in etfs_splits:
span = s["val_end"] - s["val_start"]
assert pd.Timedelta(days=330) <= span <= pd.Timedelta(days=380), s
# -----------------------------------------------------------------------------
# generate_cv_splits — crypto (24/7, calendar=None after mapping)
# -----------------------------------------------------------------------------
def test_generate_cv_splits_crypto_respects_8h_buffer_and_no_calendar() -> None:
ts = pd.date_range("2019-01-01", "2023-12-31", freq="8h")
df = pl.DataFrame({"timestamp": pl.Series(ts)})
splits = generate_cv_splits(df, case_study_id="crypto_perps_funding", label_buffer="8H")
assert len(splits) == 2
for s in splits:
# 8h buffer means val_start ≥ train_end + 8h (may be slightly larger
# because step is in 8-hour bars).
gap = s["val_start"] - s["train_end"]
assert gap >= pd.Timedelta(hours=8), s
# -----------------------------------------------------------------------------
# generate_cv_splits — input DataFrame flavors
# -----------------------------------------------------------------------------
def test_generate_cv_splits_accepts_pandas_dataframe() -> None:
"""Both pl.DataFrame and pd.DataFrame inputs produce identical splits."""
ts = pd.date_range("1999-01-01", "2023-12-31", freq="B")
pdf = pd.DataFrame({"timestamp": ts})
pldf = pl.DataFrame({"timestamp": pl.Series(ts)})
pd_splits = generate_cv_splits(pdf, case_study_id="etfs", label_buffer="21D")
pl_splits = generate_cv_splits(pldf, case_study_id="etfs", label_buffer="21D")
assert pd_splits == pl_splits
# -----------------------------------------------------------------------------
# generate_cv_splits — legacy cv_config dict path
# -----------------------------------------------------------------------------
def test_generate_cv_splits_cv_config_passthrough_of_precomputed_splits() -> None:
"""If cv_config already carries a ``splits`` list, return it unchanged."""
precomputed = [
{
"fold": 0,
"train_start": "2020-01-01",
"train_end": "2022-12-31",
"val_start": "2023-01-01",
"val_end": "2023-12-31",
}
]
df = pl.DataFrame({"timestamp": pl.Series(pd.date_range("2020", "2023", freq="D"))})
out = generate_cv_splits(df, cv_config={"splits": precomputed})
assert out is precomputed or out == precomputed
def test_generate_cv_splits_cv_config_accepts_legacy_alias_keys() -> None:
"""Legacy keys test_size / test_start / test_end must be accepted.
Old pipeline persisted cv_config.json with these aliases; the loader
must still accept them so archived runs replay correctly.
"""
cv = {
"n_splits": 2,
"train_size": "5Y",
"test_size": "1Y",
"test_start": "2023-01-01",
"test_end": "2023-12-31",
"calendar": "NYSE",
}
ts = pd.date_range("2010-01-01", "2023-12-31", freq="B")
df = pl.DataFrame({"timestamp": pl.Series(ts)})
splits = generate_cv_splits(df, cv_config=cv, label_buffer="5D")
assert len(splits) == 2
for s in splits:
assert s["train_end"] < s["val_start"]
def test_generate_cv_splits_cv_config_with_val_size_key_also_works() -> None:
"""Newer pipelines persist val_size / holdout_start — also supported."""
cv = {
"n_splits": 2,
"train_size": "5Y",
"val_size": "1Y",
"holdout_start": "2023-01-01",
"holdout_end": "2023-12-31",
"calendar": "NYSE",
}
ts = pd.date_range("2010-01-01", "2023-12-31", freq="B")
df = pl.DataFrame({"timestamp": pl.Series(ts)})
splits = generate_cv_splits(df, cv_config=cv, label_buffer="5D")
assert len(splits) == 2
# -----------------------------------------------------------------------------
# generate_cv_splits — error paths
# -----------------------------------------------------------------------------
def test_generate_cv_splits_raises_without_any_config_source() -> None:
df = pl.DataFrame({"timestamp": pl.Series(pd.date_range("2020", "2023", freq="D"))})
with pytest.raises(ValueError, match="case_study_id"):
generate_cv_splits(df)
def test_generate_cv_splits_raises_on_empty_dataset() -> None:
df = pl.DataFrame({"timestamp": pl.Series([], dtype=pl.Datetime)})
with pytest.raises(ValueError, match="No timestamps"):
generate_cv_splits(df, case_study_id="etfs", label_buffer="21D")
# -----------------------------------------------------------------------------
# make_walk_forward_config
# -----------------------------------------------------------------------------
def test_make_walk_forward_config_nyse_label_horizon_is_int_trading_days() -> None:
"""NYSE case study with a D-unit buffer passes label_horizon as int so
the library counts trading days instead of calendar days.
"""
cfg = make_walk_forward_config("etfs", label_horizon="21D")
assert isinstance(cfg.label_horizon, int)
assert cfg.label_horizon == 21
assert cfg.calendar_id == "NYSE"
assert cfg.n_splits == 8
assert cfg.train_size == "10Y"
assert cfg.test_size == "1Y" # val_size → test_size alias
assert cfg.fold_direction == "backward"
def test_make_walk_forward_config_crypto_label_horizon_is_timedelta() -> None:
"""24/7 crypto: calendar_id=None → horizon stays as string/Timedelta."""
cfg = make_walk_forward_config("crypto_perps_funding", label_horizon="8H")
assert cfg.calendar_id is None
# Library may coerce to Timedelta; never an int for calendar-less case studies.
assert not isinstance(cfg.label_horizon, int)
def test_make_walk_forward_config_holdout_dates_round_trip() -> None:
"""holdout_start / holdout_end from setup.yaml flow through to test_start / test_end."""
cfg = make_walk_forward_config("etfs", label_horizon="21D")
# Library stores as date objects
assert str(cfg.test_start) == "2024-01-01"
assert str(cfg.test_end) == "2025-12-31"
def test_make_wf_config_is_alias_of_make_walk_forward_config() -> None:
"""Backward-compat alias should delegate with identical output."""
a = make_walk_forward_config("etfs", label_horizon="21D")
b = make_wf_config("etfs", label_horizon="21D")
assert a.model_dump() == b.model_dump()
tests/test_cv_window_fold_splits.py
+162
View File
@@ -0,0 +1,162 @@
"""Tests for case_studies/utils/cv_window.py P2.6 fixes (#2471).
Covers:
1. ``_fold_splits`` raises ``ValueError`` with the actionable
"Add buffer to labels.buffer..." hint when ``label_buffer`` is
missing from setup.yaml — restores the loud-fail contract that
matches ``utils.modeling.load_modeling_dataset``.
2. ``_fold_splits`` detects the time column from the parquet schema
(``timestamp`` else ``date``), so legacy parquets that haven't
migrated to the canonical ``timestamp`` name don't crash with
``ColumnNotFoundError``.
3. ``_fold_splits`` returns ``None`` when the label parquet doesn't
exist (unchanged contract).
"""
from __future__ import annotations
from datetime import date
from pathlib import Path
import polars as pl
import pytest
import yaml
@pytest.fixture
def isolated_case_study(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
"""Redirect get_case_study_dir to tmp_path via ML4T_OUTPUT_DIR.
Also clears the _fold_splits / _load_setup_yaml / _holdout_window
lru caches so tests don't leak case-study state across runs.
"""
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
from case_studies.utils import cv_window
cv_window._fold_splits.cache_clear()
cv_window._load_setup_yaml.cache_clear()
cv_window._holdout_window.cache_clear()
yield tmp_path
cv_window._fold_splits.cache_clear()
cv_window._load_setup_yaml.cache_clear()
cv_window._holdout_window.cache_clear()
def _seed_setup_yaml(cs_dir: Path, *, with_buffer: bool, label: str) -> None:
cs_dir.mkdir(parents=True, exist_ok=True)
cfg = cs_dir / "config"
cfg.mkdir(exist_ok=True)
setup: dict = {
"strategy_id": cs_dir.name,
"labels": {"primary": label},
"evaluation": {
"n_splits": 2,
"train_size": "1Y",
"val_size": "6M",
"holdout_start": "2023-01-01",
"holdout_end": "2023-12-31",
"calendar": "NYSE",
"periods_per_year": 252,
},
}
if with_buffer:
setup["labels"]["buffer"] = "21D"
(cfg / "setup.yaml").write_text(yaml.safe_dump(setup))
def _seed_label_parquet(cs_dir: Path, *, label: str, date_col: str) -> None:
"""Write a minimal label parquet with the given time column name."""
labels_dir = cs_dir / "labels"
labels_dir.mkdir(parents=True, exist_ok=True)
dates = pl.date_range(start=date(2020, 1, 1), end=date(2023, 12, 31), interval="1d", eager=True)
df = pl.DataFrame(
{
date_col: dates,
"symbol": ["AAA"] * len(dates),
label: [0.01] * len(dates),
}
)
df.write_parquet(labels_dir / f"{label}.parquet")
def test_missing_label_buffer_raises_with_actionable_hint(
isolated_case_study: Path,
) -> None:
"""Setup.yaml without labels.buffer must raise loudly."""
from case_studies.utils.cv_window import _fold_splits
cs = "test_cs_missing_buffer"
cs_dir = isolated_case_study / cs
_seed_setup_yaml(cs_dir, with_buffer=False, label="fwd_ret_21d")
_seed_label_parquet(cs_dir, label="fwd_ret_21d", date_col="timestamp")
with pytest.raises(ValueError, match=r"No explicit label buffer found for 'fwd_ret_21d'"):
_fold_splits(cs, "fwd_ret_21d")
def test_missing_label_parquet_returns_none(isolated_case_study: Path) -> None:
"""No parquet means 'no folds derivable' — still a None return."""
from case_studies.utils.cv_window import _fold_splits
cs = "test_cs_no_parquet"
cs_dir = isolated_case_study / cs
_seed_setup_yaml(cs_dir, with_buffer=True, label="fwd_ret_21d")
# NB: no parquet written
assert _fold_splits(cs, "fwd_ret_21d") is None
def test_schema_detection_picks_timestamp_column(isolated_case_study: Path) -> None:
"""Canonical-schema parquet with 'timestamp' column resolves folds."""
from case_studies.utils.cv_window import _fold_splits
cs = "test_cs_ts"
cs_dir = isolated_case_study / cs
_seed_setup_yaml(cs_dir, with_buffer=True, label="fwd_ret_21d")
_seed_label_parquet(cs_dir, label="fwd_ret_21d", date_col="timestamp")
splits = _fold_splits(cs, "fwd_ret_21d")
assert splits is not None
assert len(splits) >= 1
fold_id, val_start, val_end = splits[0]
assert fold_id == 0
assert isinstance(val_start, date) and isinstance(val_end, date)
assert val_start <= val_end
def test_schema_detection_falls_back_to_date_column(
isolated_case_study: Path,
) -> None:
"""Legacy 'date'-column parquet still works — no ColumnNotFoundError."""
from case_studies.utils.cv_window import _fold_splits
cs = "test_cs_date"
cs_dir = isolated_case_study / cs
_seed_setup_yaml(cs_dir, with_buffer=True, label="fwd_ret_21d")
_seed_label_parquet(cs_dir, label="fwd_ret_21d", date_col="date")
splits = _fold_splits(cs, "fwd_ret_21d")
assert splits is not None
assert len(splits) >= 1
def test_schema_without_timestamp_or_date_raises(
isolated_case_study: Path,
) -> None:
"""A parquet with neither 'timestamp' nor 'date' must raise actionably."""
from case_studies.utils.cv_window import _fold_splits
cs = "test_cs_no_time_col"
cs_dir = isolated_case_study / cs
_seed_setup_yaml(cs_dir, with_buffer=True, label="fwd_ret_21d")
# Write a parquet with neither column.
labels_dir = cs_dir / "labels"
labels_dir.mkdir(parents=True, exist_ok=True)
pl.DataFrame({"symbol": ["AAA"], "fwd_ret_21d": [0.01]}).write_parquet(
labels_dir / "fwd_ret_21d.parquet"
)
with pytest.raises(ValueError, match=r"neither 'timestamp' nor 'date'"):
_fold_splits(cs, "fwd_ret_21d")
tests/test_data_exceptions.py
+163
View File
@@ -0,0 +1,163 @@
"""Contract tests for data/exceptions.py.
Every loader in data/ raises DataNotFoundError with a specific combination
of keyword arguments (download_script vs instructions vs download_url, plus
readme). The tests below pin the message shape so a reformat of _build_message
does not silently drop the reader-facing download instructions.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from data.exceptions import DataNotFoundError, DownloadError, MissingDependencyError
def test_data_not_found_with_download_script_shows_command() -> None:
err = DataNotFoundError(
dataset_name="ETF Universe",
path=Path("/data/etfs/market/etf_universe.parquet"),
download_script="data/etfs/market/download.py",
)
msg = str(err)
assert "DATA NOT FOUND: ETF Universe" in msg
assert "/data/etfs/market/etf_universe.parquet" in msg
assert "uv run python data/etfs/market/download.py" in msg
assert "data/README.md" in msg # Default readme fallback
def test_data_not_found_with_custom_readme_overrides_default() -> None:
err = DataNotFoundError(
dataset_name="SEC Filings",
path=Path("/data/equities/fundamentals/"),
download_script="data/equities/fundamentals/filings_download.py",
readme="data/equities/fundamentals/README.md",
)
msg = str(err)
assert "data/equities/fundamentals/README.md" in msg
# Default readme should not appear when custom is set
assert msg.count("README.md") == 1
def test_data_not_found_with_download_url_prefers_url_over_script() -> None:
err = DataNotFoundError(
dataset_name="AlgoSeek Options",
path=Path("/data/options/sp500_options.parquet"),
download_url="https://example.com/algoseek.zip",
download_script="data/ignored/download.py", # Should be ignored
requires_api_key="ALGOSEEK_KEY",
)
msg = str(err)
assert "https://example.com/algoseek.zip" in msg
# download_script is shadowed by download_url in the elif chain
assert "uv run python data/ignored/download.py" not in msg
assert "Extract to:" in msg
assert "ALGOSEEK_KEY" in msg
def test_data_not_found_with_instructions_overrides_other_branches() -> None:
err = DataNotFoundError(
dataset_name="Derived Daily Bars",
path=Path("/data/futures/market/continuous/daily/continuous_daily.parquet"),
instructions="Run: python 02/05_futures_session_aggregation.py",
download_script="should_not_appear.py",
)
msg = str(err)
assert "02/05_futures_session_aggregation.py" in msg
assert "should_not_appear.py" not in msg
def test_data_not_found_with_derivation_notebook_adds_pointer() -> None:
err = DataNotFoundError(
dataset_name="Derived Label",
path=Path("/data/labels/fwd_ret_5d.parquet"),
derivation_notebook="07_defining_the_learning_task/03_label_methods.py",
)
msg = str(err)
assert "How this dataset is built" in msg
assert "07_defining_the_learning_task/03_label_methods.py" in msg
def test_data_not_found_is_filenotfounderror() -> None:
"""Callers catch FileNotFoundError generically; guard the inheritance."""
err = DataNotFoundError(
dataset_name="x",
path=Path("/tmp/x"),
download_script="x.py",
)
assert isinstance(err, FileNotFoundError)
with pytest.raises(FileNotFoundError):
raise err
def test_data_not_found_accepts_str_path() -> None:
"""Some callers pass str paths; the class must normalize to Path."""
err = DataNotFoundError(
dataset_name="x",
path="/tmp/string_path.parquet",
download_script="x.py",
)
assert isinstance(err.path, Path)
assert "/tmp/string_path.parquet" in str(err)
def test_download_error_includes_reason_and_suggestion() -> None:
err = DownloadError(
dataset_name="Crypto Perps",
reason="API returned 429",
suggestion="Retry after backoff or check rate limits",
)
msg = str(err)
assert "DOWNLOAD FAILED: Crypto Perps" in msg
assert "API returned 429" in msg
assert "Retry after backoff" in msg
def test_download_error_without_suggestion_still_valid() -> None:
err = DownloadError(dataset_name="x", reason="unknown")
msg = str(err)
assert "DOWNLOAD FAILED: x" in msg
assert "Suggestion:" not in msg
def test_download_error_is_runtime_error() -> None:
err = DownloadError(dataset_name="x", reason="y")
assert isinstance(err, RuntimeError)
def test_missing_dependency_error_defaults_install_command() -> None:
err = MissingDependencyError(package="edgartools")
msg = str(err)
assert "edgartools" in msg
assert "pip install edgartools" in msg
def test_missing_dependency_error_custom_install_and_purpose() -> None:
err = MissingDependencyError(
package="torch",
install_command="uv sync --extra gpu",
purpose="LSTM training in Ch13",
)
msg = str(err)
assert "torch" in msg
assert "uv sync --extra gpu" in msg
assert "LSTM training in Ch13" in msg
def test_missing_dependency_is_import_error() -> None:
"""Callers catch ImportError to gate optional backends."""
err = MissingDependencyError(package="x")
assert isinstance(err, ImportError)
tests/test_data_quality.py
+215
View File
@@ -0,0 +1,215 @@
"""Tests for utils/data_quality.py.
Pins:
- apply_max_symbols: seed determinism + edge cases (no-op when max<=0 or >=N).
Called by every loader; non-determinism would break reproducibility of tests
and notebooks that depend on a sampled subset.
- check_ohlc_invariants: correct detection of OHLC violations, graceful
handling of null values (TAQ no-trade bars), and missing-column tolerance.
"""
from __future__ import annotations
import polars as pl
from utils.data_quality import (
apply_max_symbols,
check_ohlc_invariants,
describe_coverage,
null_rate,
)
def _make_prices(symbols: list[str], n_rows: int = 3) -> pl.DataFrame:
rows = []
for s in symbols:
for i in range(n_rows):
rows.append(
{
"symbol": s,
"timestamp": f"2024-01-{i + 1:02d}",
"open": 100.0 + i,
"high": 105.0 + i,
"low": 95.0 + i,
"close": 102.0 + i,
"volume": 1_000 * (i + 1),
}
)
return pl.DataFrame(rows)
# -----------------------------------------------------------------------------
# apply_max_symbols
# -----------------------------------------------------------------------------
def test_apply_max_symbols_zero_is_passthrough() -> None:
df = _make_prices(["A", "B", "C"])
out = apply_max_symbols(df, 0)
assert out.equals(df)
def test_apply_max_symbols_negative_is_passthrough() -> None:
df = _make_prices(["A", "B", "C"])
out = apply_max_symbols(df, -1)
assert out.equals(df)
def test_apply_max_symbols_exceeds_universe_is_passthrough() -> None:
df = _make_prices(["A", "B"])
out = apply_max_symbols(df, 10)
assert out.equals(df)
def test_apply_max_symbols_samples_requested_count() -> None:
df = _make_prices(["A", "B", "C", "D", "E"])
out = apply_max_symbols(df, 3)
assert out["symbol"].n_unique() == 3
def test_apply_max_symbols_is_seed_deterministic() -> None:
"""Same seed → same subset; critical for reproducible tests."""
df = _make_prices(["A", "B", "C", "D", "E", "F", "G", "H"])
first = apply_max_symbols(df, 3, seed=42)["symbol"].unique().sort().to_list()
second = apply_max_symbols(df, 3, seed=42)["symbol"].unique().sort().to_list()
assert first == second
def test_apply_max_symbols_different_seed_yields_different_sample() -> None:
df = _make_prices(["A", "B", "C", "D", "E", "F", "G", "H"])
s42 = set(apply_max_symbols(df, 3, seed=42)["symbol"].unique().to_list())
s7 = set(apply_max_symbols(df, 3, seed=7)["symbol"].unique().to_list())
# At least one sample differs — very high probability for k=3, n=8
assert s42 != s7
def test_apply_max_symbols_preserves_all_rows_per_symbol() -> None:
df = _make_prices(["A", "B", "C", "D"], n_rows=5)
out = apply_max_symbols(df, 2)
# Each selected symbol should keep all its rows (function filters by symbol set)
per_symbol = out.group_by("symbol").len()
assert per_symbol["len"].to_list() == [5, 5]
def test_apply_max_symbols_sort_then_sample_is_order_invariant() -> None:
"""Shuffling the input before sampling must yield the same subset: the
function sorts symbols before seeding the RNG so unstable loader order
(e.g., parquet partition order) can't perturb the selection."""
df_asc = _make_prices(["A", "B", "C", "D", "E"])
df_desc = df_asc.sort("symbol", descending=True)
s1 = apply_max_symbols(df_asc, 2, seed=42)["symbol"].unique().sort().to_list()
s2 = apply_max_symbols(df_desc, 2, seed=42)["symbol"].unique().sort().to_list()
assert s1 == s2
def test_apply_max_symbols_custom_symbol_col() -> None:
df = pl.DataFrame({"product": ["X", "Y", "Z"], "value": [1, 2, 3]})
out = apply_max_symbols(df, 2, symbol_col="product")
assert out["product"].n_unique() == 2
def test_apply_max_symbols_with_lazyframe_returns_lazyframe() -> None:
df = _make_prices(["A", "B", "C", "D"])
lf = df.lazy()
out = apply_max_symbols(lf, 2)
assert isinstance(out, pl.LazyFrame)
assert out.collect()["symbol"].n_unique() == 2
# -----------------------------------------------------------------------------
# check_ohlc_invariants
# -----------------------------------------------------------------------------
def test_check_ohlc_invariants_clean_data_is_100_percent() -> None:
df = _make_prices(["A"], n_rows=5)
invariants = check_ohlc_invariants(df)
# 6 checks expected: high_gte_low/open/close, low_lte_open/close, volume_non_negative
assert invariants.height == 6
assert (invariants["valid_pct"] == 100.0).all()
def test_check_ohlc_invariants_detects_high_below_low() -> None:
df = pl.DataFrame(
{
"open": [100.0, 100.0],
"high": [99.0, 105.0], # first row violates high >= low
"low": [100.0, 95.0],
"close": [99.5, 102.0],
"volume": [1_000, 1_000],
}
)
invariants = check_ohlc_invariants(df)
row = invariants.filter(pl.col("check") == "high_gte_low").row(0, named=True)
assert row["valid_pct"] == 50.0
assert row["applicable_rows"] == 2
def test_check_ohlc_invariants_ignores_null_rows() -> None:
"""Rows where any required col is null are excluded from the percentage."""
df = pl.DataFrame(
{
"open": [100.0, None, 100.0],
"high": [105.0, None, 102.0],
"low": [95.0, None, 95.0],
"close": [101.0, None, 101.0],
"volume": [1_000, 1_000, 1_000],
}
)
invariants = check_ohlc_invariants(df)
row = invariants.filter(pl.col("check") == "high_gte_low").row(0, named=True)
assert row["applicable_rows"] == 2 # 3 rows, 1 excluded for nulls
assert row["valid_pct"] == 100.0
def test_check_ohlc_invariants_detects_negative_volume() -> None:
df = pl.DataFrame(
{
"open": [100.0],
"high": [105.0],
"low": [95.0],
"close": [101.0],
"volume": [-5],
}
)
invariants = check_ohlc_invariants(df)
row = invariants.filter(pl.col("check") == "volume_non_negative").row(0, named=True)
assert row["valid_pct"] == 0.0
def test_check_ohlc_invariants_omits_volume_when_missing() -> None:
df = pl.DataFrame({"open": [100.0], "high": [105.0], "low": [95.0], "close": [101.0]})
invariants = check_ohlc_invariants(df)
assert "volume_non_negative" not in invariants["check"].to_list()
assert invariants.height == 5 # 5 price checks, no volume check
def test_check_ohlc_invariants_empty_df_returns_empty_result() -> None:
df = pl.DataFrame({"x": []})
invariants = check_ohlc_invariants(df)
assert invariants.height == 0
# -----------------------------------------------------------------------------
# Smaller coverage helpers
# -----------------------------------------------------------------------------
def test_describe_coverage_basic_shape() -> None:
df = _make_prices(["A", "B"], n_rows=3)
cov = describe_coverage(df)
assert cov["rows"] == 6
assert cov["assets"] == 2
assert cov["unique_times"] == 3
def test_null_rate_reports_per_column() -> None:
df = pl.DataFrame({"a": [1, None, 3], "b": [None, None, 3]})
rates = null_rate(df)
by_col = dict(zip(rates["column"], rates["null_pct"], strict=True))
# 1/3 ≈ 33.33 for a, 2/3 ≈ 66.66 for b
assert round(by_col["a"], 2) == 33.33
assert round(by_col["b"], 2) == 66.67
tests/test_docker_notebooks.py
+98
View File
@@ -0,0 +1,98 @@
"""Test notebooks that require Docker environments (py312, neo4j, benchmark).
Same as test_chapter_notebooks.py but IGNORES skip flags from overrides.yaml.
These notebooks are skipped in uv-native runs (missing modules like signatory,
gensim, esig, tfcausalimpact, or Neo4j) but CAN run inside their respective
Docker images.
The skip flag stays in overrides.yaml so the uv-native runner still skips them.
This file runs them in Docker where the dependencies are available.
Usage:
# Py312 notebooks (signatory, gensim, esig, pfhedge, tfcausalimpact, torch CUDA bug)
python -m pytest tests/test_docker_notebooks.py -v -k "03_sigcwgan or ..."
# Neo4j notebooks
python -m pytest tests/test_docker_notebooks.py -v -k "08_8k_event_extraction or ..."
# Benchmark notebooks
python -m pytest tests/test_docker_notebooks.py -v -k "18_storage_benchmark_database or ..."
"""
from pathlib import Path
import pytest
from tests.pm_helpers import (
collect_chapter_notebooks,
current_test_tier,
get_overrides,
get_tier,
run_notebook,
)
REPO_ROOT = Path(__file__).parent.parent
# Collect all chapter teaching notebooks (Ch01-Ch26) + data layer
CHAPTER_RANGE = range(1, 27)
CHAPTER_NOTEBOOKS = collect_chapter_notebooks(REPO_ROOT, CHAPTER_RANGE)
for notebook in sorted(REPO_ROOT.glob("data/**/dataset_card.py")):
CHAPTER_NOTEBOOKS.append(notebook)
@pytest.mark.parametrize(
"notebook_path",
CHAPTER_NOTEBOOKS,
ids=lambda p: p.relative_to(REPO_ROOT).as_posix().replace("/", "::"),
)
def test_docker_notebook(notebook_path, populated_data_dir, seeded_output_dir):
"""Execute a notebook via Papermill, ignoring skip flags.
Identical to test_chapter_notebook() but does NOT honor the 'skip' key
in overrides.yaml. This allows Docker-based CI jobs to run notebooks
that are skipped in the uv-native environment due to missing modules.
GPU skips are still honored (Docker CI runners have no GPU).
"""
rel_path = notebook_path.relative_to(REPO_ROOT).with_suffix("")
overrides = get_overrides(str(rel_path))
# Tier routing still applies — Docker tests are normally per_commit, but
# this honors the same env-driven gating used by the uv-native runners.
nb_tier = get_tier(overrides)
run_tier = current_test_tier()
if nb_tier != run_tier:
pytest.skip(f"Tier {nb_tier} — current run tier is {run_tier}")
# NOTE: We intentionally do NOT check overrides.get("skip") here.
# That's the whole point of this file — Docker provides the missing deps.
# GPU requirement still applies (CI runners have no GPU)
if overrides.get("gpu"):
try:
import torch
if not torch.cuda.is_available():
pytest.skip("GPU required but not available")
except ImportError:
pytest.skip("GPU required but torch not installed")
timeout = overrides.get("timeout", 300)
parameters = overrides.get("parameters", {})
result = run_notebook(
py_path=notebook_path,
parameters=parameters,
timeout=timeout,
output_dir=seeded_output_dir,
data_dir=populated_data_dir,
)
if result["status"] == "error":
pytest.fail(
f"\n{'=' * 70}\n"
f"Notebook failed: {rel_path}\n"
f"{'=' * 70}\n"
f"Error: {result['error']}\n"
f"{'=' * 70}\n"
)
tests/test_download_all_dispatch.py
+101
View File
@@ -0,0 +1,101 @@
"""Dispatch-logic tests for ``data/download_all.py``.
These run the real ``main()`` control flow with the network boundary
(``run_download_script``, which shells out to each downloader) replaced by a
recorder. No subprocesses, no network. They pin which datasets each mode
dispatches — the layer that silently broke in the #361 fixes:
- ``--free-only`` must skip the API-key datasets (macro/FX) but keep the core +
factor + firm-char downloads;
- ``--skip-firm-characteristics`` must omit the 1.5 GB academic pull;
- firm characteristics must be invoked *plainly* (no ``--convert``, which used
to run a conversion over data that had never been downloaded);
- every dispatched script must carry the selected ``--data-path``.
"""
from __future__ import annotations
import sys
import pytest
import data.download_all as da
@pytest.fixture
def recorder(monkeypatch, tmp_path):
"""Replace run_download_script with a call recorder; neutralise dotenv."""
calls: list[tuple[str, list]] = []
def _record(script_name, extra_args=None):
calls.append((script_name, list(extra_args or [])))
return True
monkeypatch.setattr(da, "run_download_script", _record)
monkeypatch.setattr(da, "load_dotenv", lambda *a, **k: None)
return calls
def _run(monkeypatch, tmp_path, *cli):
monkeypatch.setattr(sys, "argv", ["download_all.py", "--data-path", str(tmp_path), *cli])
da.main()
def _names(calls):
return {name for name, _ in calls}
def test_free_only_dispatches_core_and_factors(recorder, monkeypatch, tmp_path):
_run(monkeypatch, tmp_path, "--free-only")
names = _names(recorder)
# core case-study datasets + free factors + firm characteristics
assert {
"etfs.py",
"crypto.py",
"prediction_markets.py",
"cot.py",
"ff_factors.py",
"aqr_factors.py",
"firm_characteristics.py",
} <= names
# API-key / paid / large datasets must NOT be dispatched in free-only mode
assert "macro.py" not in names
assert "fx_pairs.py" not in names
assert "us_equities.py" not in names
def test_skip_firm_characteristics(recorder, monkeypatch, tmp_path):
_run(monkeypatch, tmp_path, "--free-only", "--skip-firm-characteristics")
assert "firm_characteristics.py" not in _names(recorder)
# the rest of the free tier still runs
assert "etfs.py" in _names(recorder)
def test_firm_characteristics_invoked_without_convert(recorder, monkeypatch, tmp_path):
"""firm-char must download plainly — never with ``--convert`` alone."""
_run(monkeypatch, tmp_path, "--free-only")
fc = [args for name, args in recorder if name == "firm_characteristics.py"]
assert len(fc) == 1
assert "--convert" not in fc[0]
assert "--data-path" in fc[0]
def test_core_mode_adds_api_key_datasets(recorder, monkeypatch, tmp_path):
"""Default (core) mode also dispatches the free-API-key datasets."""
_run(monkeypatch, tmp_path) # no --free-only
names = _names(recorder)
assert "macro.py" in names
assert "fx_pairs.py" in names
# but still not the --all-only historical equities pull
assert "us_equities.py" not in names
def test_every_dispatch_carries_selected_data_path(recorder, monkeypatch, tmp_path):
"""No dispatched script may be left to guess the data root."""
_run(monkeypatch, tmp_path, "--free-only")
for name, args in recorder:
assert str(tmp_path) in args, f"{name} did not receive the data path: {args}"
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_download_coverage.py
+90
View File
@@ -0,0 +1,90 @@
"""Phase 4 — coverage ratchet: every free dataset keeps a live drift smoke.
Phase 1 keeps ``download_all.py``'s registry in lock-step with the on-disk
scripts. Phase 3 adds a live external drift smoke per free source. This test
welds the two together so coverage can't rot as later beats ship more datasets:
a newly shipped free dataset must appear in the free-tier list (Phase 1) *and*
carry a drift test (Phase 3), or CI fails on the PR that adds it.
Deterministic and network-free — it only introspects test modules and the
download registry, so it runs in the fast per-PR ``test-unit`` job.
"""
from __future__ import annotations
import tests.test_external_drift as drift
from tests.test_download_scripts_registry import DATA_DIR, FREE_TIER_SCRIPTS
# Maps each free-tier download script (the Phase 1 SSOT) to the drift-source key
# it must be smoke-tested under in tests/test_external_drift.py. Keep this in
# lock-step with FREE_TIER_SCRIPTS — the tests below fail loudly if it drifts.
FREE_SCRIPT_TO_DRIFT_SOURCE = {
"etfs/market/download.py": "etf",
"crypto/market/download.py": "crypto",
"prediction_markets/download.py": "prediction_markets",
"futures/positioning/cot_download.py": "cot",
"factors/ff_download.py": "fama_french",
"factors/aqr_download.py": "aqr",
"equities/firm_characteristics/download.py": "firm_characteristics",
"macro/download.py": "fred",
"fx/market/download.py": "fx",
}
def _drift_test_names() -> set[str]:
return {
name for name in dir(drift) if name.startswith("test_") and callable(getattr(drift, name))
}
def test_every_free_script_is_mapped_to_a_drift_source():
"""A new free dataset can't ship without declaring its drift source.
When a later beat adds a free dataset to FREE_TIER_SCRIPTS, this fails until
the author maps it here (which in turn forces a DRIFT_SOURCES entry and a
drift test via the checks below).
"""
mapped = set(FREE_SCRIPT_TO_DRIFT_SOURCE)
free = set(FREE_TIER_SCRIPTS)
unmapped = free - mapped
assert not unmapped, (
f"free-tier download scripts with no drift-source mapping: {sorted(unmapped)} — "
f"add them to FREE_SCRIPT_TO_DRIFT_SOURCE and give each a drift test."
)
stale = mapped - free
assert not stale, (
f"drift-source mappings for scripts no longer in FREE_TIER_SCRIPTS: {sorted(stale)} — "
f"remove them (the dataset was dropped or renamed)."
)
def test_mapped_scripts_exist_on_disk():
"""Every mapped free script resolves to a real file (catches a rename)."""
missing = [rel for rel in FREE_SCRIPT_TO_DRIFT_SOURCE if not (DATA_DIR / rel).is_file()]
assert not missing, f"mapped free-tier scripts missing on disk: {missing}"
def test_drift_sources_match_mapping():
"""``DRIFT_SOURCES`` and the free-script mapping are bidirectionally consistent."""
declared = set(drift.DRIFT_SOURCES)
mapped = set(FREE_SCRIPT_TO_DRIFT_SOURCE.values())
assert declared == mapped, (
f"DRIFT_SOURCES {sorted(declared)} != mapped free sources {sorted(mapped)} — "
f"a source is declared without a script mapping (or vice versa)."
)
def test_every_drift_source_has_a_test():
"""Each declared drift source is backed by a ``test_<source>*`` function."""
names = _drift_test_names()
for source in drift.DRIFT_SOURCES:
assert any(n.startswith(f"test_{source}") for n in names), (
f"drift source '{source}' has no test_{source}* function in "
f"tests/test_external_drift.py — declared but not smoke-tested."
)
if __name__ == "__main__":
import pytest
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_download_helpers.py
+170
View File
@@ -0,0 +1,170 @@
"""Unit tests for the shared download plumbing in ``utils/downloading.py``.
These are fast, deterministic, and network-free. They pin the *contracts* that
every ``data/**/download.py`` script relies on for path resolution and output
placement — the layer where the #361-class bugs lived:
- a downloader must write under the *selected* data root (``--data-path`` /
``ML4T_DATA_PATH``), never blindly under ``<repo>/data`` (the read-only /
wrong-dir bug);
- ``~`` in a ``--config`` / ``--data-path`` must expand (the Copilot review bug);
- ``atomic_write_parquet`` must land a complete file and leave no ``.tmp`` behind.
"""
from __future__ import annotations
from pathlib import Path
import polars as pl
import pytest
from utils.downloading import (
atomic_write_parquet,
create_base_parser,
flatten_group_values,
load_section,
resolve_data_dir,
resolve_storage_path,
)
# ---------------------------------------------------------------------------
# resolve_data_dir — precedence + expansion
# ---------------------------------------------------------------------------
def test_resolve_data_dir_cli_arg_wins_over_env(tmp_path, monkeypatch):
"""An explicit ``--data-path`` must win over ``ML4T_DATA_PATH``.
This is the contract that keeps a downloader writing where the reader
asked — the mechanism behind the #361 read-only-mount fix.
"""
cli = tmp_path / "cli_root"
monkeypatch.setenv("ML4T_DATA_PATH", str(tmp_path / "env_root"))
resolved = resolve_data_dir(cli)
assert resolved == cli.resolve()
def test_resolve_data_dir_expands_user(monkeypatch, tmp_path):
"""A ``~``-prefixed CLI path is expanded, not taken literally."""
monkeypatch.setenv("HOME", str(tmp_path))
resolved = resolve_data_dir(Path("~/some_data"))
assert resolved == (tmp_path / "some_data").resolve()
assert "~" not in str(resolved)
# ---------------------------------------------------------------------------
# resolve_storage_path — the exact mechanism behind the ETF wrong-dir bug
# ---------------------------------------------------------------------------
def test_resolve_storage_path_relative_joins_data_root(tmp_path):
"""A relative configured path is joined under the selected data root."""
root = tmp_path / "data_root"
resolved = resolve_storage_path(root, "etfs/market", "etfs")
assert resolved == root / "etfs/market"
def test_resolve_storage_path_absolute_is_preserved(tmp_path):
"""An absolute configured path is preserved (not re-rooted)."""
absolute = tmp_path / "elsewhere" / "etfs"
resolved = resolve_storage_path(tmp_path / "data_root", str(absolute), "etfs")
assert resolved == absolute
def test_resolve_storage_path_falls_back_when_unconfigured(tmp_path):
"""With no configured path, the fallback is used under the data root."""
root = tmp_path / "data_root"
assert resolve_storage_path(root, None, "etfs") == root / "etfs"
def test_resolve_storage_path_expands_user(tmp_path, monkeypatch):
monkeypatch.setenv("HOME", str(tmp_path))
resolved = resolve_storage_path(tmp_path / "data_root", "~/abs_etfs", "etfs")
assert resolved == (tmp_path / "abs_etfs")
# ---------------------------------------------------------------------------
# load_section — YAML section reader used by the ETF/crypto/fx downloaders
# ---------------------------------------------------------------------------
def test_load_section_reads_named_section(tmp_path):
cfg = tmp_path / "config.yaml"
cfg.write_text("etfs:\n storage_path: etfs/market\n start: '2010-01-01'\n")
section = load_section(cfg, "etfs")
assert section == {"storage_path": "etfs/market", "start": "2010-01-01"}
def test_load_section_missing_section_returns_empty(tmp_path):
cfg = tmp_path / "config.yaml"
cfg.write_text("etfs:\n storage_path: etfs/market\n")
assert load_section(cfg, "nonexistent") == {}
def test_load_section_expands_user(tmp_path, monkeypatch):
"""``load_section`` must honour ``~`` so ``--config ~/foo.yaml`` works."""
monkeypatch.setenv("HOME", str(tmp_path))
(tmp_path / "config.yaml").write_text("etfs:\n storage_path: etfs/market\n")
section = load_section("~/config.yaml", "etfs")
assert section == {"storage_path": "etfs/market"}
# ---------------------------------------------------------------------------
# flatten_group_values — grouped symbols/pairs flattening
# ---------------------------------------------------------------------------
def test_flatten_group_values_dedups_and_preserves_order():
groups = {
"large_cap": {"symbols": ["SPY", "QQQ", "SPY"]},
"bonds": {"symbols": ["TLT", "QQQ"]},
"not_a_dict": ["ignored"],
}
assert flatten_group_values(groups, "symbols") == ["SPY", "QQQ", "TLT"]
def test_flatten_group_values_empty():
assert flatten_group_values({}, "symbols") == []
# ---------------------------------------------------------------------------
# atomic_write_parquet — complete write, no temp residue
# ---------------------------------------------------------------------------
def test_atomic_write_parquet_writes_and_roundtrips(tmp_path):
df = pl.DataFrame({"symbol": ["AAA", "BBB"], "timestamp": [1, 2]})
out = tmp_path / "nested" / "dir" / "data.parquet"
atomic_write_parquet(df, out)
assert out.exists() # parent dirs created
assert pl.read_parquet(out).equals(df)
# no leftover temp file
assert not (out.parent / f".{out.name}.tmp").exists()
assert list(out.parent.glob(".*.tmp")) == []
# ---------------------------------------------------------------------------
# create_base_parser — the standard download CLI flags
# ---------------------------------------------------------------------------
def test_create_base_parser_has_standard_flags(tmp_path):
parser = create_base_parser("test")
args = parser.parse_args(["--data-path", str(tmp_path), "--dry-run", "--force", "--verbose"])
assert args.data_path == tmp_path
assert args.dry_run is True
assert args.force is True
assert args.verbose is True
def test_create_base_parser_defaults(tmp_path):
parser = create_base_parser("test")
args = parser.parse_args([])
assert args.data_path is None
assert args.dry_run is False
assert args.force is False
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_download_scripts_registry.py
+72
View File
@@ -0,0 +1,72 @@
"""Coverage-ratchet tests for the download-script registry.
``download_all.py`` dispatches by shelling out to the paths in
``DOWNLOAD_SCRIPTS``. If a downloader is renamed or moved without updating the
map, dispatch silently degrades to "Script not found" at runtime — invisible to
CI until a reader hits it. These tests keep the map and the on-disk scripts in
lock-step, in both directions, and fail when a *new* ``download.py`` is added
without wiring it in (so coverage can't rot as later beats ship more datasets).
"""
from __future__ import annotations
from pathlib import Path
import pytest
import data.download_all as da
DATA_DIR = Path(da.__file__).parent
# The free datasets a reader gets with ``download_all.py --free-only`` (no API
# key, or a free key). Single source of truth: the registry ratchet below pins
# that these ship, and ``tests/test_download_coverage.py`` (Phase 4) pins that
# each keeps a live drift smoke — so a newly shipped free dataset can't land
# without both a registry entry and an external-drift test.
FREE_TIER_SCRIPTS = (
"etfs/market/download.py",
"crypto/market/download.py",
"prediction_markets/download.py",
"futures/positioning/cot_download.py",
"factors/ff_download.py",
"factors/aqr_download.py",
"equities/firm_characteristics/download.py",
"macro/download.py",
"fx/market/download.py",
)
def _discovered_download_scripts() -> set[str]:
"""All download scripts on disk, as paths relative to ``data/``."""
found = set(DATA_DIR.glob("**/download.py")) | set(DATA_DIR.glob("**/*_download.py"))
return {str(p.relative_to(DATA_DIR)) for p in found}
@pytest.mark.parametrize("script_name,relative_path", sorted(da.DOWNLOAD_SCRIPTS.items()))
def test_registered_script_exists(script_name, relative_path):
"""Every entry in DOWNLOAD_SCRIPTS resolves to a real file."""
assert (DATA_DIR / relative_path).is_file(), (
f"DOWNLOAD_SCRIPTS['{script_name}'] -> {relative_path} does not exist; "
"a downloader was renamed/moved without updating download_all.py"
)
def test_every_download_script_is_registered():
"""Coverage ratchet: no download script may be orphaned from the registry."""
registered = set(da.DOWNLOAD_SCRIPTS.values())
discovered = _discovered_download_scripts()
unregistered = discovered - registered
assert not unregistered, (
f"download scripts present on disk but missing from DOWNLOAD_SCRIPTS: "
f"{sorted(unregistered)} — wire them into download_all.py (or rename)."
)
def test_free_tier_scripts_present():
"""The free datasets a reader gets with `download_all.py --free-only` all ship."""
missing = [p for p in FREE_TIER_SCRIPTS if not (DATA_DIR / p).is_file()]
assert not missing, f"free-tier download scripts missing: {missing}"
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_external_drift.py
+237
View File
@@ -0,0 +1,237 @@
"""Phase 3 — live external drift smoke (weekly, flake-tolerant, auto-issue).
One tiny **real** pull per FREE data source. Each test asserts the source is
reachable, returns non-empty data, and still carries the schema our download
scripts depend on. When an upstream provider moves, renames a file, or changes
its payload shape, the matching test fails and the weekly workflow opens a
deduped GitHub issue — this is the drift detector.
Design rules (mirror ``work/2026-07-06-public-test-suite/PLAN.md``):
- **Tiny props only** — 2 symbols / a short window / a single product-year.
Never a full-universe or full-history pull.
- **No billed APIs.** Only free sources run. Key-gated free sources (FRED,
OANDA) skip cleanly when their free key is absent; geo-restricted sources
(Kalshi / Polymarket) skip when blocked rather than fail.
- **Provider-level, not script-level.** We call the same providers the
``data/**/download.py`` scripts use, so a provider/API change surfaces here
without writing files or pulling gigabytes.
Marked ``@pytest.mark.drift`` and collected only by the weekly ``drift`` job —
never per-PR (they touch the network).
``DRIFT_SOURCES`` is the public contract that ``tests/test_download_coverage.py``
(Phase 4) ratchets against: every free dataset must keep a drift smoke here.
"""
from __future__ import annotations
import os
import pytest
pytestmark = pytest.mark.drift
# The free data sources this module keeps a live drift smoke for. Phase 4's
# coverage ratchet asserts every free ``data/**/download.py`` maps to an entry
# here, so a newly shipped free dataset can't land without a drift test.
DRIFT_SOURCES: tuple[str, ...] = (
"etf", # data/etfs/market/download.py (Yahoo Finance)
"crypto", # data/crypto/market/download.py (Binance public)
"fama_french", # data/factors/ff_download.py (Ken French library)
"aqr", # data/factors/aqr_download.py (AQR research)
"cot", # data/futures/positioning/cot_download.py (CFTC)
"fred", # data/macro/download.py (FRED — free key)
"firm_characteristics", # data/equities/firm_characteristics/download.py (Google Drive)
"fx", # data/fx/market/download.py (OANDA — free key)
"prediction_markets", # data/prediction_markets/download.py (Kalshi + Polymarket)
)
# A short, recent window keeps every pull tiny and deterministic in size.
_START = "2024-01-02"
_END = "2024-01-10"
def _assert_schema(df, required: set[str], source: str) -> None:
"""A source is healthy only if it returns non-empty data with our columns."""
import polars as pl
assert isinstance(df, pl.DataFrame), f"{source}: expected a polars DataFrame"
assert not df.is_empty(), f"{source}: reachable but returned zero rows (drift?)"
missing = required - set(df.columns)
assert not missing, f"{source}: missing expected column(s) {sorted(missing)} — schema drift"
# ---------------------------------------------------------------------------
# No-key free sources — always run in the weekly job.
# ---------------------------------------------------------------------------
def test_etf_yahoo_reachable():
"""ETF universe download source (Yahoo Finance) — canonical OHLCV schema."""
from ml4t.data.providers.yahoo import YahooFinanceProvider
provider = YahooFinanceProvider()
try:
df = provider.fetch_ohlcv("SPY", _START, _END, "daily")
finally:
provider.close()
_assert_schema(df, {"timestamp", "symbol", "open", "high", "low", "close"}, "etf/yahoo")
def test_crypto_binance_reachable():
"""Crypto premium-index source (Binance public) — canonical premium schema."""
from ml4t.data.providers.binance_public import BinancePublicProvider
provider = BinancePublicProvider(market="futures")
try:
df = provider.fetch_premium_index(
"BTCUSDT", start="2024-01-01", end="2024-01-05", interval="8h"
)
finally:
provider.close()
_assert_schema(df, {"timestamp", "symbol", "premium_index_close"}, "crypto/binance")
def test_fama_french_reachable():
"""Fama-French factor source (Ken French library) — ff3 factors present."""
from ml4t.data.providers.fama_french import FamaFrenchProvider
provider = FamaFrenchProvider()
try:
df = provider.fetch("ff3", frequency="monthly", start="2023-01-01", end="2023-06-30")
finally:
provider.close()
# Ken French renames/reformats his CSV zips periodically; pin the columns
# the book's factor pipeline reads.
_assert_schema(df, {"timestamp", "Mkt-RF", "SMB", "HML", "RF"}, "fama_french")
def test_aqr_reachable():
"""AQR factor source — quality-minus-junk panel reachable (wide country panel)."""
from ml4t.data.providers.aqr import AQRFactorProvider
provider = AQRFactorProvider()
try:
# No date filter: the provider's date-string filter path is brittle and
# the monthly QMJ panel is already small. We only need reachability +
# that the Excel workbook still parses to a timestamped frame.
df = provider.fetch("qmj_factors")
finally:
provider.close()
_assert_schema(df, {"timestamp"}, "aqr")
assert df.width >= 2, "aqr: QMJ workbook parsed but carries no factor columns — layout drift"
def test_cot_cftc_reachable(tmp_path, monkeypatch):
"""CFTC Commitment-of-Traders source — one product-year panel."""
from ml4t.data.cot import COTConfig, COTFetcher
# cot_reports writes a scratch .txt into the CWD; keep it out of the repo.
monkeypatch.chdir(tmp_path)
fetcher = COTFetcher(COTConfig(products=["ES"], start_year=2024, end_year=2024))
df = fetcher.fetch_product("ES")
_assert_schema(df, {"report_date", "open_interest"}, "cot")
def test_firm_characteristics_gdrive_listing():
"""Firm-characteristics source (Google Drive folder) — listing still resolves.
Lists the folder without downloading (``skip_download=True``), catching a
moved/renamed folder or a gdown API change — the class of bug that broke the
firm-char download — for a fraction of a second and zero of the ~1.5 GB.
"""
import os
import gdown
from data.equities.firm_characteristics.download import GDRIVE_FOLDER_URL
listing = gdown.download_folder(GDRIVE_FOLDER_URL, skip_download=True, quiet=True)
assert listing, "firm_characteristics: Google Drive folder listing is empty (moved/renamed?)"
# The raw folder holds the source files the downloader fetches, then extracts:
# RetChar.csv (the ~1.1 GB characteristics table the converter reads) and
# datasets.zip (the char/macro/RF numpy splits). A rename of either breaks the
# download, so pin their presence by basename rather than the extracted layout.
names = {os.path.basename(getattr(f, "path", str(f))) for f in listing}
for required in ("RetChar.csv", "datasets.zip"):
assert required in names, (
f"firm_characteristics: '{required}' missing from Drive folder "
f"(found {sorted(names)}) — upstream dataset renamed/moved"
)
# ---------------------------------------------------------------------------
# Key-gated free sources — skip cleanly when the free key is absent (no charge).
# ---------------------------------------------------------------------------
def test_fred_reachable():
"""Treasury/macro source (FRED) — free API key required, skip if unset."""
if not os.getenv("FRED_API_KEY"):
pytest.skip("FRED_API_KEY not set — free-key source, nothing to charge")
from ml4t.data.providers.fred import FREDProvider
provider = FREDProvider()
try:
# DGS10 = 10-Year Treasury constant maturity, a stable free series.
df = provider.fetch_ohlcv("DGS10", _START, _END, "daily")
finally:
provider.close()
_assert_schema(df, {"timestamp"}, "fred")
def test_fx_oanda_reachable():
"""FX source (OANDA) — free API key required, skip if unset."""
if not os.getenv("OANDA_API_KEY"):
pytest.skip("OANDA_API_KEY not set — free-key source, nothing to charge")
from ml4t.data.providers.oanda import OandaProvider
provider = OandaProvider(api_key=os.environ["OANDA_API_KEY"])
try:
df = provider.fetch_ohlcv("EUR_USD", _START, _END, "daily")
finally:
close = getattr(provider, "close", None)
if callable(close):
close()
_assert_schema(df, {"timestamp", "open", "high", "low", "close"}, "fx/oanda")
# ---------------------------------------------------------------------------
# Geo-restricted free source — skip (not fail) when the region blocks access.
# ---------------------------------------------------------------------------
def test_prediction_markets_reachable():
"""Prediction-markets source (Kalshi) — geo-restricted, skip when blocked.
Kalshi and Polymarket restrict API access by jurisdiction, so a listing can
fail at the network layer from some regions even though the endpoints are
up. This dataset is optional; a geo/network block skips rather than fails so
the weekly job doesn't file spurious drift issues from a blocked runner.
"""
try:
from ml4t.data.providers.kalshi import KalshiProvider
except ImportError:
pytest.skip("Kalshi provider not installed")
provider = KalshiProvider()
try:
markets = provider.list_markets(limit=1)
except Exception as exc: # network/geo block — optional dataset
pytest.skip(f"Kalshi unreachable from this runner (likely geo-restricted): {exc}")
finally:
close = getattr(provider, "close", None)
if callable(close):
close()
if not markets:
pytest.skip("Kalshi reachable but returned no markets (geo-filtered)")
assert isinstance(markets, list) and markets[0].get("ticker"), (
"prediction_markets: Kalshi listing shape changed (no 'ticker') — schema drift"
)
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v", "-m", "drift"]))
tests/test_firm_characteristics_extract.py
+115
View File
@@ -0,0 +1,115 @@
"""Unit tests for the firm-characteristics archive handling.
``extract_zip`` flattens the published Chen-Pelger-Zhu archives, which wrap
their contents in a single top-level ``data/`` or ``datasets/`` directory and
(when zipped on macOS) ship ``__MACOSX`` resource forks + ``.DS_Store`` files.
A mismatch here silently lands the ``.npz`` files at the wrong depth, which
``verify_files`` then reports as "missing" — exactly the bug fixed in the
firm-char end-to-end repair. These synthetic-zip tests lock the flatten rule
without pulling the real ~1.5 GB dataset.
"""
from __future__ import annotations
import zipfile
from pathlib import Path
import pytest
from data.equities.firm_characteristics.download import EXPECTED_FILES, extract_zip
def _make_zip(zip_path: Path, entries: dict[str, bytes]) -> None:
with zipfile.ZipFile(zip_path, "w") as zf:
for arcname, content in entries.items():
zf.writestr(arcname, content)
def test_extract_zip_strips_datasets_wrapper(tmp_path):
"""A ``datasets/`` wrapper is stripped so files land directly under the dir."""
zip_path = tmp_path / "datasets.zip"
_make_zip(
zip_path,
{
"datasets/RetChar.csv": b"Date,RET\n20000101,0.01\n",
"datasets/char/Char_train.npz": b"fake-npz",
"datasets/macro/macro_test.npz": b"fake-npz",
},
)
extract_dir = tmp_path / "out"
extract_dir.mkdir()
assert extract_zip(zip_path, extract_dir) is True
assert (extract_dir / "RetChar.csv").exists()
assert (extract_dir / "char" / "Char_train.npz").exists()
assert (extract_dir / "macro" / "macro_test.npz").exists()
# wrapper directory itself must NOT survive
assert not (extract_dir / "datasets").exists()
def test_extract_zip_strips_data_wrapper(tmp_path):
"""The alternate ``data/`` wrapper is stripped identically."""
zip_path = tmp_path / "data.zip"
_make_zip(zip_path, {"data/char/Char_valid.npz": b"x"})
extract_dir = tmp_path / "out"
extract_dir.mkdir()
assert extract_zip(zip_path, extract_dir) is True
assert (extract_dir / "char" / "Char_valid.npz").exists()
assert not (extract_dir / "data").exists()
def test_extract_zip_skips_macos_cruft(tmp_path):
"""``__MACOSX`` resource forks and ``.DS_Store`` files are dropped."""
zip_path = tmp_path / "datasets.zip"
_make_zip(
zip_path,
{
"datasets/RetChar.csv": b"data",
"datasets/.DS_Store": b"junk",
"__MACOSX/datasets/._RetChar.csv": b"resource-fork",
},
)
extract_dir = tmp_path / "out"
extract_dir.mkdir()
assert extract_zip(zip_path, extract_dir) is True
assert (extract_dir / "RetChar.csv").exists()
assert not (extract_dir / ".DS_Store").exists()
assert not (extract_dir / "__MACOSX").exists()
# the resource-fork file must not have leaked in under any name
assert not any(p.name == "._RetChar.csv" for p in extract_dir.rglob("*"))
def test_extract_zip_cleans_temp_dir(tmp_path):
"""The intermediate ``_temp_extract`` scratch dir is removed."""
zip_path = tmp_path / "datasets.zip"
_make_zip(zip_path, {"datasets/RetChar.csv": b"data"})
extract_dir = tmp_path / "out"
extract_dir.mkdir()
extract_zip(zip_path, extract_dir)
assert not (extract_dir / "_temp_extract").exists()
def test_extract_zip_returns_false_on_bad_archive(tmp_path):
"""A corrupt archive fails gracefully (returns False, no raise)."""
bad = tmp_path / "broken.zip"
bad.write_bytes(b"not a real zip")
extract_dir = tmp_path / "out"
extract_dir.mkdir()
assert extract_zip(bad, extract_dir) is False
def test_expected_files_manifest_is_sane():
"""The published dataset ships 11 files across RetChar/Macro/char/macro/RF."""
assert len(EXPECTED_FILES) == 11
assert "RetChar.csv" in EXPECTED_FILES
assert all(size > 0 for size in EXPECTED_FILES.values())
# the three split families each contribute a train/valid/test triple
for prefix in ("char/", "macro/", "RF/"):
assert sum(1 for k in EXPECTED_FILES if k.startswith(prefix)) == 3
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
tests/test_futures_loader.py
+84
View File
@@ -0,0 +1,84 @@
"""Unit tests for data/futures/loader.py — list_cme_products().
The loader binds ML4T_DATA_PATH at import time (`from utils import
ML4T_DATA_PATH`), so monkeypatching the env var is insufficient —
we patch the module symbol directly.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from data.exceptions import DataNotFoundError
from data.futures import loader as futures_loader
def _build_hive_fixture(root: Path, products: list[str]) -> None:
root.mkdir(parents=True, exist_ok=True)
for p in products:
(root / f"product={p}").mkdir()
def _build_individual_fixture(root: Path, products: list[str], include_empty: bool = False) -> None:
root.mkdir(parents=True, exist_ok=True)
for p in products:
pdir = root / p
pdir.mkdir()
(pdir / "data.parquet").write_bytes(b"fake-parquet")
if include_empty:
(root / "EMPTY_DIR").mkdir() # no data.parquet — must be skipped
@pytest.fixture
def isolated_data_path(tmp_path, monkeypatch):
"""Redirect the loader's bound ML4T_DATA_PATH to a temp directory."""
monkeypatch.setattr(futures_loader, "ML4T_DATA_PATH", tmp_path)
return tmp_path
def test_list_cme_products_hourly_returns_sorted_unique(isolated_data_path) -> None:
hive_root = isolated_data_path / "futures" / "market" / "continuous" / "hourly"
_build_hive_fixture(hive_root, ["ES", "NQ", "CL", "GC", "6E"])
products = futures_loader.list_cme_products()
assert products == ["6E", "CL", "ES", "GC", "NQ"]
def test_list_cme_products_ignores_non_product_subdirs(isolated_data_path) -> None:
hive_root = isolated_data_path / "futures" / "market" / "continuous" / "hourly"
hive_root.mkdir(parents=True, exist_ok=True)
(hive_root / "product=ES").mkdir()
(hive_root / "_metadata").mkdir()
(hive_root / "random_other_dir").mkdir()
(hive_root / "readme.txt").write_text("x")
assert futures_loader.list_cme_products() == ["ES"]
def test_list_cme_products_individual_requires_data_parquet(isolated_data_path) -> None:
ind_root = isolated_data_path / "futures" / "market" / "individual"
_build_individual_fixture(ind_root, ["ES", "CL"], include_empty=True)
assert futures_loader.list_cme_products(frequency="individual") == ["CL", "ES"]
def test_list_cme_products_raises_when_hourly_root_missing(isolated_data_path) -> None:
with pytest.raises(DataNotFoundError, match="CME Futures Hourly"):
futures_loader.list_cme_products()
def test_list_cme_products_raises_when_individual_root_missing(isolated_data_path) -> None:
with pytest.raises(DataNotFoundError, match="CME Futures Individual"):
futures_loader.list_cme_products(frequency="individual")
def test_list_cme_products_rejects_unknown_frequency(isolated_data_path) -> None:
with pytest.raises(ValueError, match="frequency must be"):
futures_loader.list_cme_products(frequency="daily")
def test_list_cme_products_returns_empty_list_for_empty_hive(isolated_data_path) -> None:
(isolated_data_path / "futures" / "market" / "continuous" / "hourly").mkdir(parents=True)
assert futures_loader.list_cme_products() == []
tests/test_import_coverage.py
+172
View File
@@ -0,0 +1,172 @@
"""Scanner-driven import-coverage test.
Uses ``envs.scan_imports`` to extract every third-party top-level import
that appears anywhere in the book's source code, then verifies each one
that's expected to resolve in the current Docker image actually does.
The point of this test is **drift detection**. When a chapter adds a
new dependency, the scanner picks it up automatically — no hand-edited
list to remember to update. If the new package isn't installed in the
image a reader built (or isn't classified in ``IMAGE_OVERRIDES``), the
test fails loudly instead of waiting for a reader to hit the missing
import mid-notebook.
The current image is taken from the ``ML4T_IMAGE`` environment variable
(defaulting to ``"ml4t"``). Each Docker image's entrypoint should set
``ML4T_IMAGE=<image-id>`` so the right set of imports gets exercised.
"""
from __future__ import annotations
import importlib
import os
import pytest
from envs.scan_imports import (
IMAGE_OVERRIDES,
REPO_ROOT,
VALID_IMAGES,
classify,
scan_repo,
try_import,
)
# The default image for agent / CI runs without an explicit ML4T_IMAGE
_DEFAULT_IMAGE = "ml4t"
@pytest.fixture(scope="module")
def scanned_imports() -> set[str]:
"""Run the scanner once per module — it walks the whole tree."""
return scan_repo()
@pytest.fixture(scope="module")
def classified(scanned_imports) -> dict[str, set[str]]:
return classify(scanned_imports)
# -----------------------------------------------------------------------------
# Sanity: the scanner finds things
# -----------------------------------------------------------------------------
def test_scanner_discovers_reasonable_number_of_external_imports(scanned_imports) -> None:
"""A healthy scan finds 50-200 third-party imports across 27 chapters +
9 case studies. Outside that envelope indicates the filter is broken
(too few: stdlib/first-party leaking in; too many: everything is being
called third-party).
"""
assert 50 <= len(scanned_imports) <= 200, (
f"Scanner found {len(scanned_imports)} external imports — "
"outside the plausible range for this repo"
)
def test_scanner_finds_core_stack(scanned_imports) -> None:
"""The book's core stack must appear in every scan."""
core = {"numpy", "pandas", "polars", "matplotlib", "scipy", "sklearn", "torch"}
missing = core - scanned_imports
assert not missing, f"core stack packages not detected: {missing}"
def test_scanner_excludes_stdlib(scanned_imports) -> None:
"""Basic regression: no stdlib name should appear in external imports."""
stdlib_leak = scanned_imports & {"os", "sys", "json", "pathlib", "re", "ast"}
assert not stdlib_leak, f"stdlib names leaked into external set: {stdlib_leak}"
def test_scanner_excludes_first_party(scanned_imports) -> None:
"""First-party names must be auto-filtered by the scanner."""
first_party_leak = scanned_imports & {
"utils",
"data",
"case_studies",
"conftest",
}
assert not first_party_leak, f"first-party leaked: {first_party_leak}"
# -----------------------------------------------------------------------------
# Classification invariants
# -----------------------------------------------------------------------------
def test_every_image_override_targets_a_valid_image() -> None:
"""Every package in IMAGE_OVERRIDES must map to a recognized image id."""
invalid = {pkg: img for pkg, img in IMAGE_OVERRIDES.items() if img not in VALID_IMAGES}
assert not invalid, f"packages mapped to unknown images: {invalid}"
def test_every_override_target_appears_in_at_least_one_source_file(scanned_imports) -> None:
"""Trim dead entries: a package in IMAGE_OVERRIDES should actually be
imported somewhere. If none of the code uses it, the classification
entry is stale.
"""
stale = set(IMAGE_OVERRIDES) - scanned_imports
assert not stale, (
f"IMAGE_OVERRIDES entries no longer imported anywhere: {sorted(stale)} — "
"remove them from envs/scan_imports.py"
)
def test_classify_groups_partition_the_scanned_set(classified, scanned_imports) -> None:
"""The union of all image buckets must equal the scanned set (no orphans)."""
union: set[str] = set()
for bucket in classified.values():
union |= bucket
assert union == scanned_imports
# -----------------------------------------------------------------------------
# The real test: every expected import resolves in this image
# -----------------------------------------------------------------------------
def test_every_expected_import_resolves_in_current_image(classified) -> None:
"""Attempt to import every package the scanner classified for this image.
The image id comes from ``$ML4T_IMAGE`` (default ``ml4t``). Each Docker
entrypoint should set this variable. Running pytest locally picks up
the default.
"""
image = os.environ.get("ML4T_IMAGE", _DEFAULT_IMAGE)
expected = sorted(classified[image])
failures: list[tuple[str, str]] = []
for pkg in expected:
ok, err = try_import(pkg)
if not ok:
failures.append((pkg, err))
if failures:
lines = "\n".join(f" {pkg}: {err[:120]}" for pkg, err in failures)
pytest.fail(
f"{len(failures)} of {len(expected)} expected imports failed in "
f"image={image!r}:\n{lines}"
)
# -----------------------------------------------------------------------------
# Smoke: first-party packages we ship must import too (readers installed us)
# -----------------------------------------------------------------------------
@pytest.mark.parametrize(
"module",
[
"utils.paths",
"utils.modeling",
"data",
"case_studies.utils.analytics",
"case_studies.utils.signals",
"case_studies.utils.allocation",
"case_studies.utils.registry.metrics",
],
)
def test_first_party_modules_import(module: str) -> None:
"""First-party packages must load cleanly — regression guard for refactors
that break the data/* or case_studies/* package tree."""
ok, err = try_import(module)
assert ok, f"first-party import failed: {module}: {err[:200]}"
tests/test_latent_factors_no_leak.py
+347
View File
@@ -0,0 +1,347 @@
"""Regression tests for the latent factor forecasting contracts."""
from __future__ import annotations
from datetime import datetime
import numpy as np
import polars as pl
import pytest
from case_studies.utils.latent_factors.panel import compute_managed_portfolios
def test_managed_portfolios_are_cross_sectionally_shared() -> None:
rng = np.random.default_rng(1)
chars = rng.normal(size=(12, 8, 4)).astype(np.float32)
returns = rng.normal(size=(12, 8)).astype(np.float32)
portfolios = compute_managed_portfolios(chars, returns)
for date_idx in range(portfolios.shape[0]):
assert np.allclose(portfolios[date_idx, :1, :], portfolios[date_idx]), (
f"managed portfolios vary within date {date_idx}"
)
def test_managed_portfolios_use_current_date_only() -> None:
rng = np.random.default_rng(2)
chars = rng.normal(size=(10, 6, 3)).astype(np.float32)
returns = rng.normal(size=(10, 6)).astype(np.float32)
portfolios_a = compute_managed_portfolios(chars, returns)
perturbed = returns.copy()
perturbed[4] += 100.0
portfolios_b = compute_managed_portfolios(chars, perturbed)
changed = np.abs(portfolios_a - portfolios_b).max(axis=(1, 2))
assert changed[4] > 0.0
assert np.all(changed[np.arange(len(changed)) != 4] == 0.0)
def test_cae_validation_batch_receives_validation_returns(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from case_studies.utils.latent_factors import library_bridge
rng = np.random.default_rng(123)
chars_train = rng.normal(size=(12, 8, 4)).astype(np.float32)
returns_train = rng.normal(size=(12, 8)).astype(np.float32) * 0.02
chars_val = rng.normal(size=(5, 8, 4)).astype(np.float32)
returns_val = rng.normal(size=(5, 8)).astype(np.float32) * 0.02
captured: dict[str, np.ndarray] = {}
def capture_pipeline(**kwargs):
captured["validation_returns"] = kwargs["val_batch"].returns.copy()
return {
"checkpoint_predictions": {0: np.zeros_like(returns_val)},
"checkpoint_epochs": [0],
}
monkeypatch.setattr(library_bridge, "_run_checkpointed_latent_pipeline", capture_pipeline)
library_bridge.run_cae_fold_with_library(
chars_train,
returns_train,
chars_val=chars_val,
returns_val=returns_val,
n_factors=2,
n_epochs=2,
n_ensemble=1,
hidden_units=(8,),
checkpoint_epochs=[2],
)
assert np.array_equal(captured["validation_returns"], returns_val)
@pytest.mark.gpu
def test_cae_predictions_independent_of_validation_returns() -> None:
"""End-to-end regression: perturbing validation returns must not change predictions.
Restores the byte-identical fit-twice check that the wiring-only test above
cannot enforce — guards against future changes in
`_run_checkpointed_latent_pipeline` (or anything downstream of
`model.fit(..., validation_batch=val_batch)`) that accidentally let
validation returns influence the fitted model.
"""
pytest.importorskip("torch")
from case_studies.utils.latent_factors.cae import run_cae_fold
rng = np.random.default_rng(31)
chars_train = rng.normal(size=(12, 8, 4)).astype(np.float32)
returns_train = rng.normal(size=(12, 8)).astype(np.float32) * 0.02
chars_val = rng.normal(size=(5, 8, 4)).astype(np.float32)
returns_val = rng.normal(size=(5, 8)).astype(np.float32) * 0.02
base_preds_by_epoch, _ = run_cae_fold(
chars_train,
returns_train,
chars_val,
returns_val,
n_factors=2,
n_epochs=2,
checkpoint_epochs=[2],
hidden_units=(8,),
log_fn=lambda *args, **kwargs: None,
)
perturbed_val = returns_val.copy()
perturbed_val += 100.0
perturbed_preds_by_epoch, _ = run_cae_fold(
chars_train,
returns_train,
chars_val,
perturbed_val,
n_factors=2,
n_epochs=2,
checkpoint_epochs=[2],
hidden_units=(8,),
log_fn=lambda *args, **kwargs: None,
)
base_preds = base_preds_by_epoch[2]
perturbed_preds = perturbed_preds_by_epoch[2]
assert np.array_equal(base_preds, perturbed_preds), (
"CAE predictions changed when validation returns were perturbed by +100 — "
"validation returns are leaking into the fitted model"
)
@pytest.mark.gpu
def test_cae_classification_uses_continuous_factor_returns(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("torch")
from case_studies.utils.latent_factors import library_bridge
captured: dict[str, np.ndarray] = {}
original_cross_section_batch = library_bridge._cross_section_batch
def capture_batch(
characteristics: np.ndarray,
*,
returns: np.ndarray | None = None,
factor_returns: np.ndarray | None = None,
context_features: np.ndarray | None = None,
):
if factor_returns is not None:
captured["factor_returns"] = factor_returns.copy()
return original_cross_section_batch(
characteristics,
returns=returns,
factor_returns=factor_returns,
context_features=context_features,
)
monkeypatch.setattr(library_bridge, "_cross_section_batch", capture_batch)
rng = np.random.default_rng(7)
chars = rng.normal(size=(24, 10, 5)).astype(np.float32)
class_labels = (rng.random(size=(24, 10)) > 0.5).astype(np.float32)
factor_returns = rng.normal(size=(24, 10)).astype(np.float32) * 0.02
from case_studies.utils.latent_factors.cae import run_cae_fold
run_cae_fold(
chars[:18],
class_labels[:18],
chars[18:],
class_labels[18:],
n_factors=2,
factor_returns_train=factor_returns[:18],
n_epochs=1,
checkpoint_epochs=[1],
hidden_units=(8,),
task_type="classification",
log_fn=lambda *args, **kwargs: None,
)
assert np.array_equal(captured["factor_returns"], factor_returns[:18])
def test_reporting_epoch_defaults_to_last_checkpoint() -> None:
from case_studies.utils.latent_factors.cv import _select_reporting_epoch
metrics = pl.DataFrame(
{
"fold_id": [0, 0, 1, 1],
"epoch": [5, 10, 5, 10],
"ic_mean": [0.12, 0.03, 0.11, 0.02],
}
)
epoch, mean_ic = _select_reporting_epoch(
metrics,
checkpoint_selection_policy="fixed",
reporting_epoch=None,
)
assert epoch == 10
assert mean_ic == pytest.approx(0.025)
def test_reporting_epoch_prefers_validation_best_checkpoint_zero() -> None:
from case_studies.utils.latent_factors.cv import _select_reporting_epoch
metrics = pl.DataFrame(
{
"fold_id": [0, 0, 1, 1],
"epoch": [0, 10, 0, 10],
"ic_mean": [0.04, 0.03, 0.05, 0.02],
}
)
epoch, mean_ic = _select_reporting_epoch(
metrics,
checkpoint_selection_policy="fixed",
reporting_epoch=None,
)
assert epoch == 0
assert mean_ic == pytest.approx(0.045)
def test_prediction_frame_preserves_temporal_timestamp_dtype() -> None:
from case_studies.utils.latent_factors.cv import _build_prediction_frame
predictions = np.array([[0.1, np.nan, 0.3], [0.4, 0.5, np.nan]], dtype=np.float64)
returns_val = np.array([[0.0, np.nan, 1.0], [1.0, 2.0, np.nan]], dtype=np.float64)
val_dates = np.array(["2024-01-31", "2024-02-29"], dtype="datetime64[ns]")
val_entities = np.array(
[["A", "B", "C"], ["A", "B", "C"]],
dtype=object,
)
frame = _build_prediction_frame(
predictions=predictions,
returns_val=returns_val,
eval_returns_val=None,
val_dates=val_dates,
val_entities=val_entities,
fold_id=0,
model_name="ipca",
epoch=0,
)
assert frame is not None
assert frame["timestamp"].dtype.is_temporal()
assert frame["timestamp"].to_list() == [
datetime(2024, 1, 31),
datetime(2024, 1, 31),
datetime(2024, 2, 29),
datetime(2024, 2, 29),
]
def test_rebalance_scoring_thins_to_declared_schedule() -> None:
from case_studies.utils.latent_factors.cv import _compute_frame_ic, _score_prediction_frame
frame = pl.DataFrame(
{
"timestamp": [
datetime(2024, 1, 15),
datetime(2024, 1, 15),
datetime(2024, 1, 15),
datetime(2024, 1, 15),
datetime(2024, 1, 15),
datetime(2024, 1, 31),
datetime(2024, 1, 31),
datetime(2024, 1, 31),
datetime(2024, 1, 31),
datetime(2024, 1, 31),
datetime(2024, 2, 15),
datetime(2024, 2, 15),
datetime(2024, 2, 15),
datetime(2024, 2, 15),
datetime(2024, 2, 15),
datetime(2024, 2, 29),
datetime(2024, 2, 29),
datetime(2024, 2, 29),
datetime(2024, 2, 29),
datetime(2024, 2, 29),
],
"symbol": ["A", "B", "C", "D", "E"] * 4,
"y_true": [
0.0,
1.0,
2.0,
3.0,
4.0,
0.0,
1.0,
2.0,
3.0,
4.0,
4.0,
3.0,
2.0,
1.0,
0.0,
4.0,
3.0,
2.0,
1.0,
0.0,
],
"y_score": [
0.0,
1.0,
2.0,
3.0,
4.0,
4.0,
3.0,
2.0,
1.0,
0.0,
0.0,
1.0,
2.0,
3.0,
4.0,
4.0,
3.0,
2.0,
1.0,
0.0,
],
"fold_id": [0] * 20,
"config_name": ["cae"] * 20,
"epoch": [10] * 20,
}
)
_, full_periods = _compute_frame_ic(frame)
thinned = _score_prediction_frame(
frame,
score_dates="rebalance",
score_cadence="monthly_month_end",
score_rebalance_step=1,
)
_, thinned_periods = _compute_frame_ic(thinned)
assert full_periods == 4
assert thinned_periods == 2
assert thinned is not None
assert thinned["timestamp"].unique().sort().to_list() == [
datetime(2024, 1, 31),
datetime(2024, 2, 29),
]
tests/test_list_loaders.py
+132
View File
@@ -0,0 +1,132 @@
"""Unit tests for list_etfs / list_crypto_perps / list_fx_pairs.
Parallel to test_futures_loader.py — each loader binds ML4T_DATA_PATH at
import time, so we monkeypatch the module symbol rather than the env var.
The helpers enumerate the symbol universe from each dataset's parquet
and are the canonical way to answer "what's available locally?" for
marketing / data-inventory / Ch2 EDA notebooks.
"""
from __future__ import annotations
from pathlib import Path
import polars as pl
import pytest
from data.crypto import loader as crypto_loader
from data.etfs import loader as etfs_loader
from data.exceptions import DataNotFoundError
from data.fx import loader as fx_loader
def _write_symbol_parquet(path: Path, symbols: list[str]) -> None:
"""Minimal parquet with a ``symbol`` column plus a dummy value column."""
path.parent.mkdir(parents=True, exist_ok=True)
pl.DataFrame({"symbol": symbols, "close": [1.0] * len(symbols)}).write_parquet(path)
# -----------------------------------------------------------------------------
# list_etfs
# -----------------------------------------------------------------------------
@pytest.fixture
def etfs_isolated(tmp_path, monkeypatch):
monkeypatch.setattr(etfs_loader, "ML4T_DATA_PATH", tmp_path)
return tmp_path
def test_list_etfs_returns_sorted_unique_symbols(etfs_isolated) -> None:
_write_symbol_parquet(
etfs_isolated / "etfs" / "market" / "etf_universe.parquet",
# deliberately unsorted with a duplicate
["SPY", "QQQ", "AGG", "SPY", "IWM"],
)
assert etfs_loader.list_etfs() == ["AGG", "IWM", "QQQ", "SPY"]
def test_list_etfs_raises_when_parquet_missing(etfs_isolated) -> None:
with pytest.raises(DataNotFoundError, match="ETF Universe"):
etfs_loader.list_etfs()
# -----------------------------------------------------------------------------
# list_crypto_perps
# -----------------------------------------------------------------------------
@pytest.fixture
def crypto_isolated(tmp_path, monkeypatch):
monkeypatch.setattr(crypto_loader, "ML4T_DATA_PATH", tmp_path)
return tmp_path
def test_list_crypto_perps_returns_sorted_unique_symbols(crypto_isolated) -> None:
_write_symbol_parquet(
crypto_isolated / "crypto" / "market" / "perps_1h.parquet",
["BTCUSDT", "ETHUSDT", "ADAUSDT", "BTCUSDT"],
)
assert crypto_loader.list_crypto_perps() == ["ADAUSDT", "BTCUSDT", "ETHUSDT"]
def test_list_crypto_perps_raises_when_parquet_missing(crypto_isolated) -> None:
with pytest.raises(DataNotFoundError, match="Crypto Perpetuals"):
crypto_loader.list_crypto_perps()
# -----------------------------------------------------------------------------
# list_fx_pairs
# -----------------------------------------------------------------------------
@pytest.fixture
def fx_isolated(tmp_path, monkeypatch):
monkeypatch.setattr(fx_loader, "ML4T_DATA_PATH", tmp_path)
return tmp_path
def test_list_fx_pairs_default_probes_daily(fx_isolated) -> None:
_write_symbol_parquet(
fx_isolated / "fx" / "market" / "daily.parquet",
["EUR_USD", "GBP_USD", "USD_JPY"],
)
assert fx_loader.list_fx_pairs() == ["EUR_USD", "GBP_USD", "USD_JPY"]
def test_list_fx_pairs_4h_uses_separate_parquet(fx_isolated) -> None:
_write_symbol_parquet(
fx_isolated / "fx" / "market" / "4h.parquet",
["AUD_JPY", "AUD_NZD"],
)
assert fx_loader.list_fx_pairs(frequency="4h") == ["AUD_JPY", "AUD_NZD"]
def test_list_fx_pairs_raises_when_parquet_missing(fx_isolated) -> None:
with pytest.raises(DataNotFoundError, match="FX Pairs"):
fx_loader.list_fx_pairs()
def test_list_fx_pairs_rejects_unknown_frequency(fx_isolated) -> None:
with pytest.raises(ValueError, match="frequency must be"):
fx_loader.list_fx_pairs(frequency="hourly")
# -----------------------------------------------------------------------------
# Re-export from data/__init__.py
# -----------------------------------------------------------------------------
def test_list_helpers_are_exported_from_data_package() -> None:
"""All three list_*() helpers must be importable at ``from data import ...``
so marketing/inventory consumers don't need to know submodule layout.
"""
import data
assert hasattr(data, "list_etfs")
assert hasattr(data, "list_crypto_perps")
assert hasattr(data, "list_fx_pairs")
assert "list_etfs" in data.__all__
assert "list_crypto_perps" in data.__all__
assert "list_fx_pairs" in data.__all__
tests/test_model_registry.py
+445
View File
@@ -0,0 +1,445 @@
"""Test model notebooks produce correct registry entries in isolation.
Runs each case study model notebook (stage >= 06) with minimal parameters
in an isolated environment. Production data is read via symlinks; all
writes (registry.db, predictions, results JSON) go to a temp directory.
The production registry is NEVER opened or touched.
Design:
1. Session fixture creates temp dir with symlinked read-only data
2. Each notebook runs via Papermill with aggressive param reduction
3. ML4T_OUTPUT_DIR redirects all get_case_study_dir() writes to temp
4. After each run, query the test registry.db for expected entries
The goal is code-path coverage, not model quality. Params are set to the
absolute minimum that still exercises the training→register→predict loop:
MAX_SYMBOLS=3, MAX_FOLDS=2, N_EPOCHS=2, NUM_BOOST_ROUND=20.
Usage:
# All model notebooks (~15-20 min)
uv run pytest tests/test_model_registry.py -v
# Specific case study
uv run pytest tests/test_model_registry.py -v -k "crypto_perps_funding"
# Specific model family across all case studies
uv run pytest tests/test_model_registry.py -v -k "06_linear"
# Single notebook
uv run pytest tests/test_model_registry.py -v -k "etfs and 06_linear"
# Dry run — see what would be tested
uv run pytest tests/test_model_registry.py --collect-only
"""
import re
import sqlite3
from pathlib import Path
import pytest
from tests.pm_helpers import get_overrides, run_notebook
REPO_ROOT = Path(__file__).parent.parent
PROD_CS_DIR = REPO_ROOT / "case_studies"
# Ordered smallest-to-largest for faster feedback
CASE_STUDIES = [
"crypto_perps_funding",
"fx_pairs",
"cme_futures",
"etfs",
"sp500_options",
"nasdaq100_microstructure",
"sp500_equity_option_analytics",
"us_firm_characteristics",
"us_equities_panel",
]
# Directories containing production pipeline artifacts (read-only).
# Symlinked into the test output directory so model notebooks can read them.
# Everything else (run_log/, results/, models/) is created fresh for writes.
_READ_ONLY_DIRS = {"config", "features", "labels"}
# Minimum stage number for model notebooks
_MODEL_STAGE_MIN = 6
# Suffixes that are NOT model notebooks (backtest, strategy, diagnostics).
# These depend on upstream predictions and should be tested separately.
_EXCLUDED_SUFFIXES = frozenset(
{
"backtest",
"backtest_sweep",
"backtest_analysis",
"portfolio_management",
"costs",
"risk_management",
"model_analysis",
"strategy_analysis",
"synthesis",
"ic_diagnostic",
"prediction_ingestion",
}
)
# Latent factor models need more symbols than other families because factor
# extraction requires a cross-section wide enough for the covariance matrix.
_LATENT_FACTOR_SUFFIXES = frozenset(
{
"latent_factors",
"pca",
"ipca",
"sdf",
"cae",
"sae",
"term_structure_pca",
}
)
_LATENT_FACTOR_OVERRIDES = {
"MAX_SYMBOLS": 10,
"N_FACTORS": 3,
}
# Case studies with sparse data (monthly frequency) need more symbols
# to have enough observations for CV splits.
_SPARSE_DATA_CASE_STUDIES = frozenset({"us_firm_characteristics"})
_SPARSE_DATA_OVERRIDES = {"MAX_SYMBOLS": 20}
# Minimal parameters for code-path coverage. Applied LAST so they
# override anything from overrides.yaml — we want the absolute minimum
# that still exercises the full train→register→predict loop.
_QUICK_PARAMS = {
"MAX_SYMBOLS": 3,
"MAX_FOLDS": 2,
"N_EPOCHS": 2,
"NUM_BOOST_ROUND": 20,
"BATCH_SIZE": 64,
"LOOKBACK": 24, # PatchTST needs lookback + stride >= patch_len (≥8); 24 leaves margin
"MAX_SAMPLES": 1000,
"CV_FOLDS": 2,
"N_PLACEBO": 3,
"N_FACTORS": 2,
"FORCE_RETRAIN": True,
}
# Model suffixes known to use register=True (training_runs + prediction_sets).
# Matched against the suffix after the NN_ prefix, since notebook numbers
# vary across case studies (e.g. causal_dml is 10, 11, 12, or 13 depending
# on the case study).
# Built from: grep -l "register=True" case_studies/*/[0-9][0-9]_*.py
_REGISTERING_SUFFIXES = frozenset(
{
"linear",
"gbm",
"tabular_dl",
"dl_lstm",
"dl_patchtst",
"dl_tsmixer",
"dl_nlinear",
"dl_tcn",
# NOTE: causal_dml notebooks register to ``causal_runs`` (DML effect
# estimates), not ``training_runs`` — so they are intentionally NOT in
# this set. Likewise ``NN_latent_factors`` is a thin index notebook that
# only displays the best already-registered factor IC; the factor models
# themselves register under their own sub-stems (pca/ipca/sdf/cae/sae,
# listed below). Both still execute; only the training-run-registration
# assertion is skipped for them.
"ipca",
"pca",
"sdf",
"cae",
"sae",
}
)
# DL notebooks use entry_point = "dl_{model}" (e.g. "dl_lstm") instead of
# the full filename stem (e.g. "09_dl_lstm"). Map stage stems to actual
# entry_point values for these notebooks.
_DL_RE = re.compile(r"^\d{2}_(dl_.+)$")
def _expected_entry_point(stage: str) -> str:
"""Return the entry_point value the notebook will use in the registry."""
m = _DL_RE.match(stage)
if m:
return m.group(1) # "09_dl_lstm" → "dl_lstm"
return stage # "06_linear" → "06_linear"
_STAGE_RE = re.compile(r"^(\d{2})_")
# ---------------------------------------------------------------------------
# Test collection
# ---------------------------------------------------------------------------
def _collect_model_notebooks() -> list[tuple[str, str, Path]]:
"""Discover all model notebooks (stage >= 06) across case studies.
Returns (case_study, stage_stem, notebook_path) tuples sorted by
case study order then filename within each case study.
"""
tests = []
for cs in CASE_STUDIES:
cs_dir = PROD_CS_DIR / cs
if not cs_dir.exists():
continue
for notebook in sorted(cs_dir.glob("[0-9][0-9]_*.py")):
if notebook.name.startswith("_"):
continue
match = _STAGE_RE.match(notebook.name)
if not match:
continue
stage_num = int(match.group(1))
if stage_num < _MODEL_STAGE_MIN:
continue
# Skip non-model notebooks (backtest, strategy, diagnostics)
suffix = notebook.stem[len(match.group(0)) :]
if suffix in _EXCLUDED_SUFFIXES:
continue
tests.append((cs, notebook.stem, notebook))
return tests
MODEL_TESTS = _collect_model_notebooks()
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(scope="session")
def isolated_model_output(tmp_path_factory):
"""Create an isolated output directory with symlinked production data.
For each case study, symlinks read-only directories (config/, features/,
labels/) from the production case study directory so that model notebooks
can load upstream artifacts. Write-target directories (run_log/, results/,
models/) are NOT symlinked — they are created fresh by the notebooks.
Returns the temp root directory (passed as output_dir to run_notebook,
which sets ML4T_OUTPUT_DIR).
"""
import shutil
test_root = tmp_path_factory.mktemp("model_registry_test")
for cs in CASE_STUDIES:
prod_cs = PROD_CS_DIR / cs
if not prod_cs.exists():
continue
test_cs = test_root / cs
test_cs.mkdir()
for subdir in _READ_ONLY_DIRS:
src = prod_cs / subdir
if src.exists():
(test_cs / subdir).symlink_to(src.resolve())
# Seed the global preset library (case_studies/config/{model_type}/*.yaml).
# load_configs() resolves presets at {case_dir.parent}/config/, which maps
# to test_root/config/ when ML4T_OUTPUT_DIR is set. Without this, every
# notebook that loads GBM/DL/TabDL/latent/causal presets fails.
global_config_src = PROD_CS_DIR / "config"
global_config_dst = test_root / "config"
if global_config_src.exists():
shutil.copytree(global_config_src, global_config_dst)
# Patch presets for minimal runtime (2 epochs, etc.)
from tests.conftest import _patch_presets_for_testing
_patch_presets_for_testing(global_config_dst)
return test_root
LOG_PATH = Path("/tmp/model_registry_test.log")
@pytest.fixture(scope="session", autouse=True)
def _init_log():
"""Initialize the progress log and route Papermill cell output to it."""
import logging
import time
with open(LOG_PATH, "w") as f:
f.write(f"[{time.strftime('%H:%M:%S')}] === Model Registry Test Suite ===\n")
f.write(f"[{time.strftime('%H:%M:%S')}] {len(MODEL_TESTS)} tests collected\n")
f.flush()
# Route Papermill's cell-level progress + notebook print() output to log file.
# Papermill uses "papermill" logger (not "papermill.execute") for cell markers
# and captured output. We add a file handler so it goes to our log regardless
# of pytest's log level.
handler = logging.FileHandler(LOG_PATH)
handler.setFormatter(logging.Formatter("[%(asctime)s] %(message)s", datefmt="%H:%M:%S"))
for logger_name in ("papermill", "papermill.execute"):
logger = logging.getLogger(logger_name)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.propagate = False # Don't pollute pytest captured output
# ---------------------------------------------------------------------------
# Registry helpers
# ---------------------------------------------------------------------------
def _query_registry(db_path: Path, table: str, where: str = "") -> list[dict]:
"""Query a registry table and return rows as dicts."""
if not db_path.exists():
return []
db = sqlite3.connect(str(db_path))
db.row_factory = sqlite3.Row
try:
sql = f"SELECT * FROM {table}"
if where:
sql += f" WHERE {where}"
return [dict(r) for r in db.execute(sql).fetchall()]
except sqlite3.OperationalError:
return []
finally:
db.close()
def _registry_summary(db_path: Path) -> dict:
"""Return a summary of registry contents for reporting."""
return {
"training_runs": len(_query_registry(db_path, "training_runs")),
"prediction_sets": len(_query_registry(db_path, "prediction_sets")),
"prediction_metrics": len(_query_registry(db_path, "prediction_metrics")),
}
# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"case_study,stage,notebook_path",
MODEL_TESTS,
ids=[f"{cs}::{stage}" for cs, stage, _ in MODEL_TESTS],
)
def test_model_notebook(case_study, stage, notebook_path, isolated_model_output):
"""Run a model notebook in isolation and verify registry output.
Steps:
1. Load per-notebook overrides (timeout, parameters, skip/gpu flags)
2. Merge with default reduced parameters (MAX_SYMBOLS=15, MAX_FOLDS=2)
3. Execute via Papermill with ML4T_OUTPUT_DIR → isolated temp dir
4. Assert successful completion
5. For notebooks with register=True, assert registry entries exist
"""
# --- Skip / override handling ---
rel_path = notebook_path.relative_to(REPO_ROOT).with_suffix("")
overrides = get_overrides(str(rel_path))
if overrides.get("skip"):
pytest.skip(overrides.get("skip_reason", "marked skip in overrides"))
if overrides.get("gpu"):
try:
import torch
if not torch.cuda.is_available():
pytest.skip("GPU required but not available")
except ImportError:
pytest.skip("torch not installed")
# --- Parameters ---
# Start with overrides.yaml, then apply ALL quick-test params on top.
# Quick params win — we want minimal runtime, not overrides.yaml scale.
# Papermill warns (but doesn't error) about unknown parameters, so it's
# safe to inject all of them even if the notebook doesn't use them all.
override_params = overrides.get("parameters", {})
parameters = {**override_params, **_QUICK_PARAMS}
# Latent factor models need a wider cross-section for factor extraction
stage_match_p = _STAGE_RE.match(stage)
suffix_p = stage[len(stage_match_p.group(0)) :] if stage_match_p else stage
if suffix_p in _LATENT_FACTOR_SUFFIXES:
parameters.update(_LATENT_FACTOR_OVERRIDES)
# Sparse-data case studies (monthly frequency) need more symbols
if case_study in _SPARSE_DATA_CASE_STUDIES:
parameters.update(_SPARSE_DATA_OVERRIDES)
default_timeout = 600 if suffix_p in _LATENT_FACTOR_SUFFIXES else 300
timeout = overrides.get("timeout", default_timeout)
# --- Snapshot registry state before run ---
registry_db = isolated_model_output / case_study / "run_log" / "registry.db"
before = _registry_summary(registry_db)
# --- Execute ---
result = run_notebook(
py_path=notebook_path,
parameters=parameters,
timeout=timeout,
output_dir=isolated_model_output,
log_path=LOG_PATH,
)
assert result["status"] == "ok", (
f"\n{'=' * 70}\nFAILED: {case_study}::{stage}\n{'=' * 70}\n{result['error']}\n{'=' * 70}"
)
# --- Registry assertions (for notebooks that register) ---
after = _registry_summary(registry_db)
# Check if this notebook is expected to register (match on suffix)
stage_match = _STAGE_RE.match(stage)
suffix = stage[len(stage_match.group(0)) :] if stage_match else stage
expects_registration = suffix in _REGISTERING_SUFFIXES
if expects_registration:
new_training = after["training_runs"] - before["training_runs"]
new_predictions = after["prediction_sets"] - before["prediction_sets"]
# Check for new entries OR updated entries (upserts).
# Some notebooks (e.g. 12_pca) re-register configs that were
# already created by an earlier notebook (11_latent_factors),
# resulting in upserts with 0 net new rows but updated entry_points.
expected_ep = _expected_entry_point(stage)
runs = _query_registry(
registry_db,
"training_runs",
f"entry_point = '{expected_ep}'",
)
if new_training > 0:
assert new_predictions > 0, (
f"{case_study}::{stage} created {new_training} training_runs "
f"but 0 new prediction_sets"
)
print(
f"\n Registry OK: +{new_training} training_runs, "
f"+{new_predictions} prediction_sets"
)
elif len(runs) > 0:
print(
f"\n Registry OK: {len(runs)} training_runs with "
f"entry_point='{expected_ep}' (upserted, no net new rows)"
)
else:
# Neither new entries nor matching entry_points — real failure
msg = (
f"{case_study}::{stage} has register=True but created "
f"0 new training_runs and found 0 with "
f"entry_point='{expected_ep}' (total: {after['training_runs']})"
)
raise AssertionError(msg)
else:
# Non-registering notebook — just report what happened
new_training = after["training_runs"] - before["training_runs"]
if new_training > 0:
print(
f"\n Note: {stage} created {new_training} training_runs "
f"(not in _REGISTERING_NOTEBOOKS set — consider adding)"
)
else:
print("\n OK (no registry writes expected)")
tests/test_notebook_output_hygiene.py
+37
View File
@@ -0,0 +1,37 @@
"""Guard: committed notebooks must not leak machine-specific absolute paths.
Notebooks executed on a contributor's machine can bake ``/home/<user>/...``
paths into committed cell outputs and papermill metadata. Readers should never
see those. This test scans every tracked ``.ipynb`` and fails if any survive.
To fix a failure, run::
uv run python scripts/sanitize_notebook_paths.py
which rewrites repo-internal paths to repo-relative form. See that script and
``utils.paths.display_path`` for the source-side helper.
"""
from __future__ import annotations
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "scripts"))
from sanitize_notebook_paths import _iter_notebooks, sanitize_text # noqa: E402
def test_no_machine_specific_paths_in_committed_notebooks() -> None:
offenders: list[str] = []
for nb in _iter_notebooks():
raw = nb.read_text(encoding="utf-8")
_, n = sanitize_text(raw)
if n:
offenders.append(f"{nb.relative_to(REPO_ROOT)} ({n})")
assert not offenders, (
"Notebooks leak machine-specific absolute paths in their committed "
"outputs/metadata. Run `uv run python scripts/sanitize_notebook_paths.py` "
"to fix:\n " + "\n ".join(offenders)
)
tests/test_notebook_sync.py
+40
View File
@@ -0,0 +1,40 @@
"""Gate: a committed notebook must be its current .py executed in production.
Stamped notebooks carry ``metadata.ml4t_provenance`` recording the git blob of the
paired ``.py`` they were executed from and whether the run used production
parameters. This test fails if any *stamped* notebook is stale (its ``.py`` changed
since execution) or was committed from a TEST-mode run.
Unstamped notebooks are not failed here (adoption is gradual — stamp notebooks as
they are re-run through the canonical path). See
``scripts/notebook_provenance.py`` for the stamp/check tool. To stamp::
uv run python scripts/notebook_provenance.py stamp <nb.ipynb> --executor <env>
"""
from __future__ import annotations
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "scripts"))
from notebook_provenance import check_all # noqa: E402
def test_stamped_notebooks_are_current_and_production() -> None:
stale, testmode, _unverified = check_all(strict=False)
assert not stale and not testmode, (
"Committed notebooks are out of sync with their source .py:\n"
+ (
" STALE (re-run in the canonical env):\n " + "\n ".join(stale) + "\n"
if stale
else ""
)
+ (
" TEST-MODE (must be a production run):\n " + "\n ".join(testmode)
if testmode
else ""
)
)
tests/test_paths.py
+151
View File
@@ -0,0 +1,151 @@
"""Tests for utils/paths.py — chapter/case-study path resolution.
Pins the redirection semantics of ML4T_OUTPUT_DIR (test isolation env var)
and the input validation on the chapter/strategy registries. The redirection
behavior is load-bearing: every case-study pipeline notebook depends on it
to avoid overwriting production artifacts during tests.
"""
from __future__ import annotations
import pytest
from utils.paths import (
CHAPTERS,
REPO_ROOT,
STRATEGY_IDS,
get_case_study_dir,
get_chapter_dir,
get_output_dir,
)
# -----------------------------------------------------------------------------
# Registry invariants
# -----------------------------------------------------------------------------
def test_chapters_registry_covers_1_through_27() -> None:
"""The book ships 27 chapters; the registry must enumerate all of them."""
assert set(CHAPTERS.keys()) == set(range(1, 28))
def test_chapter_directory_names_are_prefixed_by_number() -> None:
"""Notebook discovery in tests/pm_helpers relies on the NN_ prefix pattern."""
for n, dirname in CHAPTERS.items():
assert dirname.startswith(f"{n:02d}_"), f"chapter {n} dir {dirname!r} missing prefix"
def test_strategy_ids_match_case_studies_dir() -> None:
"""STRATEGY_IDS should mirror case_studies/<id>/ on disk (structure invariant)."""
cs_root = REPO_ROOT / "case_studies"
on_disk = {
p.name
for p in cs_root.iterdir()
if p.is_dir() and not p.name.startswith(("_", ".")) and p.name not in {"utils", "config"}
}
# STRATEGY_IDS is the enforced registry; on_disk may contain extras (ignored_subdirs)
# but every declared id must exist on disk.
missing = STRATEGY_IDS - on_disk
assert not missing, f"STRATEGY_IDS declare non-existent case studies: {missing}"
# -----------------------------------------------------------------------------
# get_chapter_dir
# -----------------------------------------------------------------------------
def test_get_chapter_dir_returns_absolute_path() -> None:
path = get_chapter_dir(7)
assert path.is_absolute()
assert path.name == "07_defining_the_learning_task"
def test_get_chapter_dir_rejects_out_of_range() -> None:
with pytest.raises(ValueError, match="Invalid chapter"):
get_chapter_dir(99)
def test_get_chapter_dir_rejects_zero() -> None:
with pytest.raises(ValueError, match="Invalid chapter"):
get_chapter_dir(0)
# -----------------------------------------------------------------------------
# get_output_dir — test-mode redirection
# -----------------------------------------------------------------------------
def test_get_output_dir_production_path(tmp_path, monkeypatch) -> None:
"""No env var → writes under the chapter dir."""
monkeypatch.delenv("ML4T_OUTPUT_DIR", raising=False)
monkeypatch.delenv("ML4T_CHAPTER_OUTPUT_DIR", raising=False)
# Use create=False to avoid making a directory in the real repo.
path = get_output_dir(7, "etfs", create=False)
assert path == get_chapter_dir(7) / "output" / "etfs"
def test_get_output_dir_redirects_under_ml4t_output_dir(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
monkeypatch.delenv("ML4T_CHAPTER_OUTPUT_DIR", raising=False)
path = get_output_dir(7, "etfs")
assert path == tmp_path / "ch07_etfs"
assert path.exists()
def test_get_output_dir_chapter_specific_env_wins_over_global(tmp_path, monkeypatch) -> None:
"""ML4T_CHAPTER_OUTPUT_DIR should take precedence over ML4T_OUTPUT_DIR."""
ch_dir = tmp_path / "chapter-only"
global_dir = tmp_path / "global"
monkeypatch.setenv("ML4T_CHAPTER_OUTPUT_DIR", str(ch_dir))
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(global_dir))
path = get_output_dir(11, "etfs")
assert path.parent == ch_dir
assert not str(path).startswith(str(global_dir))
def test_get_output_dir_zero_pads_chapter_number(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
monkeypatch.delenv("ML4T_CHAPTER_OUTPUT_DIR", raising=False)
assert get_output_dir(3, "x").name == "ch03_x"
assert get_output_dir(11, "x").name == "ch11_x"
# -----------------------------------------------------------------------------
# get_case_study_dir — test-mode redirection
# -----------------------------------------------------------------------------
def test_get_case_study_dir_production_path(monkeypatch) -> None:
"""No env var → writes under case_studies/{id}/."""
monkeypatch.delenv("ML4T_OUTPUT_DIR", raising=False)
path = get_case_study_dir("etfs", create=False)
assert path == REPO_ROOT / "case_studies" / "etfs"
def test_get_case_study_dir_redirects_under_ml4t_output_dir(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
path = get_case_study_dir("etfs")
assert path == tmp_path / "etfs"
assert path.exists()
def test_get_case_study_dir_create_false_does_not_mkdir(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
path = get_case_study_dir("etfs", create=False)
assert not path.exists()
def test_get_case_study_dir_create_true_is_idempotent(tmp_path, monkeypatch) -> None:
monkeypatch.setenv("ML4T_OUTPUT_DIR", str(tmp_path))
first = get_case_study_dir("etfs", create=True)
second = get_case_study_dir("etfs", create=True)
assert first == second
assert first.exists()
tests/test_pm_helpers.py
+98
View File
@@ -0,0 +1,98 @@
from pathlib import Path
import pytest
from tests.pm_helpers import (
RECORD_REPLAY,
RECORD_REWRITE,
TIER_ON_DEMAND,
TIER_PER_COMMIT,
TIER_WEEKLY,
collect_chapter_notebooks,
current_test_tier,
get_record_mode,
get_reruns,
get_tier,
)
def test_collect_chapter_notebooks_keeps_real_notebooks_and_skips_helpers() -> None:
notebooks = {path.as_posix() for path in collect_chapter_notebooks(Path("."), range(1, 28))}
assert "06_strategy_definition/02_case_study_overview.py" in notebooks
assert "08_financial_features/case_study_feature_summary.py" in notebooks
assert "11_ml_pipeline/08_ml_backtest_intro.py" in notebooks
assert "16_strategy_simulation/01_backtest_first_principles.py" in notebooks
assert "21_rl_execution_hedging/07_backtest_with_impact.py" in notebooks
assert "03_market_microstructure/filter_itch_symbol.py" not in notebooks
assert "03_market_microstructure/lob_utils.py" not in notebooks
assert "03_market_microstructure/lob_generator.py" not in notebooks
assert "13_dl_time_series/dl_utils.py" not in notebooks
# ---------------------------------------------------------------------------
# Tier / reruns / record_mode helpers
# ---------------------------------------------------------------------------
def test_get_tier_defaults_to_per_commit() -> None:
assert get_tier({}) == TIER_PER_COMMIT
assert get_tier({"tier": None}) == TIER_PER_COMMIT
def test_get_tier_accepts_valid_values() -> None:
assert get_tier({"tier": "per_commit"}) == TIER_PER_COMMIT
assert get_tier({"tier": "weekly"}) == TIER_WEEKLY
assert get_tier({"tier": "on_demand"}) == TIER_ON_DEMAND
def test_get_tier_rejects_invalid() -> None:
with pytest.raises(ValueError, match="Invalid tier"):
get_tier({"tier": "nightly"})
def test_current_test_tier_defaults_to_per_commit(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("ML4T_TEST_TIER", raising=False)
assert current_test_tier() == TIER_PER_COMMIT
def test_current_test_tier_reads_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("ML4T_TEST_TIER", "weekly")
assert current_test_tier() == TIER_WEEKLY
monkeypatch.setenv("ML4T_TEST_TIER", "on_demand")
assert current_test_tier() == TIER_ON_DEMAND
def test_current_test_tier_rejects_invalid(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("ML4T_TEST_TIER", "bogus")
with pytest.raises(ValueError, match="Invalid ML4T_TEST_TIER"):
current_test_tier()
def test_get_reruns_default_zero() -> None:
assert get_reruns({}) == 0
def test_get_reruns_returns_int() -> None:
assert get_reruns({"reruns": 3}) == 3
def test_get_reruns_rejects_negative_or_nonint() -> None:
with pytest.raises(ValueError):
get_reruns({"reruns": -1})
with pytest.raises(ValueError):
get_reruns({"reruns": "2"})
def test_get_record_mode_defaults_to_replay() -> None:
assert get_record_mode({}) == RECORD_REPLAY
def test_get_record_mode_accepts_rewrite() -> None:
assert get_record_mode({"record_mode": "rewrite"}) == RECORD_REWRITE
def test_get_record_mode_rejects_invalid() -> None:
with pytest.raises(ValueError, match="Invalid record_mode"):
get_record_mode({"record_mode": "none"})
tests/test_registry_completeness.py
+331
View File
@@ -0,0 +1,331 @@
"""Tests for case_studies/utils/registry/completeness.py.
The skip-if-complete invariant is load-bearing: wrong answers either
waste hours of compute (should have skipped) or silently reuse stale
partial artifacts (should have retrained). Tests:
- missing training_run → exists=False, not complete
- present training_run but no prediction_sets → partial, missing list
- complete run → exists=True, complete=True
- partial backtest_run (no daily_returns.parquet) → partial
- require_metrics=False relaxes the completeness rule as advertised
- skip_* wrappers return the same status (behavior pin for callers)
"""
from __future__ import annotations
import sqlite3
import time
from pathlib import Path
import pytest
from case_studies.utils.registry.completeness import (
BacktestRunStatus,
TrainingRunStatus,
backtest_run_status,
skip_backtest_if_complete,
skip_training_if_complete,
training_run_status,
)
from case_studies.utils.registry.specs import (
backtest_hash_from_parts,
training_hash_from_spec,
)
from case_studies.utils.registry.store import (
REGISTRY_SCHEMA_SQL,
_backtest_dir,
_prediction_dir,
_registry_db_path,
)
@pytest.fixture
def case_dir(tmp_path) -> Path:
"""Create a minimal case study dir with an empty registry.db."""
case = tmp_path / "etfs"
case.mkdir()
db_path = _registry_db_path(case)
db_path.parent.mkdir(parents=True, exist_ok=True)
db = sqlite3.connect(str(db_path))
db.executescript(REGISTRY_SCHEMA_SQL)
db.commit()
db.close()
return case
@pytest.fixture
def canonical_spec() -> dict:
return {"family": "linear", "label": "fwd_ret_21d", "seed": 42, "n_folds": 5}
def _insert_training_run(case_dir: Path, spec: dict) -> str:
"""Insert a training_runs row. Returns the training_hash."""
t_hash = training_hash_from_spec(spec)
db = sqlite3.connect(str(_registry_db_path(case_dir)))
now = time.strftime("%Y-%m-%dT%H:%M:%S")
db.execute(
"INSERT INTO training_runs (training_hash, family, label, config_name, spec_json, created_at) "
"VALUES (?, ?, ?, ?, ?, ?)",
(t_hash, spec["family"], spec["label"], "test", "{}", now),
)
db.commit()
db.close()
return t_hash
def _insert_prediction_set(case_dir: Path, t_hash: str, split: str = "val") -> str:
"""Insert a prediction_sets row. Returns the prediction_hash."""
from case_studies.utils.registry.specs import prediction_hash_from_parts
p_hash = prediction_hash_from_parts(t_hash, None, split)
db = sqlite3.connect(str(_registry_db_path(case_dir)))
now = time.strftime("%Y-%m-%dT%H:%M:%S")
db.execute(
"INSERT INTO prediction_sets (prediction_hash, training_hash, split, created_at) "
"VALUES (?, ?, ?, ?)",
(p_hash, t_hash, split, now),
)
db.commit()
db.close()
return p_hash
def _insert_prediction_metric(case_dir: Path, p_hash: str, ic_mean: float = 0.01) -> None:
db = sqlite3.connect(str(_registry_db_path(case_dir)))
now = time.strftime("%Y-%m-%dT%H:%M:%S")
db.execute(
"INSERT INTO prediction_metrics (prediction_hash, computed_at, ic_mean) VALUES (?, ?, ?)",
(p_hash, now, ic_mean),
)
db.commit()
db.close()
def _touch_predictions_file(case_dir: Path, p_hash: str) -> None:
d = _prediction_dir(case_dir, p_hash)
d.mkdir(parents=True, exist_ok=True)
(d / "predictions.parquet").write_bytes(b"fake")
# -----------------------------------------------------------------------------
# training_run_status
# -----------------------------------------------------------------------------
def test_training_status_missing_run_is_not_complete(case_dir, canonical_spec) -> None:
status = training_run_status("etfs", canonical_spec, case_dir=case_dir)
assert not status.exists
assert not status.complete
assert not status.partial # Neither complete nor partial when nothing exists
assert "training_run" in status.missing
assert status.training_hash == training_hash_from_spec(canonical_spec)
def test_training_status_run_without_predictions_is_partial(case_dir, canonical_spec) -> None:
_insert_training_run(case_dir, canonical_spec)
status = training_run_status("etfs", canonical_spec, case_dir=case_dir)
assert status.exists
assert status.partial
assert not status.complete
assert "prediction_sets" in status.missing
def test_training_status_run_without_metrics_is_partial(case_dir, canonical_spec) -> None:
t_hash = _insert_training_run(case_dir, canonical_spec)
p_hash = _insert_prediction_set(case_dir, t_hash)
_touch_predictions_file(case_dir, p_hash)
# Deliberately skip metric insertion
status = training_run_status("etfs", canonical_spec, case_dir=case_dir)
assert status.partial
assert not status.complete
assert "ic_mean" in status.missing
def test_training_status_complete_when_all_artifacts_present(case_dir, canonical_spec) -> None:
t_hash = _insert_training_run(case_dir, canonical_spec)
p_hash = _insert_prediction_set(case_dir, t_hash)
_insert_prediction_metric(case_dir, p_hash)
_touch_predictions_file(case_dir, p_hash)
status = training_run_status("etfs", canonical_spec, case_dir=case_dir)
assert status.complete
assert not status.partial
assert status.missing == ()
def test_training_status_require_metrics_false_relaxes_completeness(
case_dir, canonical_spec
) -> None:
"""With require_metrics=False, a run is complete even without ic_mean."""
t_hash = _insert_training_run(case_dir, canonical_spec)
p_hash = _insert_prediction_set(case_dir, t_hash)
_touch_predictions_file(case_dir, p_hash)
# No metric insertion
status = training_run_status("etfs", canonical_spec, case_dir=case_dir, require_metrics=False)
assert status.complete
assert "ic_mean" not in status.missing
def test_training_status_require_predictions_file_false_relaxes(case_dir, canonical_spec) -> None:
t_hash = _insert_training_run(case_dir, canonical_spec)
p_hash = _insert_prediction_set(case_dir, t_hash)
_insert_prediction_metric(case_dir, p_hash)
# Deliberately skip writing predictions.parquet
status = training_run_status(
"etfs", canonical_spec, case_dir=case_dir, require_predictions_file=False
)
assert status.complete
def test_training_status_summary_formats() -> None:
missing = TrainingRunStatus(
training_hash="abcdef1234567890",
exists=False,
has_predictions=False,
has_predictions_file=False,
has_metrics=False,
missing=("training_run",),
)
assert "no training_run" in missing.summary()
assert "abcdef123456" in missing.summary()
partial = TrainingRunStatus(
training_hash="abcdef1234567890",
exists=True,
has_predictions=True,
has_predictions_file=False,
has_metrics=False,
missing=("predictions.parquet", "ic_mean"),
)
assert "partial" in partial.summary()
assert "predictions.parquet" in partial.summary()
complete = TrainingRunStatus(
training_hash="abcdef1234567890",
exists=True,
has_predictions=True,
has_predictions_file=True,
has_metrics=True,
)
assert "complete" in complete.summary()
# -----------------------------------------------------------------------------
# skip_training_if_complete (thin wrapper)
# -----------------------------------------------------------------------------
def test_skip_training_returns_same_status_as_direct_call(case_dir, canonical_spec) -> None:
direct = training_run_status("etfs", canonical_spec, case_dir=case_dir)
wrapped = skip_training_if_complete("etfs", canonical_spec, case_dir=case_dir, verbose=False)
assert direct.training_hash == wrapped.training_hash
assert direct.complete == wrapped.complete
# -----------------------------------------------------------------------------
# backtest_run_status
# -----------------------------------------------------------------------------
def test_backtest_status_missing_is_not_complete(case_dir) -> None:
strategy = {"signal": {"method": "equal_weight_top_k", "top_k": 10}}
status = backtest_run_status("etfs", "pred123", strategy, case_dir=case_dir)
assert not status.exists
assert not status.complete
assert "backtest_run" in status.missing
def test_backtest_status_partial_when_returns_missing(case_dir) -> None:
strategy = {"signal": {"method": "equal_weight_top_k", "top_k": 10}}
b_hash = backtest_hash_from_parts("pred123", strategy)
# Insert backtest_runs row but no returns file
db = sqlite3.connect(str(_registry_db_path(case_dir)))
# Need to satisfy FK: insert a synthetic prediction first.
# The schema is ON CASCADE default, but FK references must exist.
db.execute("PRAGMA foreign_keys=OFF") # Tests: skip FK check to simplify fixture
now = time.strftime("%Y-%m-%dT%H:%M:%S")
db.execute(
"INSERT INTO backtest_runs (backtest_hash, prediction_hash, spec_json, created_at) "
"VALUES (?, ?, ?, ?)",
(b_hash, "pred123", "{}", now),
)
db.execute(
"INSERT INTO backtest_metrics (backtest_hash, computed_at, sharpe) VALUES (?, ?, ?)",
(b_hash, now, 1.5),
)
db.commit()
db.close()
status = backtest_run_status("etfs", "pred123", strategy, case_dir=case_dir)
assert status.exists
assert status.partial
assert not status.complete
assert "daily_returns.parquet" in status.missing
def test_backtest_status_complete_when_all_present(case_dir) -> None:
strategy = {"signal": {"method": "equal_weight_top_k", "top_k": 10}}
b_hash = backtest_hash_from_parts("pred123", strategy)
db = sqlite3.connect(str(_registry_db_path(case_dir)))
db.execute("PRAGMA foreign_keys=OFF")
now = time.strftime("%Y-%m-%dT%H:%M:%S")
db.execute(
"INSERT INTO backtest_runs (backtest_hash, prediction_hash, spec_json, created_at) "
"VALUES (?, ?, ?, ?)",
(b_hash, "pred123", "{}", now),
)
db.execute(
"INSERT INTO backtest_metrics (backtest_hash, computed_at, sharpe) VALUES (?, ?, ?)",
(b_hash, now, 1.5),
)
db.commit()
db.close()
d = _backtest_dir(case_dir, b_hash)
d.mkdir(parents=True, exist_ok=True)
(d / "daily_returns.parquet").write_bytes(b"fake")
status = backtest_run_status("etfs", "pred123", strategy, case_dir=case_dir)
assert status.complete
def test_backtest_status_hash_is_deterministic(case_dir) -> None:
strategy = {"signal": {"method": "x", "top_k": 10}, "allocation": {"method": "eq"}}
s1 = backtest_run_status("etfs", "pred123", strategy, case_dir=case_dir)
s2 = backtest_run_status("etfs", "pred123", dict(strategy), case_dir=case_dir)
assert s1.backtest_hash == s2.backtest_hash
def test_backtest_status_summary_formats() -> None:
missing = BacktestRunStatus(
backtest_hash="abc123def456000",
exists=False,
has_returns=False,
has_metrics=False,
missing=("backtest_run",),
)
assert "no backtest_run" in missing.summary()
# -----------------------------------------------------------------------------
# skip_backtest_if_complete (thin wrapper)
# -----------------------------------------------------------------------------
def test_skip_backtest_wraps_backtest_run_status(case_dir) -> None:
strategy = {"signal": {"method": "x"}}
direct = backtest_run_status("etfs", "pred123", strategy, case_dir=case_dir)
wrapped = skip_backtest_if_complete(
"etfs", "pred123", strategy, case_dir=case_dir, verbose=False
)
assert direct.backtest_hash == wrapped.backtest_hash
assert direct.complete == wrapped.complete
tests/test_registry_metrics.py
+308
View File
@@ -0,0 +1,308 @@
"""Tests for case_studies/utils/registry/metrics.py — prediction metric aggregation.
The critical contract this pins is the *classification IC rule*: when
``task_type='classification'``, IC is computed against the continuous
return column named by ``eval_col``, never against the binary label.
Computing IC against the binary label degenerates to ``2·(AUC − 0.5)``
and is not a valid Spearman rank correlation against returns — the April
classification IC backfill was needed precisely because this was wrong.
These tests lock in:
- Regression path: IC is cross-sectional rank correlation of y_score
vs y_true (continuous); RMSE / MAE are computed on valid pairs.
- Classification path: IC uses ``eval_col`` (continuous return), and
AUC / log_loss / accuracy use the binary ``y_true_col``.
- Classification IC on y_score + y_ret equals the regression IC on the
same y_score + y_ret — i.e., the classification branch does not
silently fall back to using the binary label for IC.
- Missing ``eval_col`` (or a column that isn't on the DataFrame) raises
loudly rather than silently collapsing to AUC-disguised-as-IC.
- Headline aggregation: ``ic_mean`` = mean across folds, ``ic_t`` =
Newey-West-free pooled t, ``pct_positive`` = fraction of folds with
IC > 0, ``n_folds`` = count, ``task_type`` = 'classification' for classification.
All fixtures are hermetic — no real data, no setup.yaml.
"""
from __future__ import annotations
import math
import numpy as np
import polars as pl
import pytest
from case_studies.utils.registry.metrics import compute_prediction_fold_metrics
# -----------------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------------
@pytest.fixture(scope="module")
def regression_predictions() -> pl.DataFrame:
"""2 folds × 10 dates × 10 entities with y_score ≈ y_true (high IC)."""
rng = np.random.default_rng(0)
rows = []
for fold in (0, 1):
for d in range(10):
for e in range(10):
y_true = float(rng.normal())
y_score = 0.8 * y_true + 0.2 * float(rng.normal())
rows.append(
{
"timestamp": f"2024-{fold + 1:02d}-{d + 1:02d}",
"symbol": f"S{e}",
"fold_id": fold,
"y_true": y_true,
"y_score": y_score,
}
)
return pl.DataFrame(rows).with_columns(pl.col("timestamp").str.to_date())
@pytest.fixture(scope="module")
def classification_predictions(regression_predictions) -> pl.DataFrame:
"""Classification variant: y_true is the sign of the continuous return.
- ``y_ret`` preserves the continuous return (the eval_col target)
- ``y_true`` is the binary label (1 if return > 0)
- ``y_score`` is a probability-style score from a logistic squash of
the original continuous score, so it is still monotone in the
continuous return
"""
return regression_predictions.rename({"y_score": "y_score_cont"}).with_columns(
y_ret=pl.col("y_true"),
y_true=pl.when(pl.col("y_true") > 0).then(1).otherwise(0).cast(pl.Int8),
y_score=1.0 / (1.0 + (-pl.col("y_score_cont")).exp()),
)
# -----------------------------------------------------------------------------
# Regression path
# -----------------------------------------------------------------------------
def test_regression_computes_rmse_mae_and_ic(regression_predictions) -> None:
headline, folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
assert set(folds.keys()) == {0, 1}
for fm in folds.values():
assert "rmse" in fm and "mae" in fm
assert "ic" in fm
# y_score ≈ 0.8 * y_true so IC should be high
assert fm["ic"] > 0.5
# RMSE / MAE on standard normals with 0.2σ noise should be tiny-ish
assert fm["rmse"] >= 0
assert fm["mae"] >= 0
def test_regression_headline_task_type_is_regression(regression_predictions) -> None:
headline, _ = compute_prediction_fold_metrics(regression_predictions, task_type="regression")
assert headline["task_type"] == "regression"
def test_regression_headline_ic_mean_equals_fold_ic_mean(regression_predictions) -> None:
headline, folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
expected = float(np.mean([fm["ic"] for fm in folds.values()]))
assert math.isclose(headline["ic_mean"], expected, rel_tol=1e-12)
def test_regression_pct_positive_matches_fraction_of_positive_ic_folds(
regression_predictions,
) -> None:
headline, folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
expected = float(np.mean([fm["ic"] > 0 for fm in folds.values()]))
assert headline["pct_positive"] == expected
def test_regression_headline_ic_t_is_mean_over_stderr(regression_predictions) -> None:
headline, folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
fold_ics = np.array([fm["ic"] for fm in folds.values()])
expected_t = float(np.mean(fold_ics) / (np.std(fold_ics) / np.sqrt(len(fold_ics))))
assert math.isclose(headline["ic_t"], expected_t, rel_tol=1e-12)
def test_regression_n_folds_matches_unique_fold_ids(regression_predictions) -> None:
headline, _ = compute_prediction_fold_metrics(regression_predictions, task_type="regression")
assert headline["n_folds"] == 2
# -----------------------------------------------------------------------------
# Classification path — the load-bearing contract
# -----------------------------------------------------------------------------
def test_classification_without_eval_col_raises_value_error(classification_predictions) -> None:
"""The defensive check that saved us from re-introducing the IC-on-binary bug."""
with pytest.raises(ValueError, match="eval_col"):
compute_prediction_fold_metrics(
classification_predictions, task_type="classification", class_values=[0, 1]
)
def test_classification_missing_eval_col_raises_key_error(classification_predictions) -> None:
with pytest.raises(KeyError, match="does_not_exist"):
compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
eval_col="does_not_exist",
class_values=[0, 1],
)
def test_classification_ic_is_computed_vs_continuous_return(
regression_predictions, classification_predictions
) -> None:
"""The classification IC on (y_score_cont, y_ret) must equal the regression
IC on the same (y_score, y_true) — proving classification did not silently
fall back to IC-vs-binary.
We compare against a regression run on the ORIGINAL continuous pair, to
establish what the IC should be. Then we compare the classification IC
on (y_score_cont, y_ret) against that reference. Classification's IC
uses the CONTINUOUS score column via ``y_score_col`` override.
"""
# Reference: regression on the continuous ground truth
ref_headline, _ = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
# Classification run — pass the continuous score column as ``y_score_col``
# and point ``eval_col`` at the continuous return. That pairing should
# reproduce the regression IC exactly.
cls_headline, _ = compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret",
class_values=[0, 1],
)
assert math.isclose(cls_headline["ic_mean"], ref_headline["ic_mean"], rel_tol=1e-12)
assert math.isclose(cls_headline["ic_std"], ref_headline["ic_std"], rel_tol=1e-12)
def test_classification_ic_differs_from_ic_on_binary_label(classification_predictions) -> None:
"""Sanity: IC on continuous return is materially different from what you'd
get if you wrongly computed IC on the binary label.
We simulate the wrong behavior by aliasing y_true as eval_col and check
that IC differs from the correct run.
"""
correct, _ = compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret",
class_values=[0, 1],
)
# Build a frame where eval_col points at the BINARY label (simulating the bug).
wrong_df = classification_predictions.with_columns(y_ret_bin=pl.col("y_true"))
wrong, _ = compute_prediction_fold_metrics(
wrong_df,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret_bin",
class_values=[0, 1],
)
# The two IC values MUST differ materially — if they matched, IC-on-binary
# would be indistinguishable from IC-on-continuous, defeating the rule.
assert abs(correct["ic_mean"] - wrong["ic_mean"]) > 0.05
def test_classification_adds_auc_accuracy_logloss_to_headline(
classification_predictions,
) -> None:
headline, _ = compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret",
class_values=[0, 1],
)
for m in ("auc_roc", "auc_pr", "log_loss", "brier_score", "accuracy", "balanced_accuracy"):
assert m in headline, f"missing classification metric {m!r} in headline"
def test_classification_headline_task_type_is_one(classification_predictions) -> None:
headline, _ = compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret",
class_values=[0, 1],
)
assert headline["task_type"] == "classification"
def test_classification_auc_is_computed_on_binary_label(classification_predictions) -> None:
"""AUC / accuracy / log_loss go against the binary y_true, not the continuous
eval_col. A classification score that is perfectly monotone in the binary
label should yield AUC=1.0 (well above the chance baseline of 0.5).
"""
headline, _ = compute_prediction_fold_metrics(
classification_predictions,
task_type="classification",
y_score_col="y_score_cont",
eval_col="y_ret",
class_values=[0, 1],
)
# Strongly monotone score ⇒ high AUC (not 0.5)
assert headline["auc_roc"] > 0.9
# -----------------------------------------------------------------------------
# Edge cases
# -----------------------------------------------------------------------------
def test_single_fold_returns_zero_ic_std_and_zero_t(regression_predictions) -> None:
"""With only one fold, cross-fold stddev is undefined; the function reports 0."""
fold0_only = regression_predictions.filter(pl.col("fold_id") == 0)
headline, _ = compute_prediction_fold_metrics(fold0_only, task_type="regression")
assert headline["n_folds"] == 1
assert headline["ic_std"] == 0.0
assert headline["ic_t"] == 0.0
def test_accepts_pandas_dataframe(regression_predictions) -> None:
"""pd.DataFrame input is converted to polars internally."""
import pandas as pd
pdf = regression_predictions.to_pandas()
assert isinstance(pdf, pd.DataFrame)
headline, folds = compute_prediction_fold_metrics(pdf, task_type="regression")
assert set(folds.keys()) == {0, 1}
assert "ic_mean" in headline
def test_deterministic_across_calls(regression_predictions) -> None:
"""Repeated calls produce numerically equivalent output.
BLAS threading can introduce <1e-13 jitter in rank-correlation summations,
so we use approximate equality rather than bit-exact.
"""
a_head, a_folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
b_head, b_folds = compute_prediction_fold_metrics(
regression_predictions, task_type="regression"
)
for key in a_head:
assert a_head[key] == pytest.approx(b_head[key], abs=1e-10), key
for fold_id in a_folds:
for key in a_folds[fold_id]:
assert a_folds[fold_id][key] == pytest.approx(b_folds[fold_id][key], abs=1e-10), (
f"fold {fold_id} / {key}"
)
tests/test_registry_queries_cohort.py
+194
View File
@@ -0,0 +1,194 @@
"""Tests for case_studies/utils/registry/queries.py cohort_metrics overrides.
Pins three behaviors of ``load_backtest_metrics``:
1. ER values from ``cohort_metrics`` override the legacy ``dsr`` /
``dsr_pvalue`` / ``expected_max_sharpe`` / ``min_trl_periods`` columns
for rows that are the family leader; non-leaders pass through with
NULLs (legacy backtest_metrics columns were dropped post-Phase-H).
2. The pre-migration fallback — when ``cohort_metrics`` doesn't exist on
the registry — returns raw ``backtest_metrics`` rows with null
placeholder columns for the override columns. The fallback is keyed
on a ``sqlite_master`` probe (``_table_exists``) so the narrow case
stays narrow.
3. Duplicate-leader_hash defense — if two family cohorts somehow share
a leader_hash, the join keeps the first row and emits a warning
instead of fanning out and silently changing row cardinality.
"""
from __future__ import annotations
import sqlite3
from pathlib import Path
import polars as pl
import pytest
def _bootstrap_registry(db_path: Path, *, with_cohort_metrics: bool = True) -> None:
db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(db_path))
try:
# Minimal backtest_metrics schema — only the columns the override
# touches plus backtest_hash + a passthrough column.
conn.execute(
"""
CREATE TABLE backtest_metrics (
backtest_hash TEXT PRIMARY KEY,
sharpe REAL
)
"""
)
conn.executemany(
"INSERT INTO backtest_metrics(backtest_hash, sharpe) VALUES (?, ?)",
[("hash_leader", 1.2), ("hash_other", 0.6)],
)
if with_cohort_metrics:
conn.execute(
"""
CREATE TABLE cohort_metrics (
cohort_type TEXT NOT NULL,
stage TEXT,
label TEXT NOT NULL,
family TEXT,
leader_hash TEXT NOT NULL,
k_variants INTEGER,
dsr_er REAL,
dsr_er_pvalue REAL,
expected_max_sharpe_er REAL,
min_trl_periods_er REAL,
leader_min_trl REAL,
pbo REAL,
pbo_n_combinations REAL,
pbo_median_oos_rank REAL,
pbo_mean_degradation REAL,
pbo_n_folds REAL,
reality_check_pvalue REAL,
reality_check_statistic REAL,
reality_check_k REAL
)
"""
)
conn.commit()
finally:
conn.close()
def _seed_cohort_row(
db_path: Path,
*,
leader_hash: str,
stage: str = "signal",
label: str = "fwd_ret_21d",
family: str = "linear",
dsr_er: float = 0.85,
dsr_er_pvalue: float = 0.02,
) -> None:
conn = sqlite3.connect(str(db_path))
try:
conn.execute(
"""
INSERT INTO cohort_metrics(
cohort_type, stage, label, family, leader_hash, k_variants,
dsr_er, dsr_er_pvalue, expected_max_sharpe_er, min_trl_periods_er,
leader_min_trl, pbo, pbo_n_combinations, pbo_median_oos_rank,
pbo_mean_degradation, pbo_n_folds, reality_check_pvalue,
reality_check_statistic, reality_check_k
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
"family",
stage,
label,
family,
leader_hash,
30,
dsr_er,
dsr_er_pvalue,
1.5,
40.0,
40.0,
0.18,
20.0,
3.5,
0.1,
6.0,
0.04,
0.7,
30.0,
),
)
conn.commit()
finally:
conn.close()
@pytest.fixture
def registry_with_cohort(tmp_path: Path) -> Path:
case_dir = tmp_path / "case_x"
db_path = case_dir / "run_log" / "registry.db"
_bootstrap_registry(db_path, with_cohort_metrics=True)
_seed_cohort_row(db_path, leader_hash="hash_leader")
return case_dir
@pytest.fixture
def registry_without_cohort(tmp_path: Path) -> Path:
case_dir = tmp_path / "case_y"
db_path = case_dir / "run_log" / "registry.db"
_bootstrap_registry(db_path, with_cohort_metrics=False)
return case_dir
def test_cohort_override_applies_to_leader_and_skips_non_leader(
registry_with_cohort: Path,
) -> None:
from case_studies.utils.registry.queries import load_backtest_metrics
df = load_backtest_metrics("case_x", case_dir=registry_with_cohort)
assert df.height == 2 # row count preserved
leader = df.filter(pl.col("backtest_hash") == "hash_leader")
other = df.filter(pl.col("backtest_hash") == "hash_other")
assert leader["dsr"].item() == pytest.approx(0.85)
assert leader["dsr_pvalue"].item() == pytest.approx(0.02)
assert leader["k_variants"].item() == 30
assert leader["pbo"].item() == pytest.approx(0.18)
# Non-leader carries NULLs for the override columns.
assert other["dsr"].item() is None
assert other["pbo"].item() is None
def test_missing_cohort_table_returns_raw_rows_with_null_overrides(
registry_without_cohort: Path,
) -> None:
from case_studies.utils.registry.queries import load_backtest_metrics
df = load_backtest_metrics("case_y", case_dir=registry_without_cohort)
assert df.height == 2
assert {"dsr", "pbo", "reality_check_pvalue"}.issubset(df.columns)
# All override columns null when cohort_metrics doesn't exist.
assert df["dsr"].is_null().all()
assert df["pbo"].is_null().all()
def test_duplicate_leader_hash_keeps_row_cardinality_and_warns(
tmp_path: Path, caplog: pytest.LogCaptureFixture
) -> None:
"""Two family cohorts with the same leader_hash must not fan out."""
import logging
from case_studies.utils.registry.queries import load_backtest_metrics
case_dir = tmp_path / "case_z"
db_path = case_dir / "run_log" / "registry.db"
_bootstrap_registry(db_path, with_cohort_metrics=True)
_seed_cohort_row(db_path, leader_hash="hash_leader", family="linear", dsr_er=0.85)
_seed_cohort_row(db_path, leader_hash="hash_leader", family="gbm", dsr_er=0.99)
with caplog.at_level(logging.WARNING, logger="case_studies.utils.registry.queries"):
df = load_backtest_metrics("case_z", case_dir=case_dir)
# Row count preserved; the join didn't fan out.
assert df.height == 2
assert any("duplicate family-cohort leader_hash" in r.message for r in caplog.records)
tests/test_registry_specs.py
+227
View File
@@ -0,0 +1,227 @@
"""Tests for case_studies/utils/registry/specs.py.
Hashing determinism is load-bearing: the registry is a content-addressed
store, so any perturbation to hash computation (key ordering, separator
choice, seed handling) silently duplicates runs and corrupts lineage.
These tests pin the exact byte-for-byte hash output so a reformat of
canonical_json or compute_hash cannot change the addresses of existing runs.
"""
from __future__ import annotations
import hashlib
import json
import pytest
from case_studies.utils.registry.specs import (
DEFAULT_SEED,
HASH_LENGTH,
_validate_spec,
backtest_hash_from_parts,
canonical_json,
compute_hash,
prediction_hash_from_parts,
training_hash_from_spec,
)
# -----------------------------------------------------------------------------
# canonical_json — deterministic serialization
# -----------------------------------------------------------------------------
def test_canonical_json_sorts_keys() -> None:
a = canonical_json({"b": 2, "a": 1})
b = canonical_json({"a": 1, "b": 2})
assert a == b
assert a == '{"a":1,"b":2}'
def test_canonical_json_uses_compact_separators() -> None:
assert canonical_json({"x": 1}) == '{"x":1}'
def test_canonical_json_stringifies_unserializable_via_default() -> None:
"""Path/Enum/datetime-like fields fall through `default=str` so a spec
never fails serialization."""
from pathlib import Path
out = canonical_json({"path": Path("/tmp/x.parquet")})
assert "/tmp/x.parquet" in out
def test_canonical_json_is_deterministic_across_nested_structures() -> None:
spec = {
"outer": {"b": [3, 2, 1], "a": {"z": 9, "y": 8}},
"flat": 42,
}
first = canonical_json(spec)
second = canonical_json(spec)
assert first == second
# Keys are sorted at every level
assert first.index('"a":') < first.index('"b":')
assert first.index('"y":') < first.index('"z":')
# -----------------------------------------------------------------------------
# compute_hash — sha256 truncation invariant
# -----------------------------------------------------------------------------
def test_compute_hash_default_length_is_12() -> None:
h = compute_hash("anything")
assert len(h) == HASH_LENGTH == 12
def test_compute_hash_is_prefix_of_full_sha256() -> None:
content = "some_training_content"
expected_prefix = hashlib.sha256(content.encode()).hexdigest()[:12]
assert compute_hash(content) == expected_prefix
def test_compute_hash_length_override_respects_arg() -> None:
assert len(compute_hash("x", length=6)) == 6
assert len(compute_hash("x", length=64)) == 64
# -----------------------------------------------------------------------------
# training_hash_from_spec
# -----------------------------------------------------------------------------
def _base_spec(**overrides) -> dict:
spec = {
"family": "linear",
"label": "fwd_ret_21d",
"seed": 42,
"n_folds": 5,
}
spec.update(overrides)
return spec
def test_training_hash_is_deterministic() -> None:
spec = _base_spec()
assert training_hash_from_spec(spec) == training_hash_from_spec(dict(spec))
def test_training_hash_differs_when_seed_changes() -> None:
assert training_hash_from_spec(_base_spec(seed=1)) != training_hash_from_spec(
_base_spec(seed=2)
)
def test_training_hash_differs_when_family_changes() -> None:
assert training_hash_from_spec(_base_spec(family="gbm")) != training_hash_from_spec(
_base_spec(family="linear")
)
def test_training_hash_differs_when_label_changes() -> None:
assert training_hash_from_spec(_base_spec(label="fwd_ret_5d")) != training_hash_from_spec(
_base_spec(label="fwd_ret_21d")
)
def test_training_hash_invariant_under_key_order() -> None:
"""Client code may build the spec dict in arbitrary order; hash must be stable."""
spec_a = {"family": "gbm", "label": "fwd_ret_21d", "seed": 42}
spec_b = {"seed": 42, "label": "fwd_ret_21d", "family": "gbm"}
assert training_hash_from_spec(spec_a) == training_hash_from_spec(spec_b)
# -----------------------------------------------------------------------------
# Spec validation
# -----------------------------------------------------------------------------
def test_validate_spec_missing_seed_injects_default(caplog) -> None:
"""Missing-seed-only case: warn and inject DEFAULT_SEED."""
with caplog.at_level("WARNING"):
enriched = _validate_spec({"family": "gbm", "label": "fwd_ret_5d"})
assert enriched["seed"] == DEFAULT_SEED
assert "missing 'seed'" in caplog.text
def test_validate_spec_missing_multiple_fields_raises() -> None:
"""Anything beyond a missing seed is a hard error."""
with pytest.raises(ValueError, match="missing required fields"):
_validate_spec({"family": "gbm"}) # missing label + seed
def test_validate_spec_does_not_mutate_original() -> None:
original = {"family": "gbm", "label": "x"}
enriched = _validate_spec(original)
assert "seed" not in original # original untouched
assert enriched["seed"] == DEFAULT_SEED
# -----------------------------------------------------------------------------
# prediction_hash_from_parts
# -----------------------------------------------------------------------------
def test_prediction_hash_combines_training_hash_checkpoint_split() -> None:
h = prediction_hash_from_parts("abc123", 100, "val")
# Reconstruct exact content and compare to the public API
assert h == compute_hash("abc123|100|val")
def test_prediction_hash_none_checkpoint_becomes_final() -> None:
h_none = prediction_hash_from_parts("abc", None, "val")
h_final_str = prediction_hash_from_parts("abc", None, "val") # Same call
assert h_none == compute_hash("abc|final|val")
assert h_none == h_final_str
def test_prediction_hash_distinct_on_split() -> None:
assert prediction_hash_from_parts("abc", 1, "val") != prediction_hash_from_parts(
"abc", 1, "test"
)
def test_prediction_hash_distinct_on_checkpoint() -> None:
assert prediction_hash_from_parts("abc", 10, "val") != prediction_hash_from_parts(
"abc", 20, "val"
)
# -----------------------------------------------------------------------------
# backtest_hash_from_parts
# -----------------------------------------------------------------------------
def test_backtest_hash_combines_prediction_hash_and_strategy_spec() -> None:
strategy = {"signal": {"method": "equal_weight_top_k", "top_k": 10}}
h = backtest_hash_from_parts("pred123", strategy)
assert h == compute_hash(f"pred123|{canonical_json(strategy)}")
def test_backtest_hash_sensitive_to_strategy_change() -> None:
base = {"top_k": 10}
variant = {"top_k": 20}
assert backtest_hash_from_parts("p1", base) != backtest_hash_from_parts("p1", variant)
def test_backtest_hash_invariant_under_strategy_key_order() -> None:
a = {"signal": {"method": "x", "top_k": 10}, "allocation": {"method": "eq"}}
b = {"allocation": {"method": "eq"}, "signal": {"top_k": 10, "method": "x"}}
assert backtest_hash_from_parts("p", a) == backtest_hash_from_parts("p", b)
# -----------------------------------------------------------------------------
# Regression pin — the exact hash for a canonical spec
# -----------------------------------------------------------------------------
def test_training_hash_regression_pin_for_canonical_spec() -> None:
"""Pin the exact hash of a minimal valid spec. Changing this value
invalidates every existing registry entry — so any change should be an
explicit migration, not an accidental refactor."""
spec = {"family": "linear", "label": "fwd_ret_21d", "seed": 42}
content = json.dumps(spec, sort_keys=True, separators=(",", ":"), default=str)
expected = hashlib.sha256(content.encode()).hexdigest()[:12]
assert training_hash_from_spec(spec) == expected
tests/test_sequence_dataset.py
+202
View File
@@ -0,0 +1,202 @@
"""Correctness tests for case_studies/utils/sequence_dataset.py.
These tests encode the methodology property that every DL case study
depends on: the first validation sequence must predict the target at
val_start, using an input window that may extend back into train (this
is legal because features at times ≤ val_start are already known at
val_start; only labels after val_start are held out).
A test failure here means validation sequences have a warmup-drop bug
where the first `lookback` trading days of each val fold are silently
discarded — this inflates DL Sharpe on adversarial sample-period
exclusions and diverges from how the model would be deployed in
production.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
import pytest
def _synthetic_fold_df(
*,
n_symbols: int = 3,
train_start: str = "2020-01-01",
train_end: str = "2020-12-31",
val_start: str = "2021-01-01",
val_end: str = "2021-06-30",
freq: str = "B",
) -> tuple[pd.DataFrame, pd.Series, pd.Series, pd.Timestamp, pd.Timestamp]:
"""Build a synthetic panel: N symbols × business days train+val.
Returns (df, train_mask, val_mask, val_start_ts, val_end_ts).
"""
all_dates = pd.date_range(train_start, val_end, freq=freq)
rows = []
for i, sym in enumerate([f"S{j}" for j in range(n_symbols)]):
for dt in all_dates:
rows.append(
{
"symbol": sym,
"timestamp": dt,
"feat0": float(i) + dt.toordinal() / 1e6,
"feat1": float(i) * 2 + dt.toordinal() / 1e6,
"y": float(i) + np.sin(dt.toordinal() / 10.0),
}
)
df = pd.DataFrame(rows)
ts_start = pd.Timestamp(train_start)
ts_train_end = pd.Timestamp(train_end)
ts_val_start = pd.Timestamp(val_start)
ts_val_end = pd.Timestamp(val_end)
train_mask = (df["timestamp"] >= ts_start) & (df["timestamp"] <= ts_train_end)
val_mask = (df["timestamp"] >= ts_val_start) & (df["timestamp"] <= ts_val_end)
return df, train_mask, val_mask, ts_val_start, ts_val_end
def test_val_sequence_starts_at_val_start():
"""Every symbol's first val sequence should have target == val_start.
This is the core correctness property: in production, on val_start
we have all pre-val features available and must emit a prediction
for val_start. The prior (buggy) implementation discards the first
`lookback` rows of each val fold.
"""
from case_studies.utils.sequence_dataset import prepare_fold_sequence_stores
df, train_mask, val_mask, val_start_ts, _ = _synthetic_fold_df()
lookback = 20
_, val_store, fold_info = prepare_fold_sequence_stores(
df,
train_mask=train_mask,
val_mask=val_mask,
feature_names=["feat0", "feat1"],
label_col="y",
date_col="timestamp",
entity_col="symbol",
lookback=lookback,
val_start=val_start_ts,
)
assert fold_info["val_sequences"] > 0, "No val sequences generated"
# For each symbol, find the first sequence's target timestamp
for symbol_id in range(val_store.n_symbols):
end_positions = val_store.end_idx[val_store.symbol_idx == symbol_id]
if len(end_positions) == 0:
continue
first_end = end_positions.min()
first_target_ts = val_store.timestamps[symbol_id][first_end]
assert pd.Timestamp(first_target_ts) == val_start_ts, (
f"Symbol {val_store.entities[symbol_id]!r}: first val sequence "
f"predicts {first_target_ts}, expected {val_start_ts}. "
f"This indicates the warmup-drop bug — the first {lookback} "
f"trading days of val are being silently skipped."
)
def test_val_sequence_count_matches_val_calendar_days():
"""Number of val sequences per symbol == number of val-period rows."""
from case_studies.utils.sequence_dataset import prepare_fold_sequence_stores
df, train_mask, val_mask, val_start_ts, val_end_ts = _synthetic_fold_df()
lookback = 20
_, val_store, fold_info = prepare_fold_sequence_stores(
df,
train_mask=train_mask,
val_mask=val_mask,
feature_names=["feat0", "feat1"],
label_col="y",
date_col="timestamp",
entity_col="symbol",
lookback=lookback,
val_start=val_start_ts,
)
expected_per_symbol = int(
df[(df["timestamp"] >= val_start_ts) & (df["timestamp"] <= val_end_ts)]
.groupby("symbol")
.size()
.iloc[0]
)
actual_per_symbol = fold_info["val_sequences"] // val_store.n_symbols
assert actual_per_symbol == expected_per_symbol, (
f"Each symbol should have {expected_per_symbol} val sequences "
f"(one per val trading day); got {actual_per_symbol}. "
f"Shortfall indicates warmup drop."
)
def test_val_sequence_targets_never_include_train_period():
"""No val sequence should have a target timestamp < val_start.
Train-tail rows are used for priming input features only; their
labels must not appear as val targets (that would be leakage).
"""
from case_studies.utils.sequence_dataset import prepare_fold_sequence_stores
df, train_mask, val_mask, val_start_ts, _ = _synthetic_fold_df()
lookback = 20
_, val_store, _ = prepare_fold_sequence_stores(
df,
train_mask=train_mask,
val_mask=val_mask,
feature_names=["feat0", "feat1"],
label_col="y",
date_col="timestamp",
entity_col="symbol",
lookback=lookback,
val_start=val_start_ts,
)
for symbol_id in range(val_store.n_symbols):
end_positions = val_store.end_idx[val_store.symbol_idx == symbol_id]
for pos in end_positions:
target_ts = val_store.timestamps[symbol_id][pos]
assert pd.Timestamp(target_ts) >= val_start_ts, (
f"Val sequence target {target_ts} predates val_start "
f"{val_start_ts} — train-tail priming is leaking into "
f"predictions."
)
def test_backwards_compatible_without_val_start():
"""Omitting val_start should preserve the legacy behavior exactly.
This ensures existing callers that don't pass val_start get the
same (buggy, but known) output — the fix is opt-in via val_start.
The legacy path may be removed in a later commit.
"""
from case_studies.utils.sequence_dataset import prepare_fold_sequence_stores
df, train_mask, val_mask, _, _ = _synthetic_fold_df()
lookback = 20
_, val_store, fold_info = prepare_fold_sequence_stores(
df,
train_mask=train_mask,
val_mask=val_mask,
feature_names=["feat0", "feat1"],
label_col="y",
date_col="timestamp",
entity_col="symbol",
lookback=lookback,
# val_start intentionally omitted — legacy behavior
)
# In legacy mode, first val sequence should be at position `lookback`
# within the val slice (the bug we're documenting).
for symbol_id in range(val_store.n_symbols):
end_positions = val_store.end_idx[val_store.symbol_idx == symbol_id]
if len(end_positions) == 0:
continue
assert int(end_positions.min()) == lookback, (
"Legacy path should start sequences at position=lookback"
)
tests/test_signals.py
+494
View File
@@ -0,0 +1,494 @@
"""Tests for case_studies/utils/signals.py — prediction → weight contracts.
Signal construction sits on the critical path between every model and every
backtest. A silent behavior change here would corrupt every Ch16-20
strategy result. These tests pin the observable contracts:
- threshold / percentile cutoffs are applied with the documented
inequality semantics (``>`` for fixed threshold, ``>=`` for cross-
sectional percentile, ``>`` for rolling)
- long-short variants produce symmetric signals and weights
- equal-weight top-K weights sum to 1 (or 0 for excluded assets), and
score-weighted weights sum to 1 with score-proportional magnitudes
- the config dispatcher routes every documented method and raises on
unknowns
- ``direction=short_only`` is a pure sign flip of the weight column
- zero-weight rows are filtered from the output
- outputs are deterministic across repeated calls on the same input
"""
from __future__ import annotations
import numpy as np
import polars as pl
import pytest
from polars.testing import assert_frame_equal
from case_studies.utils.signals import (
build_target_weights,
build_target_weights_from_config,
cross_sectional_percentile_signal,
fixed_threshold_signal,
per_symbol_rolling_percentile_signal,
rolling_percentile_signal,
)
# -----------------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------------
@pytest.fixture
def predictions_2d5s() -> pl.DataFrame:
"""2 timestamps × 5 symbols (A–E), y_score ascending per date."""
return pl.DataFrame(
{
"timestamp": ["2024-01-01"] * 5 + ["2024-01-02"] * 5,
"symbol": list("ABCDE") * 2,
"y_score": [0.1, 0.3, 0.5, 0.7, 0.9, 0.2, 0.4, 0.6, 0.8, 1.0],
}
).with_columns(pl.col("timestamp").str.to_date())
@pytest.fixture
def predictions_2d6s() -> pl.DataFrame:
"""2 timestamps × 6 symbols (A–F) for even-split top/bottom tests."""
return pl.DataFrame(
{
"timestamp": ["2024-01-01"] * 6 + ["2024-01-02"] * 6,
"symbol": list("ABCDEF") * 2,
"y_score": [0.1, 0.2, 0.3, 0.7, 0.8, 0.9, 0.1, 0.3, 0.5, 0.6, 0.8, 1.0],
}
).with_columns(pl.col("timestamp").str.to_date())
@pytest.fixture
def predictions_rolling() -> pl.DataFrame:
"""50 timestamps × 2 symbols for rolling-window tests."""
rng = np.random.default_rng(42)
ts = pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 2, 19), "1d", eager=True)
rows = [(t, s, float(rng.random())) for s in ("A", "B") for t in ts]
return pl.DataFrame(rows, schema=["timestamp", "symbol", "y_score"], orient="row").sort(
"timestamp", "symbol"
)
# -----------------------------------------------------------------------------
# fixed_threshold_signal
# -----------------------------------------------------------------------------
def test_fixed_threshold_long_only_strict_greater_than(predictions_2d5s) -> None:
"""signal=1 iff score > threshold (strict). At-threshold scores get 0."""
out = fixed_threshold_signal(predictions_2d5s, threshold=0.5, signal_type="long_only")
# 0.5 → 0 (not strictly >); 0.7/0.9/0.6/0.8/1.0 → 1
expected = [0, 0, 0, 1, 1, 0, 0, 1, 1, 1]
assert out["signal"].to_list() == expected
def test_fixed_threshold_signal_is_int8(predictions_2d5s) -> None:
out = fixed_threshold_signal(predictions_2d5s, threshold=0.5)
assert out["signal"].dtype == pl.Int8
def test_fixed_threshold_preserves_row_count_and_columns(predictions_2d5s) -> None:
out = fixed_threshold_signal(predictions_2d5s, threshold=0.5)
assert out.height == predictions_2d5s.height
assert set(predictions_2d5s.columns) <= set(out.columns)
def test_fixed_threshold_long_short_uses_symmetric_mirror() -> None:
"""long_short with threshold=0.7 → above 0.7 → 1, below (1-0.7)=0.3 → -1."""
df = pl.DataFrame({"y_score": [0.1, 0.4, 0.5, 0.6, 0.9]})
out = fixed_threshold_signal(df, threshold=0.7, signal_type="long_short")
assert out["signal"].to_list() == [-1, 0, 0, 0, 1]
def test_fixed_threshold_deterministic_across_calls(predictions_2d5s) -> None:
a = fixed_threshold_signal(predictions_2d5s, threshold=0.5)
b = fixed_threshold_signal(predictions_2d5s, threshold=0.5)
assert_frame_equal(a, b)
# -----------------------------------------------------------------------------
# rolling_percentile_signal
# -----------------------------------------------------------------------------
def test_rolling_percentile_adds_threshold_column(predictions_rolling) -> None:
out = rolling_percentile_signal(predictions_rolling, window=10, percentile=80.0)
assert "rolling_threshold" in out.columns
def test_rolling_percentile_early_window_has_null_threshold(predictions_rolling) -> None:
"""First window-1 rows per asset have insufficient history → null threshold."""
out = rolling_percentile_signal(predictions_rolling, window=10, percentile=80.0)
# 2 symbols × (window-1=9) early rows = 18 nulls
assert out["rolling_threshold"].null_count() == 18
def test_rolling_percentile_long_short_adds_both_thresholds(predictions_rolling) -> None:
out = rolling_percentile_signal(
predictions_rolling, window=10, percentile=80.0, signal_type="long_short"
)
assert "rolling_threshold" in out.columns
assert "rolling_lower_threshold" in out.columns
# Must produce at least one long and one short signal with random data
counts = dict(out.group_by("signal").len().iter_rows())
assert counts.get(1, 0) > 0
assert counts.get(-1, 0) > 0
def test_rolling_percentile_per_asset_independence() -> None:
"""Each asset computes its own rolling quantile — asset ordering shouldn't
change its own signal sequence."""
ts = pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 1, 20), "1d", eager=True)
rows_a = [(t, "A", float(i)) for i, t in enumerate(ts)]
rows_b = [(t, "B", float(-i)) for i, t in enumerate(ts)]
df = pl.DataFrame(rows_a + rows_b, schema=["timestamp", "symbol", "y_score"], orient="row")
a_only_thresholds = rolling_percentile_signal(
df.filter(pl.col("symbol") == "A"), window=5, percentile=80.0
)["rolling_threshold"]
with_both = rolling_percentile_signal(df, window=5, percentile=80.0).filter(
pl.col("symbol") == "A"
)["rolling_threshold"]
assert a_only_thresholds.to_list() == with_both.to_list()
# -----------------------------------------------------------------------------
# cross_sectional_percentile_signal
# -----------------------------------------------------------------------------
def test_cs_percentile_long_only_at_or_above_cutoff(predictions_2d5s) -> None:
"""cs_percentile uses ``>=`` — score equal to the threshold gets a signal.
At percentile=80 with 5 symbols, the 80th percentile interpolates to the
second-highest score. With ascending scores D=0.7, E=0.9 for date 1,
cs_threshold=0.7 and both D and E get signal=1.
"""
out = cross_sectional_percentile_signal(predictions_2d5s, percentile=80.0).sort(
"timestamp", "symbol"
)
# Per date, top 2 by score should be selected
assert out.filter(pl.col("signal") == 1).height == 4 # 2 dates × 2 winners
def test_cs_percentile_threshold_differs_per_timestamp(predictions_2d5s) -> None:
"""Different dates have different score distributions → different thresholds."""
out = cross_sectional_percentile_signal(predictions_2d5s, percentile=80.0)
per_date = out.group_by("timestamp").agg(pl.col("cs_threshold").first()).sort("timestamp")
thresholds = per_date["cs_threshold"].to_list()
assert thresholds[0] != thresholds[1]
def test_cs_percentile_long_short_produces_both_signs(predictions_2d5s) -> None:
out = cross_sectional_percentile_signal(
predictions_2d5s, percentile=80.0, signal_type="long_short"
)
signs = set(out["signal"].to_list())
assert 1 in signs
assert -1 in signs
# -----------------------------------------------------------------------------
# build_target_weights — equal_weight_top_k
# -----------------------------------------------------------------------------
def test_equal_weight_top_k_long_only_weights_sum_to_1(predictions_2d5s) -> None:
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
per_date = out.group_by("timestamp").agg(pl.col("weight").sum()).sort("timestamp")
for w in per_date["weight"].to_list():
assert abs(w - 1.0) < 1e-9
def test_equal_weight_top_k_selects_exactly_k_assets_per_date(predictions_2d5s) -> None:
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
per_date = out.group_by("timestamp").agg(pl.col("symbol").count().alias("n")).sort("timestamp")
assert per_date["n"].to_list() == [2, 2]
def test_equal_weight_top_k_picks_highest_scores(predictions_2d5s) -> None:
"""With ascending scores A..E, top 2 should always be D and E."""
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
selected = set(out["symbol"].unique().to_list())
assert selected == {"D", "E"}
def test_equal_weight_top_k_long_short_weights_are_symmetric(predictions_2d6s) -> None:
"""long_short top_k=2 with 6 symbols: 2 long @+0.5, 2 short @-0.5, 2 zero (dropped)."""
out = build_target_weights(
predictions_2d6s, method="equal_weight_top_k", top_k=2, long_short=True
)
longs = out.filter(pl.col("weight") > 0)
shorts = out.filter(pl.col("weight") < 0)
assert longs.height == 4 # 2 dates × 2 longs
assert shorts.height == 4
# Magnitudes equal
assert all(abs(w - 0.5) < 1e-9 for w in longs["weight"])
assert all(abs(w + 0.5) < 1e-9 for w in shorts["weight"])
def test_equal_weight_top_k_clamps_when_k_exceeds_n_assets(predictions_2d5s) -> None:
"""Asking for top_k=100 with 5 assets per date → selects all 5, weights = 1/5."""
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=100)
assert out.height == 10 # all rows survive
assert all(abs(w - 0.2) < 1e-9 for w in out["weight"])
def test_equal_weight_top_k_filters_zero_weights(predictions_2d5s) -> None:
"""The helper strips zero-weight rows from the output."""
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
assert (out["weight"] == 0.0).sum() == 0
def test_equal_weight_top_k_output_sorted_by_time_then_asset(predictions_2d5s) -> None:
out = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
pairs = list(zip(out["timestamp"].to_list(), out["symbol"].to_list(), strict=True))
assert pairs == sorted(pairs)
# -----------------------------------------------------------------------------
# build_target_weights — score_weighted_top_k
# -----------------------------------------------------------------------------
def test_score_weighted_top_k_long_only_weights_sum_to_1(predictions_2d5s) -> None:
out = build_target_weights(predictions_2d5s, method="score_weighted_top_k", top_k=2)
per_date = out.group_by("timestamp").agg(pl.col("weight").sum())
for w in per_date["weight"].to_list():
assert abs(w - 1.0) < 1e-9
def test_score_weighted_top_k_weight_proportional_to_abs_score(predictions_2d6s) -> None:
"""Top 2 of [0.8, 0.9] → weights 0.8/1.7 ≈ 0.4706 and 0.9/1.7 ≈ 0.5294."""
out = build_target_weights(predictions_2d6s, method="score_weighted_top_k", top_k=2).sort(
"timestamp", "symbol"
)
date1 = out.filter(pl.col("timestamp") == pl.date(2024, 1, 1)).sort("symbol")
weights = dict(zip(date1["symbol"].to_list(), date1["weight"].to_list(), strict=True))
assert abs(weights["E"] - 0.8 / 1.7) < 1e-9
assert abs(weights["F"] - 0.9 / 1.7) < 1e-9
def test_score_weighted_top_k_deterministic(predictions_2d6s) -> None:
a = build_target_weights(predictions_2d6s, method="score_weighted_top_k", top_k=2)
b = build_target_weights(predictions_2d6s, method="score_weighted_top_k", top_k=2)
assert_frame_equal(a.sort("timestamp", "symbol"), b.sort("timestamp", "symbol"))
# -----------------------------------------------------------------------------
# build_target_weights — inverse_vol (placeholder path: equal weight)
# -----------------------------------------------------------------------------
def test_inverse_vol_placeholder_uses_equal_weight(predictions_2d5s) -> None:
"""inverse_vol is documented as a placeholder — same output as equal_weight_top_k."""
eq = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
iv = build_target_weights(predictions_2d5s, method="inverse_vol", top_k=2)
assert_frame_equal(
eq.sort("timestamp", "symbol"),
iv.sort("timestamp", "symbol"),
)
# -----------------------------------------------------------------------------
# build_target_weights_from_config — dispatcher
# -----------------------------------------------------------------------------
def test_from_config_equal_weight_top_k_matches_direct_call(predictions_2d5s) -> None:
direct = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
via = build_target_weights_from_config(
predictions_2d5s, {"method": "equal_weight_top_k", "top_k": 2}
)
assert_frame_equal(direct.sort("timestamp", "symbol"), via.sort("timestamp", "symbol"))
def test_from_config_decile_long_short_on_small_universe(predictions_2d6s) -> None:
"""6 symbols, decile → top_cutoff=floor(6/10)=0 clipped to 1 → 1 long, 1 short."""
out = build_target_weights_from_config(predictions_2d6s, {"method": "decile_long_short"}).sort(
"timestamp", "symbol"
)
# Per date: 1 long @ +1.0 (top score), 1 short @ -1.0 (bottom score)
assert out.height == 4
assert sorted(out["weight"].unique().to_list()) == [-1.0, 1.0]
def test_from_config_cross_sectional_percentile(predictions_2d5s) -> None:
out = build_target_weights_from_config(
predictions_2d5s,
{"method": "cross_sectional_percentile", "percentile": 80.0},
)
# Top 2 assets per date → weights sum to 1 per date
per_date = out.group_by("timestamp").agg(pl.col("weight").sum())
for w in per_date["weight"].to_list():
assert abs(w - 1.0) < 1e-9
def test_from_config_fixed_threshold_selects_above_cutoff(predictions_2d5s) -> None:
out = build_target_weights_from_config(
predictions_2d5s, {"method": "fixed_threshold", "threshold": 0.5}
)
# Per date: D (0.7), E (0.9) → 2 assets @ 0.5 each, sum to 1 on date 1.
# Date 2: C (0.6), D (0.8), E (1.0) → 3 assets @ 1/3 each.
date1 = out.filter(pl.col("timestamp") == pl.date(2024, 1, 1))
date2 = out.filter(pl.col("timestamp") == pl.date(2024, 1, 2))
assert date1.height == 2 and abs(date1["weight"].sum() - 1.0) < 1e-9
assert date2.height == 3 and abs(date2["weight"].sum() - 1.0) < 1e-9
def test_from_config_short_only_negates_weights(predictions_2d5s) -> None:
"""direction=short_only flips signs; magnitudes identical to long_only."""
long_w = build_target_weights_from_config(
predictions_2d5s, {"method": "equal_weight_top_k", "top_k": 2}
)
short_w = build_target_weights_from_config(
predictions_2d5s,
{"method": "equal_weight_top_k", "top_k": 2, "direction": "short_only"},
)
# Sort and pair up, then verify the negation contract
long_sorted = long_w.sort("timestamp", "symbol")
short_sorted = short_w.sort("timestamp", "symbol")
assert long_sorted["symbol"].to_list() == short_sorted["symbol"].to_list()
for lw, sw in zip(
long_sorted["weight"].to_list(), short_sorted["weight"].to_list(), strict=True
):
assert abs(lw + sw) < 1e-9
def test_from_config_rejects_unknown_method(predictions_2d5s) -> None:
with pytest.raises(ValueError, match="Unknown signal method"):
build_target_weights_from_config(predictions_2d5s, {"method": "bogus"})
def test_from_config_rejects_unknown_direction(predictions_2d5s) -> None:
with pytest.raises(ValueError, match="Unknown signal direction"):
build_target_weights_from_config(
predictions_2d5s,
{"method": "equal_weight_top_k", "top_k": 2, "direction": "bogus"},
)
def test_from_config_quintile_long_short_uses_5_buckets(predictions_2d5s) -> None:
"""quintile with 5 assets → top_cutoff=1 → 1 long, 1 short per date."""
out = build_target_weights_from_config(predictions_2d5s, {"method": "quintile_long_short"})
per_date = out.group_by("timestamp").agg(pl.col("symbol").count().alias("n"))
assert per_date["n"].to_list() == [2, 2] # 1 long + 1 short each date
# -----------------------------------------------------------------------------
# Determinism
# -----------------------------------------------------------------------------
def test_cross_sectional_percentile_deterministic(predictions_2d5s) -> None:
a = cross_sectional_percentile_signal(predictions_2d5s, percentile=80.0)
b = cross_sectional_percentile_signal(predictions_2d5s, percentile=80.0)
assert_frame_equal(a, b)
def test_rolling_percentile_deterministic(predictions_rolling) -> None:
a = rolling_percentile_signal(predictions_rolling, window=10, percentile=80.0)
b = rolling_percentile_signal(predictions_rolling, window=10, percentile=80.0)
assert_frame_equal(a, b)
def test_build_target_weights_deterministic(predictions_2d5s) -> None:
a = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
b = build_target_weights(predictions_2d5s, method="equal_weight_top_k", top_k=2)
assert_frame_equal(a, b)
# -----------------------------------------------------------------------------
# per_symbol_rolling_percentile_signal — stay_q extension
# -----------------------------------------------------------------------------
@pytest.fixture
def per_symbol_intraday() -> pl.DataFrame:
"""30 days × 2 symbols × 14 bars/day; deterministic seeded scores."""
from datetime import datetime, timedelta
rng = np.random.default_rng(11)
rows = []
for d in range(30):
for i in range(14):
ts = datetime(2024, 1, 2, 9, 30) + timedelta(days=d, minutes=15 * i)
for sym in ("AAA", "BBB"):
rows.append((ts, sym, float(rng.standard_normal())))
return pl.DataFrame(rows, schema=["timestamp", "symbol", "y_score"], orient="row").sort(
"symbol", "timestamp"
)
def test_per_symbol_default_excludes_stay_thresh(per_symbol_intraday) -> None:
"""When stay_q is None, stay_thresh column is NOT present (back-compat)."""
out = per_symbol_rolling_percentile_signal(
per_symbol_intraday,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
)
assert "stay_thresh" not in out.columns
assert "signal" in out.columns
def test_per_symbol_stay_q_adds_stay_thresh(per_symbol_intraday) -> None:
"""When stay_q is set, stay_thresh column is added; non-null after warm-up."""
out = per_symbol_rolling_percentile_signal(
per_symbol_intraday,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
stay_q=0.40,
)
assert "stay_thresh" in out.columns
# After warm-up (~5 sessions = 70 bars per symbol), stay_thresh should be non-null
by_sym = out.group_by("symbol").agg(pl.col("stay_thresh").is_not_null().sum().alias("n"))
for n in by_sym["n"].to_list():
assert n > 100 # well past warm-up of W//2 = 70
def test_per_symbol_stay_thresh_monotonic_in_stay_q(per_symbol_intraday) -> None:
"""stay_thresh must increase monotonically with stay_q, and thus stay below
the entry threshold (long_q) for any stay_q < long_q.
long_thresh is dropped from the output, and the function forbids
``stay_q == long_q``, so we cannot read long_thresh directly. Instead we
verify the underlying invariant black-box: a lower stay_q must yield a
stay_thresh at or below a higher stay_q's on the same row. Since the entry
threshold is the long_q quantile, monotonicity transitively guarantees
every ``stay_q < long_q`` threshold sits below it. This catches a sign flip
between stay_thresh and long_thresh, unlike the previous tautological
"score exceeds the threshold that fired it" check.
"""
kw = dict(long_q=0.80, lookback_days=10, bars_per_day=14)
lo = per_symbol_rolling_percentile_signal(per_symbol_intraday, stay_q=0.40, **kw)
hi = per_symbol_rolling_percentile_signal(per_symbol_intraday, stay_q=0.79, **kw)
joined = (
lo.select(["symbol", "timestamp", "stay_thresh"])
.join(
hi.select(["symbol", "timestamp", pl.col("stay_thresh").alias("stay_thresh_hi")]),
on=["symbol", "timestamp"],
how="inner",
)
.filter(pl.col("stay_thresh").is_not_null() & pl.col("stay_thresh_hi").is_not_null())
)
assert joined.height > 100 # well past warm-up
# q=0.40 quantile must never exceed the q=0.79 quantile (< the q=0.80 entry).
assert (joined["stay_thresh"] - joined["stay_thresh_hi"]).max() <= 1e-9
def test_per_symbol_rejects_stay_q_at_or_above_long_q(per_symbol_intraday) -> None:
with pytest.raises(ValueError, match="stay_q must be < long_q"):
per_symbol_rolling_percentile_signal(
per_symbol_intraday,
long_q=0.60,
lookback_days=10,
bars_per_day=14,
stay_q=0.60,
)
tests/test_slot_strategy.py
+326
View File
@@ -0,0 +1,326 @@
"""Tests for case_studies/utils/slot_strategy.py — persistent-slot selection.
The slot mechanism is a new selection method introduced for intraday case
studies where per-symbol score distributions and signal-based exits matter.
These tests pin the observable contracts of the high-level
``build_persistent_slot_weights_hybrid`` entry plus the underlying
``_run_slot_simulation`` mechanism:
- max-hold caps position age regardless of score
- signal-exit fires when current score < stay threshold
- capacity is respected (max_slots concurrent holdings)
- new entries are score-ordered when capacity is constrained
- short_only flips the weight sign
- stale-pred rows (older than freshness tolerance) are dropped before entry
- empty input returns empty frame with canonical schema
"""
from __future__ import annotations
from datetime import datetime, timedelta
import numpy as np
import polars as pl
import pytest
from case_studies.utils.slot_strategy import (
_align_predictions_to_bars,
_run_slot_simulation,
build_persistent_slot_weights_hybrid,
)
# -----------------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------------
def _bars(n: int, start: datetime | None = None, step: timedelta | None = None) -> list[datetime]:
"""Generate ``n`` evenly spaced bar timestamps."""
start = start or datetime(2024, 1, 2, 9, 30)
step = step or timedelta(minutes=15)
return [start + i * step for i in range(n)]
# -----------------------------------------------------------------------------
# _run_slot_simulation — pure mechanism
# -----------------------------------------------------------------------------
def test_simulation_single_symbol_fill_then_maxhold_exit() -> None:
"""One symbol enters at bar 0, must exit at bar 4 when hold_bars=4."""
bars = _bars(8)
signals = {bars[0]: [("AAA", 1.0)]}
weights, stats = _run_slot_simulation(
signals_by_ts=signals,
all_bars_sorted=bars,
max_slots=1,
weight_per_slot=1.0,
hold_bars=4,
score_by_ts_sym=None,
stay_threshold_by_ts_sym=None,
)
# Held bars 0..3 inclusive, exits at bar 4 (entry_i=0, i-entry_i=4 >= hold_bars)
held_ts = weights["timestamp"].to_list()
assert held_ts == bars[:4]
assert (weights["symbol"] == "AAA").all()
assert stats["n_entries"] == 1
assert stats["n_exits_maxhold"] == 1
assert stats["n_exits_signal"] == 0
def test_simulation_max_slots_caps_concurrent_holdings() -> None:
"""5 symbols signal simultaneously, max_slots=2 keeps top-2 by score."""
bars = _bars(3)
signals = {bars[0]: [(s, sc) for s, sc in zip("ABCDE", [0.1, 0.9, 0.5, 0.7, 0.3])]}
weights, stats = _run_slot_simulation(
signals_by_ts=signals,
all_bars_sorted=bars,
max_slots=2,
weight_per_slot=0.5,
hold_bars=10,
score_by_ts_sym=None,
stay_threshold_by_ts_sym=None,
)
held_first_bar = set(weights.filter(pl.col("timestamp") == bars[0])["symbol"].to_list())
assert held_first_bar == {"B", "D"} # top-2 scores 0.9 and 0.7
assert stats["n_entries"] == 2
def test_simulation_signal_exit_fires_when_score_below_stay() -> None:
"""Signal-exit triggers when current score drops below stay threshold."""
bars = _bars(5)
signals = {bars[0]: [("AAA", 0.9)]}
score_lookup = {(bars[i], "AAA"): 0.9 if i < 2 else 0.1 for i in range(5)}
stay_lookup = {(bars[i], "AAA"): 0.5 for i in range(5)}
weights, stats = _run_slot_simulation(
signals_by_ts=signals,
all_bars_sorted=bars,
max_slots=1,
weight_per_slot=1.0,
hold_bars=10,
score_by_ts_sym=score_lookup,
stay_threshold_by_ts_sym=stay_lookup,
)
held_ts = weights["timestamp"].to_list()
# Held at bars 0, 1; at bar 2 score (0.1) < stay (0.5) → exit at start of bar 2
assert held_ts == bars[:2]
assert stats["n_exits_signal"] == 1
assert stats["n_exits_maxhold"] == 0
def test_simulation_signal_exit_skipped_when_score_unknown() -> None:
"""If a (ts,sym) is missing from score_lookup, signal-exit must not fire."""
bars = _bars(4)
signals = {bars[0]: [("AAA", 0.9)]}
score_lookup = {(bars[0], "AAA"): 0.9} # only bar 0
stay_lookup = {(bars[i], "AAA"): 0.5 for i in range(4)}
weights, stats = _run_slot_simulation(
signals_by_ts=signals,
all_bars_sorted=bars,
max_slots=1,
weight_per_slot=1.0,
hold_bars=10,
score_by_ts_sym=score_lookup,
stay_threshold_by_ts_sym=stay_lookup,
)
# Missing scores at bars 1,2,3 → never signal-exit. Held all 4 bars.
assert weights.height == 4
assert stats["n_exits_signal"] == 0
def test_simulation_validates_positive_max_slots() -> None:
with pytest.raises(ValueError, match="max_slots must be positive"):
_run_slot_simulation(
signals_by_ts={},
all_bars_sorted=[],
max_slots=0,
weight_per_slot=1.0,
hold_bars=1,
score_by_ts_sym=None,
stay_threshold_by_ts_sym=None,
)
def test_simulation_validates_weight_per_slot_range() -> None:
with pytest.raises(ValueError, match="weight_per_slot must be in"):
_run_slot_simulation(
signals_by_ts={},
all_bars_sorted=[],
max_slots=1,
weight_per_slot=1.5,
hold_bars=1,
score_by_ts_sym=None,
stay_threshold_by_ts_sym=None,
)
def test_simulation_empty_signals_returns_empty_frame_with_schema() -> None:
weights, stats = _run_slot_simulation(
signals_by_ts={},
all_bars_sorted=_bars(3),
max_slots=1,
weight_per_slot=1.0,
hold_bars=5,
score_by_ts_sym=None,
stay_threshold_by_ts_sym=None,
)
assert weights.is_empty()
assert weights.columns == ["timestamp", "symbol", "weight"]
assert weights.schema["timestamp"] == pl.Datetime("us")
assert stats["n_entries"] == 0
assert stats["n_exits_total"] == 0
# -----------------------------------------------------------------------------
# _align_predictions_to_bars — backward-asof staleness filter
# -----------------------------------------------------------------------------
def test_align_drops_predictions_older_than_freshness_tolerance() -> None:
"""Predictions older than ``pred_freshness_max_min`` are filtered out."""
bar_grid = pl.DataFrame(
{
"symbol": ["AAA"] * 3,
"timestamp": _bars(3), # 09:30, 09:45, 10:00
}
)
# One prediction at 09:30 (fresh), one at 09:20 (stale for 09:45 bar with 14m tol)
preds = pl.DataFrame(
{
"symbol": ["AAA", "AAA"],
"timestamp": [datetime(2024, 1, 2, 9, 30), datetime(2024, 1, 2, 9, 20)],
"y_score": [0.5, 0.3],
}
)
aligned = _align_predictions_to_bars(
preds,
bar_grid,
pred_freshness_max_min=14,
score_col="y_score",
time_col="timestamp",
asset_col="symbol",
)
# 09:30 bar: 09:30 pred (0m stale) -> 0.5
# 09:45 bar: 09:30 pred (15m stale) -> dropped; 09:20 also too old
# 10:00 bar: same — all preds >14m stale
aligned_ts = aligned["timestamp"].to_list()
assert aligned_ts == [datetime(2024, 1, 2, 9, 30)]
assert aligned["y_score"].to_list() == [0.5]
# -----------------------------------------------------------------------------
# build_persistent_slot_weights_hybrid — high-level entry
# -----------------------------------------------------------------------------
@pytest.fixture
def predictions_dense() -> pl.DataFrame:
"""50 days × 3 symbols × 14 bars/day; deterministic seeded scores."""
rng = np.random.default_rng(7)
rows = []
for d in range(50):
for i in range(14):
ts = datetime(2024, 1, 2, 9, 30) + timedelta(days=d, minutes=15 * i)
for sym in ("AAA", "BBB", "CCC"):
rows.append((ts, sym, float(rng.standard_normal())))
return pl.DataFrame(rows, schema=["timestamp", "symbol", "y_score"], orient="row").sort(
"symbol", "timestamp"
)
def test_build_weights_returns_canonical_schema(predictions_dense) -> None:
prices = predictions_dense.select(["symbol", "timestamp"]).with_columns(close=pl.lit(100.0))
weights, stats = build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.90,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=4,
)
assert set(weights.columns) == {"timestamp", "symbol", "weight"}
assert weights.schema["timestamp"] == pl.Datetime("us")
# weight defaults to 1/max_slots
if not weights.is_empty():
assert (weights["weight"] - 0.5).abs().max() < 1e-9
assert stats["max_slots"] == 2
assert stats["direction"] == "long_only"
def test_build_weights_short_only_flips_sign(predictions_dense) -> None:
prices = predictions_dense.select(["symbol", "timestamp"]).with_columns(close=pl.lit(100.0))
long_w, _ = build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=4,
direction="long_only",
)
short_w, _ = build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=4,
direction="short_only",
)
assert long_w.shape == short_w.shape
if not long_w.is_empty():
# short_only is a pure sign flip
merged = long_w.join(short_w, on=["timestamp", "symbol"], suffix="_s")
assert (merged["weight"] + merged["weight_s"]).abs().max() < 1e-9
def test_build_weights_rejects_long_short_direction(predictions_dense) -> None:
prices = predictions_dense.select(["symbol", "timestamp"]).with_columns(close=pl.lit(100.0))
with pytest.raises(ValueError, match="long_short is not supported"):
build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=4,
direction="long_short", # type: ignore[arg-type]
)
def test_build_weights_rejects_stay_q_at_or_above_long_q(predictions_dense) -> None:
prices = predictions_dense.select(["symbol", "timestamp"]).with_columns(close=pl.lit(100.0))
with pytest.raises(ValueError, match="must be < long_q"):
build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.50,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=4,
exit_signal_q=0.50,
)
def test_build_weights_with_stay_threshold_runs_clean(predictions_dense) -> None:
"""End-to-end with signal-exit enabled — schema + non-degenerate stats."""
prices = predictions_dense.select(["symbol", "timestamp"]).with_columns(close=pl.lit(100.0))
_weights, stats = build_persistent_slot_weights_hybrid(
predictions_dense,
prices,
long_q=0.80,
lookback_days=10,
bars_per_day=14,
max_slots=2,
hold_bars=8,
exit_signal_q=0.40,
)
assert stats["exit_signal_q"] == 0.40
# Either path can produce zero entries on a tiny synthetic sample, but the
# mechanism must not crash and stats must be coherent.
assert stats["n_exits_total"] == stats["n_exits_maxhold"] + stats["n_exits_signal"]
tests/test_sweep_config_seam.py
+249
View File
@@ -0,0 +1,249 @@
"""Drift-detector tests for the setup.yaml-driven Ch16-19 sweep loader.
Pins the contract between ``case_studies/{cs}/config/setup.yaml::backtest.sweep``
and the helpers in ``case_studies.utils.sweep_config``:
1. **Loader shape** — ``load_sweep`` and the ``*_for`` / ``get_*`` helpers
return the expected types and values for migrated case studies.
2. **Registry reconciliation** — the declared sweep covers the rank-1
``(method, top_k)`` for every label in ``labels.{primary, variants}``.
3. **Quarantine policy** — V3/V4 deprecated classes (``score_weighted_top_k``
on the signal stage, ``cross_sectional_percentile``) must not appear in
the declared sweep.
All 9 case studies have shipped ``backtest.sweep``; ``MIGRATED_CASES`` is
now the full set.
"""
from __future__ import annotations
import sqlite3
from pathlib import Path
import pytest
import yaml
from case_studies.utils.sweep_config import (
get_allocators,
get_cost_grid_bps,
get_entry_schemes_for,
get_portfolio_risk_controls,
get_position_risk_controls,
get_top_k_values_for,
load_sweep,
)
from utils import CASE_STUDIES_DIR
# All case studies ship a ``backtest.sweep`` block in setup.yaml.
MIGRATED_CASES: tuple[str, ...] = (
"us_firm_characteristics",
"etfs",
"fx_pairs",
"cme_futures",
"nasdaq100_microstructure",
"us_equities_panel",
"sp500_equity_option_analytics",
"crypto_perps_funding",
"sp500_options",
)
# Signal-stage methods that must never appear in a declared sweep. These
# correspond to the V3/V4 quarantine list: ``score_weighted_top_k`` is an
# allocator (Ch17), not a signal scheme; ``cross_sectional_percentile``
# was retired during V3 cleanup.
QUARANTINED_SIGNAL_METHODS: frozenset[str] = frozenset(
{"score_weighted_top_k", "cross_sectional_percentile"}
)
# ---------------------------------------------------------------------------
# Loader contract — us_firm_characteristics (the first migrated case study)
# ---------------------------------------------------------------------------
class TestUsFirmLoader:
"""Pin the loader output for us_firm_characteristics."""
CS = "us_firm_characteristics"
def test_load_sweep_returns_expected_keys(self):
sweep = load_sweep(self.CS)
assert set(sweep.keys()) >= {
"top_k_grid",
"allocators",
"cost_grid_bps",
"risk_controls",
}
@pytest.mark.parametrize("label", ["fwd_ret_1m", "fwd_ret_1m_win", "fwd_class_1m"])
def test_top_k_grid_per_label(self, label):
assert get_top_k_values_for(self.CS, label, n_assets=2500) == [5, 10, 20, 50]
@pytest.mark.parametrize("label", ["fwd_ret_1m", "fwd_ret_1m_win", "fwd_class_1m"])
def test_entry_schemes_per_label(self, label):
schemes = get_entry_schemes_for(self.CS, label, n_assets=2500, long_short=True)
# Exactly the four equal_weight_top_k schemes declared in setup.yaml.
assert [s["method"] for s in schemes] == ["equal_weight_top_k"] * 4
assert [s["top_k"] for s in schemes] == [5, 10, 20, 50]
assert all(s["long_short"] is True for s in schemes)
def test_allocators_strip_name(self):
allocators = get_allocators(self.CS)
methods = [a["method"] for a in allocators]
# us_firm is a returns-only firm-characteristics panel with no per-symbol
# price series, so the moment-based allocators (inverse_vol, risk_parity,
# hrp, mvo_ledoit_wolf) are intentionally excluded — only the
# lookback-free allocators are declared (see setup.yaml allocators block).
assert methods == [
"equal_weight",
"score_weighted",
"conformal_weighted",
]
# No allocator dict should carry the human-readable ``name`` key —
# the spec hash is computed from the shape that the dispatcher sees.
assert all("name" not in a for a in allocators)
def test_cost_grid_bps(self):
assert get_cost_grid_bps(self.CS) == [0, 1, 2, 3, 5, 7, 10, 15, 20, 30, 50]
def test_risk_control_counts(self):
# Position grid: stop_loss (4) + trailing_stop (7) + time_exit (3) = 14
assert len(get_position_risk_controls(self.CS)) == 14
# Portfolio-level overlays were removed from all setup.yaml (2026-05-17);
# only position-level controls are retained.
assert len(get_portfolio_risk_controls(self.CS)) == 0
# ---------------------------------------------------------------------------
# Quarantine policy — V3/V4 retired classes must not appear in any declared
# sweep
# ---------------------------------------------------------------------------
class TestQuarantinePolicy:
"""Declared sweep must not include V3/V4 retired selection classes."""
@pytest.mark.parametrize("case_study", MIGRATED_CASES)
def test_no_score_weighted_top_k_in_signal_sweep(self, case_study):
sweep = load_sweep(case_study)
# ``score_weighted_top_k`` should never appear as a top_k_grid /
# percentile_grid / quantile_grid axis — it belongs in
# ``allocators`` only.
# The synthesized entry schemes are the SUT.
for label in (
(sweep.get("top_k_grid") or {}).keys()
| (sweep.get("percentile_grid") or {}).keys()
| (sweep.get("quantile_grid") or {}).keys()
):
schemes = get_entry_schemes_for(case_study, label, n_assets=10_000, long_short=False)
methods = {s["method"] for s in schemes}
assert methods.isdisjoint(QUARANTINED_SIGNAL_METHODS), (
f"{case_study}/{label}: quarantined signal method appeared: "
f"{methods & QUARANTINED_SIGNAL_METHODS}"
)
# ---------------------------------------------------------------------------
# Registry reconciliation — declared sweep covers rank-1 per (CS, label)
# ---------------------------------------------------------------------------
def _registry_path(case_study: str) -> Path:
return CASE_STUDIES_DIR / case_study / "run_log" / "registry.db"
def _labels_for(case_study: str) -> list[str]:
setup = yaml.safe_load((CASE_STUDIES_DIR / case_study / "config" / "setup.yaml").read_text())
labels_block = setup.get("labels") or {}
primary = labels_block.get("primary")
variants = labels_block.get("variants") or []
return [primary, *variants] if primary else list(variants)
class TestRegistryReconciliation:
"""For every label in setup.yaml::labels.{primary,variants}, the rank-1
signal-stage row in ``backtest_runs`` must use a ``(method, top_k)`` that
is in the declared sweep.
Skips a case study if its registry has no signal-stage rows (e.g.,
immediately after a registry cleanup), since there is nothing to
reconcile against yet. As each case study completes its Ch16-19 wrap-up,
its registry rank-1 should reappear and this test should pass.
"""
@pytest.mark.parametrize("case_study", MIGRATED_CASES)
def test_rank1_signal_method_in_declared_sweep(self, case_study):
reg_path = _registry_path(case_study)
if not reg_path.exists():
pytest.skip(f"{case_study}: registry.db not present")
sweep = load_sweep(case_study)
top_k_by_label = sweep.get("top_k_grid") or {}
qnt_by_label = sweep.get("quantile_grid") or {}
pct_by_label = sweep.get("percentile_grid") or {}
labels = _labels_for(case_study)
with sqlite3.connect(reg_path) as conn:
cur = conn.cursor()
for label in labels:
row = cur.execute(
"""
SELECT json_extract(r.spec_json, '$.strategy.signal.method'),
json_extract(r.spec_json, '$.strategy.signal.top_k'),
json_extract(r.spec_json, '$.strategy.signal.n_quantiles'),
json_extract(r.spec_json, '$.strategy.signal.max_slots'),
bm.sharpe
FROM backtest_metrics bm
JOIN backtest_runs r ON r.backtest_hash = bm.backtest_hash
JOIN prediction_sets p ON p.prediction_hash = r.prediction_hash
JOIN training_runs t ON t.training_hash = p.training_hash
WHERE t.label = ? AND p.split = 'validation' AND r.stage = 'signal'
ORDER BY bm.sharpe DESC LIMIT 1
""",
(label,),
).fetchone()
if row is None:
pytest.skip(f"{case_study}/{label}: no signal-stage rows in registry")
method, top_k, n_quantiles, max_slots, _sharpe = row
if method == "equal_weight_top_k":
declared_ks = list(top_k_by_label.get(label, []))
assert top_k in declared_ks, (
f"{case_study}/{label}: registry rank-1 top_k={top_k} "
f"not in declared top_k_grid={declared_ks}"
)
elif method == "slot_persistent_signal_exit":
# Slot-mechanism signal (nasdaq100 v4 microstructure sweep).
# Slots ARE the allocation, so the swept parameter is
# ``max_slots`` (declared under the ``signal_nasdaq100``
# block), not top_k / n_quantiles.
declared_slots = list(
(sweep.get("signal_nasdaq100") or {}).get("max_slots", [])
)
assert max_slots in declared_slots, (
f"{case_study}/{label}: registry rank-1 max_slots={max_slots} "
f"not in declared signal_nasdaq100.max_slots={declared_slots}"
)
elif method in ("quintile_long_short", "decile_long_short"):
declared_qs = list(qnt_by_label.get(label, []))
assert n_quantiles in declared_qs, (
f"{case_study}/{label}: registry rank-1 n_quantiles="
f"{n_quantiles} not in declared quantile_grid={declared_qs}"
)
elif method in QUARANTINED_SIGNAL_METHODS:
# The registry still has V3/V4 debris — Ch16-19 sweep
# cleanup hasn't run yet for this (case_study, label).
# Skip rather than fail; once cleanup runs, the rank-1
# method will be canonical and the assertion above takes
# over.
pytest.skip(
f"{case_study}/{label}: rank-1 is {method!r} (V3/V4 "
f"quarantine class). Registry cleanup pending; "
f"re-rank after task-6/task-7 land."
)
else:
pytest.fail(
f"{case_study}/{label}: rank-1 method {method!r} is "
f"unrecognized by the seam test — extend the test or "
f"the loader."
)
tests/test_uncertainty_cscv.py
+110
View File
@@ -0,0 +1,110 @@
"""Tests for case_studies/utils/uncertainty.py CSCV partition + PBO smoke.
Covers two pieces P2.5 added:
1. ``_cscv_split_pairs`` — IS/OOS partition shape and balance for
``n_folds`` in {2, 3, 4}, including the asymmetric odd-fold case.
2. ``compute_cohort_metrics`` end-to-end with a ``fold_returns_by_hash``
argument, asserting that ``pbo`` / ``pbo_median_oos_rank`` /
``pbo_mean_degradation`` come back populated (i.e. the
``compute_pbo`` field-name and partition wiring is intact).
"""
from __future__ import annotations
from math import comb
import numpy as np
import polars as pl
import pytest
@pytest.mark.parametrize(
"n_folds, is_half, oos_half",
[
(2, 1, 1), # balanced
(3, 1, 2), # asymmetric — OOS gets the extra fold
(4, 2, 2), # balanced
],
)
def test_cscv_split_pairs_partition_shape(n_folds: int, is_half: int, oos_half: int) -> None:
from case_studies.utils.uncertainty import _cscv_split_pairs
rng = np.random.default_rng(0)
k_variants = 5
fold_sharpes = rng.normal(size=(n_folds, k_variants))
is_perf, oos_perf = _cscv_split_pairs(fold_sharpes)
expected_n = comb(n_folds, n_folds // 2)
assert is_perf.shape == (expected_n, k_variants)
assert oos_perf.shape == (expected_n, k_variants)
# Every row must be the mean of `is_half` folds (IS) and
# `oos_half` folds (OOS) of the original matrix — verified by
# reconstructing the underlying sums.
for row_is, row_oos in zip(is_perf, oos_perf, strict=True):
# IS mean × is_half + OOS mean × oos_half == sum of all folds
total = fold_sharpes.sum(axis=0)
reconstructed = row_is * is_half + row_oos * oos_half
np.testing.assert_allclose(reconstructed, total, atol=1e-12)
def test_cscv_split_pairs_single_fold_returns_empty() -> None:
from case_studies.utils.uncertainty import _cscv_split_pairs
is_perf, oos_perf = _cscv_split_pairs(np.array([[1.0, 2.0, 3.0]]))
assert is_perf.shape == (0, 3)
assert oos_perf.shape == (0, 3)
def test_compute_cohort_metrics_populates_pbo_with_fold_returns() -> None:
"""End-to-end smoke: PBO fields must come back non-null when
``fold_returns_by_hash`` is supplied for >=2 variants with >=2 folds.
The pre-P2.5 code called ``compute_pbo(fs, fs)`` and read the wrong
PBOResult attribute names — both bugs would surface here as NULLs.
"""
from case_studies.utils.uncertainty import compute_cohort_metrics
rng = np.random.default_rng(7)
n_periods = 252
timestamps = pl.datetime_range(
start=pl.datetime(2020, 1, 1),
end=pl.datetime(2020, 12, 31),
interval="1d",
eager=True,
).head(n_periods)
def _make_frame(mu: float) -> pl.DataFrame:
ret = rng.normal(loc=mu / 252, scale=0.01, size=n_periods)
return pl.DataFrame({"timestamp": timestamps, "ret": ret})
# Three "variants" with hash-shaped keys (32 hex chars satisfies any
# downstream FK convention; here we just need stable dict keys).
returns_by_hash = {f"{i:032x}": _make_frame(mu=mu) for i, mu in enumerate([0.05, 0.08, 0.12])}
n_folds = 4
fold_returns_by_hash = {
h: rng.normal(loc=0.0, scale=1.0, size=n_folds) for h in returns_by_hash
}
out = compute_cohort_metrics(
returns_by_hash,
periods_per_year=252.0,
fold_returns_by_hash=fold_returns_by_hash,
rademacher_n_simulations=50,
rademacher_seed=0,
)
assert out, "compute_cohort_metrics returned empty dict — alignment failed"
assert out["leader_hash"] in returns_by_hash
assert out["k_variants"] == 3
# PBO fields must be populated (the bug-surface check).
assert out["pbo"] is not None
assert 0.0 <= out["pbo"] <= 1.0
assert out["pbo_n_combinations"] == float(comb(n_folds, n_folds // 2))
assert out["pbo_median_oos_rank"] is not None
assert out["pbo_mean_degradation"] is not None
assert out["pbo_n_folds"] == float(n_folds)
tests/test_warmup_periods_for.py
+103
View File
@@ -0,0 +1,103 @@
"""Tests for ``case_studies.utils.backtest_loaders.warmup_periods_for`` +
``_calendar_days_per_period`` — the helpers that replaced the hardcoded
``warmup_periods=126`` constant duplicated across 16 call-sites in 5 CSes.
These tests close P2.8 of the roborev cleanup (review #2510 / #2511).
"""
from __future__ import annotations
import yaml
from case_studies.utils.backtest_loaders import (
_calendar_days_per_period,
_load_case_setup_yaml,
warmup_periods_for,
)
# The expected per-CS warmup is the max over ``execution.allocator_lookback``
# and any per-sweep allocator ``vol_window`` / ``lookback`` overrides. These
# expectations are anchored on the current setup.yaml values; if a CS
# tunes its allocator lookbacks, update the expected value here.
_EXPECTED_WARMUP: dict[str, int] = {
"etfs": 63,
"crypto_perps_funding": 240,
"nasdaq100_microstructure": 520,
"us_equities_panel": 126, # mvo_ledoit_wolf lookback=126 > allocator_lookback=63
"us_firm_characteristics": 12,
"fx_pairs": 63,
"cme_futures": 63,
"sp500_options": 63,
"sp500_equity_option_analytics": 126, # mvo_ledoit_wolf lookback=126
}
def test_warmup_periods_for_matches_setup_yaml() -> None:
for cs, expected in _EXPECTED_WARMUP.items():
actual = warmup_periods_for(cs)
assert actual == expected, (
f"warmup_periods_for({cs}) = {actual}, expected {expected} "
f"(max of execution.allocator_lookback + sweep allocator overrides)"
)
def test_warmup_periods_for_unknown_returns_zero(tmp_path) -> None:
# No setup.yaml → defaults to 0 (the unbounded fallback inside
# load_backtest_prices_for then skips the prefix-day computation).
assert warmup_periods_for("__nonexistent_cs__") == 0
def test_warmup_periods_for_picks_max_over_overrides(tmp_path, monkeypatch) -> None:
"""When a per-allocator override exceeds allocator_lookback, the helper
must surface the override rather than the CS-level default."""
fake_setup = {
"execution": {"allocator_lookback": 50},
"backtest": {
"sweep": {
"allocators": [
{"name": "equal_weight"},
{"name": "iv", "vol_window": 200},
{"name": "mvo_lw", "lookback": 100},
]
}
},
}
cs_dir = tmp_path / "fake_cs" / "config"
cs_dir.mkdir(parents=True)
(cs_dir / "setup.yaml").write_text(yaml.safe_dump(fake_setup))
# Drop the cache so the synthetic CS gets a fresh read.
_load_case_setup_yaml.cache_clear()
from case_studies.utils.backtest_loaders import warmup_periods_for as wpf
from utils.paths import get_case_study_dir as orig_get_dir
def fake_get_dir(cs: str):
if cs == "fake_cs":
return tmp_path / "fake_cs"
return orig_get_dir(cs)
monkeypatch.setattr("case_studies.utils.backtest_loaders.get_case_study_dir", fake_get_dir)
_load_case_setup_yaml.cache_clear()
assert wpf("fake_cs") == 200
def test_calendar_days_per_period_cadence_aware() -> None:
# Daily NYSE cadence: 1.5× (weekend + holiday allowance)
assert abs(_calendar_days_per_period("fx_pairs") - 1.5) < 1e-9
assert abs(_calendar_days_per_period("us_equities_panel") - 1.5) < 1e-9
# Weekly cadence: 7 calendar days per bar
assert abs(_calendar_days_per_period("cme_futures") - 7.0) < 1e-9
assert abs(_calendar_days_per_period("sp500_equity_option_analytics") - 7.0) < 1e-9
# 8-hour cadence: ~0.333 day per bar (3 bars / 24h day)
assert abs(_calendar_days_per_period("crypto_perps_funding") - 1.0 / 3.0) < 1e-9
# 15-minute cadence: ~0.054 day per bar (1/26 trading day × 1.4 calendar buffer)
assert _calendar_days_per_period("nasdaq100_microstructure") < 0.1
# Monthly cadence: ~31 calendar days per bar
assert abs(_calendar_days_per_period("us_firm_characteristics") - 31.0) < 1e-9
def test_calendar_days_per_period_default_for_unknown_cs() -> None:
# Falls back to the daily 1.5× heuristic when no setup.yaml is present
# or the cadence token isn't in the lookup table.
assert _calendar_days_per_period("__nonexistent_cs__") == 1.5