Coverage for book/marimo/notebooks/optimize.py: 100%
150 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-31 10:05 +0000
1"""Optuna parameter optimization for the CTA experiments.
3Each ``ExperimentN.py`` marimo notebook defines a signal function ``f(...)`` and
4builds a :class:`jquantstats.Portfolio` whose Sharpe ratio its sliders tune by
5hand. This module replays the same portfolio construction (reusing each notebook's
6``f`` via :func:`runpy.run_path` and the TinyCTA API) and hands the search to
7Optuna, maximizing the Sharpe over each experiment's parameter space. Strategy
8logic thus lives once in the notebooks; only the search space lives here. Run it as
9``python optimize.py --experiment {1..5|all} [--trials N]``.
10"""
12from __future__ import annotations
14import argparse
15import warnings
16from collections.abc import Callable
17from functools import cache
18from pathlib import Path
19from typing import Any, cast
21import numpy as np
22import optuna
23import polars as pl
24from jquantstats import Portfolio
26# ``preamble`` (the notebooks' shared loader) resolves because this script's own
27# directory is already on ``sys.path`` — put there by the interpreter for
28# ``python optimize.py`` and by the test ``conftest`` for the ``runpy`` path.
29from preamble import date_col, load_notebook, load_prices
30from tinycta.linalg import inv_a_norm, solve
31from tinycta.signal import shrink2id
33NOTEBOOK_DIR = Path(__file__).resolve().parent
35# Data/signal accessors are ``@cache``d, so importing this module reads no CSV and
36# runs no notebook. ``clip`` is a fixed winsorizing level, not a search dimension.
37CLIP = 4.2
40@cache
41def _prices() -> pl.DataFrame:
42 """Load the price frame once (date column plus one column per asset)."""
43 return load_prices(str(NOTEBOOK_DIR / "optimize.py"))
46@cache
47def _prices_only() -> pl.DataFrame:
48 """The price frame with the date column dropped."""
49 return _prices().drop(date_col)
52@cache
53def _assets() -> list[str]:
54 """Column names of the tradable assets (price columns, no date)."""
55 return _prices_only().columns
58@cache
59def _notebook(name: str) -> dict[str, Any]:
60 """Execute (and cache) a notebook's namespace (its ``f`` / ``dcc_correlation``)."""
61 return load_notebook(name)
64def _signal(notebook: str) -> Callable[..., Any]:
65 """The signal function ``f`` defined by an experiment notebook."""
66 return cast("Callable[..., Any]", _notebook(notebook)["f"])
69def _sharpe(portfolio: Portfolio) -> float:
70 """Return the portfolio's annualized Sharpe ratio (the optimization target)."""
71 value = portfolio.stats.sharpe()["returns"]
72 return float(value) if np.isfinite(value) else float("-inf")
75# Portfolio builders — one per experiment, mirroring each notebook cell.
76def build_exp1(*, fast: int, slow: int) -> Portfolio:
77 """CTA 1.0 — sign of the fast-minus-slow EWM crossover."""
78 f = _signal("Experiment1.py")
79 signals = _prices_only().select(f(pl.all(), fast=fast, slow=slow).fill_null(0.0) * 5e6)
80 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)
83def build_exp2(*, fast: int, slow: int, volatility: int) -> Portfolio:
84 """CTA 2.0 — volatility-scaled crossover."""
85 f = _signal("Experiment2.py")
86 signals = _prices_only().select(f(pl.all(), fast=fast, slow=slow, volatility=volatility).fill_null(0.0) * 1e5)
87 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)
90def build_exp3(*, fast: int, slow: int, vola: int, clip: float) -> Portfolio:
91 """CTA 3.0 — tanh oscillator on vol-adjusted prices, divided by volatility."""
92 f = _signal("Experiment3.py")
93 signals = _prices_only().select(
94 (f(pl.all(), fast=fast, slow=slow, vola=vola, clip=clip) * 1e5).fill_nan(0.0).fill_null(0.0)
95 )
96 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)
99def build_exp4(*, fast: int, slow: int, vola: int, clip: float) -> Portfolio:
100 """CTA 4.0 — Euclidean risk-scaling across assets (optimization 1.0)."""
101 prices_only = _prices_only()
102 f = _signal("Experiment4.py")
103 mu_np = prices_only.select(f(pl.all(), fast=fast, slow=slow, vola=vola, clip=clip)).to_numpy()
104 volax_np = prices_only.select(pl.all().fill_nan(None).pct_change().ewm_std(com=vola, min_samples=vola)).to_numpy()
105 euclid_norm = np.sqrt(np.nansum(mu_np**2, axis=1, keepdims=True))
106 euclid_norm[euclid_norm == 0] = np.nan
107 risk_scaled_np = mu_np / euclid_norm
108 pos_np = np.nan_to_num(5e5 * risk_scaled_np / volax_np, nan=0.0)
109 return _portfolio_from_matrix(pos_np)
112def _day_position(matrix: np.ndarray, expected_mu: np.ndarray, expected_vo: np.ndarray) -> np.ndarray | None:
113 """Risk-scaled position for one day, or ``None`` for a degenerate day (singular / zero norm)."""
114 # inv_a_norm/solve resolve from module globals so tests can monkeypatch them.
115 try:
116 norm = inv_a_norm(expected_mu, matrix)
117 except ValueError:
118 # Singular correlation matrix on this day; skip the day, not the trial.
119 return None
120 if norm == 0 or np.isnan(norm):
121 return None
122 return cast("np.ndarray", np.nan_to_num(1e6 * (solve(matrix, expected_mu) / norm) / expected_vo, nan=0.0))
125def _solve_positions(
126 cor_3d: np.ndarray, mu: np.ndarray, vo: np.ndarray, prices_np: np.ndarray, *, shrinkage: float
127) -> np.ndarray:
128 """Per-day risk-parity positions from the shrunk DCC tensor; ill-conditioned days go flat."""
129 n_rows, n_assets = prices_np.shape
130 pos_matrix = np.zeros((n_rows, n_assets))
131 # Silence the numerical noise from the ill-conditioned days the search explores
132 # (matched by message to avoid importing cvx's internal warning classes).
133 with warnings.catch_warnings():
134 warnings.filterwarnings("ignore", message="Matrix condition number")
135 warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in sqrt")
136 for _n in range(n_rows):
137 _mask = np.isfinite(prices_np[_n])
138 if _mask.sum() == 0:
139 continue
140 _matrix = shrink2id(cor_3d[_n], lamb=shrinkage)[_mask, :][:, _mask]
141 _pos = _day_position(_matrix, np.nan_to_num(mu[_n][_mask]), np.nan_to_num(vo[_n][_mask]))
142 if _pos is not None:
143 pos_matrix[_n, _mask] = _pos
144 return pos_matrix
147def build_exp5(*, vola: int, clip: float, corr: int, shrinkage: float) -> Portfolio:
148 """CTA 5.0 — DCC correlation + shrinkage optimization; fast/slow fixed at 32/96 (optimization 2.0)."""
149 prices_only = _prices_only()
150 dcc_correlation = _notebook("Experiment5.py")["dcc_correlation"]
151 cor_3d = dcc_correlation(prices_only, vola=vola, clip=clip, corr=corr)
152 f = _signal("Experiment5.py")
153 mu = prices_only.select(f(pl.all(), fast=32, slow=96, vola=vola, clip=clip)).to_numpy()
154 vo = prices_only.select(pl.all().fill_nan(None).pct_change().ewm_std(com=vola, min_samples=int(vola))).to_numpy()
155 pos_matrix = _solve_positions(cor_3d, mu, vo, prices_only.to_numpy(), shrinkage=shrinkage)
156 return _portfolio_from_matrix(pos_matrix)
159def _portfolio_from_matrix(pos_np: np.ndarray) -> Portfolio:
160 """Wrap a (rows x assets) position matrix into a Portfolio with the date column."""
161 prices = _prices()
162 cash_position = pl.concat(
163 [prices.select(date_col), pl.from_numpy(pos_np, schema=dict.fromkeys(_assets(), pl.Float64))],
164 how="horizontal_extend",
165 )
166 return Portfolio.from_cash_position(prices=prices, cash_position=cash_position, aum=1e8)
169# Optuna search spaces. Ranges mirror the marimo sliders (4..192 step 4); ``clip`` is
170# fixed at ``CLIP``; ``slow`` is drawn strictly above ``fast`` so the oscillator holds.
171def _suggest_fast_slow(trial: optuna.Trial) -> tuple[int, int]:
172 """Sample a (fast, slow) pair with slow strictly above fast."""
173 fast = trial.suggest_int("fast", 4, 96, step=4)
174 slow = trial.suggest_int("slow", fast + 4, 192, step=4)
175 return fast, slow
178def objective_exp1(trial: optuna.Trial) -> float:
179 """Sharpe of CTA 1.0 for a sampled (fast, slow)."""
180 fast, slow = _suggest_fast_slow(trial)
181 return _sharpe(build_exp1(fast=fast, slow=slow))
184def objective_exp2(trial: optuna.Trial) -> float:
185 """Sharpe of CTA 2.0 for a sampled (fast, slow, volatility)."""
186 fast, slow = _suggest_fast_slow(trial)
187 volatility = trial.suggest_int("volatility", 4, 192, step=4)
188 return _sharpe(build_exp2(fast=fast, slow=slow, volatility=volatility))
191def objective_exp3(trial: optuna.Trial) -> float:
192 """Sharpe of CTA 3.0 for a sampled (fast, slow, vola); clip is fixed."""
193 fast, slow = _suggest_fast_slow(trial)
194 vola = trial.suggest_int("vola", 4, 192, step=4)
195 return _sharpe(build_exp3(fast=fast, slow=slow, vola=vola, clip=CLIP))
198def objective_exp4(trial: optuna.Trial) -> float:
199 """Sharpe of CTA 4.0 for a sampled (fast, slow, vola); clip is fixed."""
200 fast, slow = _suggest_fast_slow(trial)
201 vola = trial.suggest_int("vola", 4, 192, step=4)
202 return _sharpe(build_exp4(fast=fast, slow=slow, vola=vola, clip=CLIP))
205def objective_exp5(trial: optuna.Trial) -> float:
206 """Sharpe of CTA 5.0 for a sampled (vola, corr, shrinkage); clip is fixed."""
207 vola = trial.suggest_int("vola", 4, 192, step=4)
208 corr = trial.suggest_int("corr", 50, 500, step=10)
209 shrinkage = trial.suggest_float("shrinkage", 0.0, 1.0, step=0.05)
210 return _sharpe(build_exp5(vola=vola, clip=CLIP, corr=corr, shrinkage=shrinkage))
213# Experiment registry.
214class Experiment:
215 """Bundles an objective with its baseline (notebook default) for reporting."""
217 def __init__(
218 self, name: str, objective: Callable[..., Any], default_params: dict[str, Any], baseline: Callable[..., Any]
219 ):
220 """Store the experiment's name, Optuna objective, notebook defaults and builder."""
221 self.name = name
222 self.objective = objective
223 self.default_params = default_params
224 self._baseline = baseline
226 def default_sharpe(self) -> float:
227 """Sharpe of the strategy run with the notebook's default (slider) parameters."""
228 return _sharpe(self._baseline(**self.default_params))
231EXPERIMENTS: dict[str, Experiment] = {
232 "1": Experiment("Experiment 1", objective_exp1, {"fast": 32, "slow": 96}, build_exp1),
233 "2": Experiment("Experiment 2", objective_exp2, {"fast": 32, "slow": 96, "volatility": 32}, build_exp2),
234 "3": Experiment("Experiment 3", objective_exp3, {"fast": 32, "slow": 96, "vola": 32, "clip": 4.2}, build_exp3),
235 "4": Experiment("Experiment 4", objective_exp4, {"fast": 32, "slow": 96, "vola": 32, "clip": 4.2}, build_exp4),
236 "5": Experiment(
237 "Experiment 5", objective_exp5, {"vola": 32, "clip": 4.2, "corr": 200, "shrinkage": 0.5}, build_exp5
238 ),
239}
241# Exp 5 evaluates a per-row matrix solve (~3 s/trial); the rest are sub-second.
242DEFAULT_TRIALS: dict[str, int] = {"1": 200, "2": 200, "3": 150, "4": 150, "5": 40}
245def optimize(key: str, *, n_trials: int, seed: int) -> optuna.Study:
246 """Run an Optuna study for a single experiment and print a summary."""
247 experiment = EXPERIMENTS[key]
248 baseline = experiment.default_sharpe()
249 study = optuna.create_study(
250 direction="maximize",
251 study_name=experiment.name,
252 sampler=optuna.samplers.TPESampler(seed=seed),
253 )
254 study.optimize(experiment.objective, n_trials=n_trials, show_progress_bar=False)
255 print(f"\n{'═' * 60}")
256 print(f"{experiment.name} ({n_trials} trials)")
257 print(f"{'─' * 60}")
258 print(f" baseline Sharpe (defaults {experiment.default_params}): {baseline:.4f}")
259 print(f" best Sharpe: {study.best_value:.4f}")
260 print(f" best params: {study.best_params}")
261 improvement = study.best_value - baseline
262 if baseline == 0:
263 # A zero baseline Sharpe has no meaningful percentage; report absolute only.
264 print(f" improvement: {improvement:+.4f} (n/a — zero baseline)")
265 else:
266 print(f" improvement: {improvement:+.4f} ({improvement / abs(baseline):+.1%})")
267 print(f"{'═' * 60}")
268 return study
271def main(argv: list[str] | None = None) -> None:
272 """Parse command-line arguments and run the requested Optuna study/studies."""
273 parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
274 parser.add_argument(
275 "--experiment",
276 "-e",
277 default="all",
278 choices=[*EXPERIMENTS.keys(), "all"],
279 help="Which experiment to optimize (1-5, or 'all'). Default: all.",
280 )
281 parser.add_argument(
282 "--trials",
283 "-n",
284 type=int,
285 default=None,
286 help="Number of Optuna trials. Default: per-experiment (see DEFAULT_TRIALS).",
287 )
288 parser.add_argument("--seed", "-s", type=int, default=42, help="Sampler seed for reproducibility.")
289 parser.add_argument("--verbose", action="store_true", help="Show Optuna's per-trial logging.")
290 args = parser.parse_args(argv)
291 if not args.verbose:
292 optuna.logging.set_verbosity(optuna.logging.WARNING)
293 keys = list(EXPERIMENTS.keys()) if args.experiment == "all" else [args.experiment]
294 for key in keys:
295 n_trials = args.trials if args.trials is not None else DEFAULT_TRIALS[key]
296 optimize(key, n_trials=n_trials, seed=args.seed)
299if __name__ == "__main__":
300 main()