Coverage for book/marimo/notebooks/optimize.py: 100%

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 sys

16import warnings

17from collections.abc import Callable

18from functools import cache

19from pathlib import Path

20from typing import Any, cast

22import numpy as np

23import optuna

24import polars as pl

25from jquantstats import Portfolio

27# ``preamble`` (the notebooks' shared loader) resolves because this script's own

28# directory is already on ``sys.path`` — put there by the interpreter for

29# ``python optimize.py`` and by the test ``conftest`` for the ``runpy`` path.

30from preamble import date_col, load_notebook, load_prices

31from tinycta.linalg import inv_a_norm, solve

32from tinycta.signal import shrink2id

34NOTEBOOK_DIR = Path(__file__).resolve().parent

36# Data/signal accessors are ``@cache``d, so importing this module reads no CSV and

37# runs no notebook. ``clip`` is a fixed winsorizing level, not a search dimension.

38CLIP = 4.2

41@cache

42def _prices() -> pl.DataFrame:

43 """Load the price frame once (date column plus one column per asset)."""

44 return load_prices(str(NOTEBOOK_DIR / "optimize.py"))

47@cache

48def _prices_only() -> pl.DataFrame:

49 """The price frame with the date column dropped."""

50 return _prices().drop(date_col)

53@cache

54def _assets() -> list[str]:

55 """Column names of the tradable assets (price columns, no date)."""

56 return _prices_only().columns

59@cache

60def _notebook(name: str) -> dict[str, Any]:

61 """Execute (and cache) a notebook's namespace (its ``f`` / ``dcc_correlation``)."""

62 return load_notebook(name)

65def _signal(notebook: str) -> Callable[..., Any]:

66 """The signal function ``f`` defined by an experiment notebook."""

67 return cast("Callable[..., Any]", _notebook(notebook)["f"])

70def _sharpe(portfolio: Portfolio) -> float:

71 """Return the portfolio's annualized Sharpe ratio (the optimization target)."""

72 value = portfolio.stats.sharpe()["returns"]

73 return float(value) if np.isfinite(value) else float("-inf")

76# Portfolio builders — one per experiment, mirroring each notebook cell.

77def build_exp1(*, fast: int, slow: int) -> Portfolio:

78 """CTA 1.0 — sign of the fast-minus-slow EWM crossover."""

79 f = _signal("Experiment1.py")

80 signals = _prices_only().select(f(pl.all(), fast=fast, slow=slow).fill_null(0.0) * 5e6)

81 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)

84def build_exp2(*, fast: int, slow: int, volatility: int) -> Portfolio:

85 """CTA 2.0 — volatility-scaled crossover."""

86 f = _signal("Experiment2.py")

87 signals = _prices_only().select(f(pl.all(), fast=fast, slow=slow, volatility=volatility).fill_null(0.0) * 1e5)

88 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)

91def build_exp3(*, fast: int, slow: int, vola: int, clip: float) -> Portfolio:

92 """CTA 3.0 — tanh oscillator on vol-adjusted prices, divided by volatility."""

93 f = _signal("Experiment3.py")

94 signals = _prices_only().select(

95 (f(pl.all(), fast=fast, slow=slow, vola=vola, clip=clip) * 1e5).fill_nan(0.0).fill_null(0.0)

96 )

97 return Portfolio.from_cash_position(prices=_prices(), cash_position=signals, aum=1e8)

100def build_exp4(*, fast: int, slow: int, vola: int, clip: float) -> Portfolio:

101 """CTA 4.0 — Euclidean risk-scaling across assets (optimization 1.0)."""

102 prices_only = _prices_only()

103 f = _signal("Experiment4.py")

104 mu_np = prices_only.select(f(pl.all(), fast=fast, slow=slow, vola=vola, clip=clip)).to_numpy()

105 volax_np = prices_only.select(pl.all().fill_nan(None).pct_change().ewm_std(com=vola, min_samples=vola)).to_numpy()

106 euclid_norm = np.sqrt(np.nansum(mu_np**2, axis=1, keepdims=True))

107 euclid_norm[euclid_norm == 0] = np.nan

108 risk_scaled_np = mu_np / euclid_norm

109 pos_np = np.nan_to_num(5e5 * risk_scaled_np / volax_np, nan=0.0)

110 return _portfolio_from_matrix(pos_np)

111

112

113def _day_position(matrix: np.ndarray, expected_mu: np.ndarray, expected_vo: np.ndarray) -> np.ndarray | None:

114 """Risk-scaled position for one day, or ``None`` for a degenerate day (singular / zero norm)."""

115 # inv_a_norm/solve resolve from module globals so tests can monkeypatch them.

116 try:

117 norm = inv_a_norm(expected_mu, matrix)

118 except ValueError:

119 # Singular correlation matrix on this day; skip the day, not the trial.

120 return None

121 if norm == 0 or np.isnan(norm):

122 return None

123 return cast("np.ndarray", np.nan_to_num(1e6 * (solve(matrix, expected_mu) / norm) / expected_vo, nan=0.0))

124

125

126def _solve_positions(

127 cor_3d: np.ndarray, mu: np.ndarray, vo: np.ndarray, prices_np: np.ndarray, *, shrinkage: float

128) -> np.ndarray:

129 """Per-day risk-parity positions from the shrunk DCC tensor; ill-conditioned days go flat."""

130 n_rows, n_assets = prices_np.shape

131 pos_matrix = np.zeros((n_rows, n_assets))

132 # Silence the numerical noise from the ill-conditioned days the search explores

133 # (matched by message to avoid importing cvx's internal warning classes).

134 with warnings.catch_warnings():

135 warnings.filterwarnings("ignore", message="Matrix condition number")

136 warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in sqrt")

137 for _n in range(n_rows):

138 _mask = np.isfinite(prices_np[_n])

139 if _mask.sum() == 0:

140 continue

141 _matrix = shrink2id(cor_3d[_n], lamb=shrinkage)[_mask, :][:, _mask]

142 _pos = _day_position(_matrix, np.nan_to_num(mu[_n][_mask]), np.nan_to_num(vo[_n][_mask]))

143 if _pos is not None:

144 pos_matrix[_n, _mask] = _pos

145 return pos_matrix

146

147

148def build_exp5(*, vola: int, clip: float, corr: int, shrinkage: float) -> Portfolio:

149 """CTA 5.0 — DCC correlation + shrinkage optimization; fast/slow fixed at 32/96 (optimization 2.0)."""

150 prices_only = _prices_only()

151 dcc_correlation = _notebook("Experiment5.py")["dcc_correlation"]

152 cor_3d = dcc_correlation(prices_only, vola=vola, clip=clip, corr=corr)

153 f = _signal("Experiment5.py")

154 mu = prices_only.select(f(pl.all(), fast=32, slow=96, vola=vola, clip=clip)).to_numpy()

155 vo = prices_only.select(pl.all().fill_nan(None).pct_change().ewm_std(com=vola, min_samples=int(vola))).to_numpy()

156 pos_matrix = _solve_positions(cor_3d, mu, vo, prices_only.to_numpy(), shrinkage=shrinkage)

157 return _portfolio_from_matrix(pos_matrix)

158

159

160def _portfolio_from_matrix(pos_np: np.ndarray) -> Portfolio:

161 """Wrap a (rows x assets) position matrix into a Portfolio with the date column."""

162 prices = _prices()

163 cash_position = pl.concat(

164 [prices.select(date_col), pl.from_numpy(pos_np, schema=dict.fromkeys(_assets(), pl.Float64))],

165 how="horizontal_extend",

166 )

167 return Portfolio.from_cash_position(prices=prices, cash_position=cash_position, aum=1e8)

168

169

170# Optuna search spaces. Ranges mirror the marimo sliders (4..192 step 4); ``clip`` is

171# fixed at ``CLIP``; ``slow`` is drawn strictly above ``fast`` so the oscillator holds.

172def _suggest_fast_slow(trial: optuna.Trial) -> tuple[int, int]:

173 """Sample a (fast, slow) pair with slow strictly above fast.

174

175 ``slow`` is drawn from ``fast + 4`` upwards rather than from a fixed low, so the

176 oscillator can never be handed an inverted pair. Feeding a fixed trial shows the

177 pair coming back in order:

178

179 >>> fast, slow = _suggest_fast_slow(optuna.trial.FixedTrial({"fast": 32, "slow": 96}))

180 >>> (fast, slow)

181 (32, 96)

182 >>> slow > fast

183 True

184 """

185 fast = trial.suggest_int("fast", 4, 96, step=4)

186 slow = trial.suggest_int("slow", fast + 4, 192, step=4)

187 return fast, slow

188

189

190def objective_exp1(trial: optuna.Trial) -> float:

191 """Sharpe of CTA 1.0 for a sampled (fast, slow)."""

192 fast, slow = _suggest_fast_slow(trial)

193 return _sharpe(build_exp1(fast=fast, slow=slow))

194

195

196def objective_exp2(trial: optuna.Trial) -> float:

197 """Sharpe of CTA 2.0 for a sampled (fast, slow, volatility)."""

198 fast, slow = _suggest_fast_slow(trial)

199 volatility = trial.suggest_int("volatility", 4, 192, step=4)

200 return _sharpe(build_exp2(fast=fast, slow=slow, volatility=volatility))

201

202

203def objective_exp3(trial: optuna.Trial) -> float:

204 """Sharpe of CTA 3.0 for a sampled (fast, slow, vola); clip is fixed."""

205 fast, slow = _suggest_fast_slow(trial)

206 vola = trial.suggest_int("vola", 4, 192, step=4)

207 return _sharpe(build_exp3(fast=fast, slow=slow, vola=vola, clip=CLIP))

208

209

210def objective_exp4(trial: optuna.Trial) -> float:

211 """Sharpe of CTA 4.0 for a sampled (fast, slow, vola); clip is fixed."""

212 fast, slow = _suggest_fast_slow(trial)

213 vola = trial.suggest_int("vola", 4, 192, step=4)

214 return _sharpe(build_exp4(fast=fast, slow=slow, vola=vola, clip=CLIP))

215

216

217def objective_exp5(trial: optuna.Trial) -> float:

218 """Sharpe of CTA 5.0 for a sampled (vola, corr, shrinkage); clip is fixed."""

219 vola = trial.suggest_int("vola", 4, 192, step=4)

220 corr = trial.suggest_int("corr", 50, 500, step=10)

221 shrinkage = trial.suggest_float("shrinkage", 0.0, 1.0, step=0.05)

222 return _sharpe(build_exp5(vola=vola, clip=CLIP, corr=corr, shrinkage=shrinkage))

223

224

225# Experiment registry.

226class Experiment:

227 """Bundles an objective with its baseline (notebook default) for reporting.

228

229 The five experiments are registered in :data:`EXPERIMENTS` under the same keys

230 ``--experiment`` accepts, and each carries the parameters its notebook's sliders

231 default to — the values :meth:`default_sharpe` scores to produce the baseline the

232 search is measured against:

233

234 >>> sorted(EXPERIMENTS)

235 ['1', '2', '3', '4', '5']

236 >>> EXPERIMENTS["1"].name

237 'Experiment 1'

238 >>> EXPERIMENTS["1"].default_params

239 {'fast': 32, 'slow': 96}

240 """

241

242 def __init__(

243 self, name: str, objective: Callable[..., Any], default_params: dict[str, Any], baseline: Callable[..., Any]

244 ):

245 """Store the experiment's name, Optuna objective, notebook defaults and builder."""

246 self.name = name

247 self.objective = objective

248 self.default_params = default_params

249 self._baseline = baseline

250

251 def default_sharpe(self) -> float:

252 """Sharpe of the strategy run with the notebook's default (slider) parameters."""

253 return _sharpe(self._baseline(**self.default_params))

254

255

256EXPERIMENTS: dict[str, Experiment] = {

257 "1": Experiment("Experiment 1", objective_exp1, {"fast": 32, "slow": 96}, build_exp1),

258 "2": Experiment("Experiment 2", objective_exp2, {"fast": 32, "slow": 96, "volatility": 32}, build_exp2),

259 "3": Experiment("Experiment 3", objective_exp3, {"fast": 32, "slow": 96, "vola": 32, "clip": 4.2}, build_exp3),

260 "4": Experiment("Experiment 4", objective_exp4, {"fast": 32, "slow": 96, "vola": 32, "clip": 4.2}, build_exp4),

261 "5": Experiment(

262 "Experiment 5", objective_exp5, {"vola": 32, "clip": 4.2, "corr": 200, "shrinkage": 0.5}, build_exp5

263 ),

264}

265

266# Exp 5 evaluates a per-row matrix solve (~3 s/trial); the rest are sub-second.

267DEFAULT_TRIALS: dict[str, int] = {"1": 200, "2": 200, "3": 150, "4": 150, "5": 40}

268

269

270def optimize(key: str, *, n_trials: int, seed: int) -> optuna.Study:

271 """Run an Optuna study for a single experiment and print a summary."""

272 experiment = EXPERIMENTS[key]

273 baseline = experiment.default_sharpe()

274 study = optuna.create_study(

275 direction="maximize",

276 study_name=experiment.name,

277 sampler=optuna.samplers.TPESampler(seed=seed),

278 )

279 study.optimize(experiment.objective, n_trials=n_trials, show_progress_bar=False)

280 print(f"\n{'═' * 60}")

281 print(f"{experiment.name} ({n_trials} trials)")

282 print(f"{'─' * 60}")

283 print(f" baseline Sharpe (defaults {experiment.default_params}): {baseline:.4f}")

284 print(f" best Sharpe: {study.best_value:.4f}")

285 print(f" best params: {study.best_params}")

286 improvement = study.best_value - baseline

287 if baseline == 0:

288 # A zero baseline Sharpe has no meaningful percentage; report absolute only.

289 print(f" improvement: {improvement:+.4f} (n/a — zero baseline)")

290 else:

291 print(f" improvement: {improvement:+.4f} ({improvement / abs(baseline):+.1%})")

292 print(f"{'═' * 60}")

293 return study

294

295

296def main(argv: list[str] | None = None) -> int:

297 """Parse command-line arguments and run the requested Optuna study/studies.

298

299 Returns the process exit code: ``0`` on success, ``1`` when the price data the

300 experiments read is missing. That case is caught here rather than allowed to

301 propagate because it is the one failure a user hits through no fault of the

302 search — a setup problem deserving a one-line message naming the file, not a

303 traceback out of the loader.

304 """

305 parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

306 parser.add_argument(

307 "--experiment",

308 "-e",

309 default="all",

310 choices=[*EXPERIMENTS.keys(), "all"],

311 help="Which experiment to optimize (1-5, or 'all'). Default: all.",

312 )

313 parser.add_argument(

314 "--trials",

315 "-n",

316 type=int,

317 default=None,

318 help="Number of Optuna trials. Default: per-experiment (see DEFAULT_TRIALS).",

319 )

320 parser.add_argument("--seed", "-s", type=int, default=42, help="Sampler seed for reproducibility.")

321 parser.add_argument("--verbose", action="store_true", help="Show Optuna's per-trial logging.")

322 args = parser.parse_args(argv)

323 if not args.verbose:

324 optuna.logging.set_verbosity(optuna.logging.WARNING)

325 keys = list(EXPERIMENTS.keys()) if args.experiment == "all" else [args.experiment]

326 try:

327 for key in keys:

328 n_trials = args.trials if args.trials is not None else DEFAULT_TRIALS[key]

329 optimize(key, n_trials=n_trials, seed=args.seed)

330 except FileNotFoundError as error:

331 print(f"error: {error}", file=sys.stderr)

332 return 1

333 return 0

334

335

336if __name__ == "__main__":

337 sys.exit(main())