Coverage for src/tinycta/hyper/_study.py: 100%

1"""Frozen Study result and Optuna-based hyperparameter optimisation.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from dataclasses import dataclass, field 

7from pathlib import Path 

8from typing import Any 

9 

10import optuna 

11from jquantstats import Portfolio 

12from loguru import logger 

13 

14 

15@dataclass(frozen=True) 

16class Study: 

17 """Frozen wrapper around a completed Optuna study.""" 

18 

19 best_params: dict[str, Any] 

20 best_value: float 

21 n_completed: int 

22 n_trials: int 

23 optuna_study: optuna.Study = field(repr=False) 

24 

25 def __str__(self) -> str: 

26 """Return a human-readable summary of the best trial.""" 

27 if self.n_completed == 0: 

28 return "No completed trials — all returned NaN Sharpe." 

29 lines = ["=== Best parameters ==="] 

30 for k, v in self.best_params.items(): 

31 lines.append(f" {k:<12} = {v}") 

32 lines.append(f" {'Sharpe':<12} = {self.best_value:.4f}") 

33 lines.append(f" {'Completed':<12} = {self.n_completed} / {self.n_trials} trials") 

34 return "\n".join(lines) 

35 

36 @classmethod 

37 def from_optuna(cls, s: optuna.Study) -> Study: 

38 """Wrap a completed optuna.Study in a frozen Study.""" 

39 n_completed = sum(1 for t in s.trials if t.state == optuna.trial.TrialState.COMPLETE) 

40 if n_completed == 0: 

41 best_params, best_value = {}, float("nan") 

42 else: 

43 best_params, best_value = s.best_params, s.best_value 

44 return cls( 

45 best_params=best_params, 

46 best_value=best_value, 

47 n_completed=n_completed, 

48 n_trials=len(s.trials), 

49 optuna_study=s, 

50 ) 

51 

52 def plot(self, output_dir: Path) -> None: 

53 """Write Optuna visualisation plots to output_dir (HTML, PNG if kaleido available).""" 

54 output_dir.mkdir(parents=True, exist_ok=True) 

55 figures = { 

56 "optuna_history": optuna.visualization.plot_optimization_history(self.optuna_study), 

57 "optuna_importance": optuna.visualization.plot_param_importances(self.optuna_study), 

58 "optuna_parallel": optuna.visualization.plot_parallel_coordinate(self.optuna_study), 

59 "optuna_contour": optuna.visualization.plot_contour(self.optuna_study), 

60 } 

61 for name, fig in figures.items(): 

62 fig.write_html(str(output_dir / f"{name}.html")) 

63 try: 

64 fig.write_image(str(output_dir / f"{name}.png"), scale=2) 

65 except (ValueError, ImportError) as exc: 

66 # PNG export needs the optional `kaleido` backend; skip if unavailable. 

67 logger.debug(f"Skipping PNG export for {name}: {exc}") 

68 

69 

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

71 """Compute Sharpe ratio, raising TrialPruned if the result is NaN or None.""" 

72 result = portfolio.stats.sharpe() 

73 sharpe = result["returns"] if isinstance(result, dict) else float(result) 

74 if sharpe is None or sharpe != sharpe: 

75 raise optuna.exceptions.TrialPruned() 

76 return sharpe 

77 

78 

79def _run_study( 

80 objective: Callable[[optuna.Trial], float], 

81 *, 

82 n_trials: int = 100, 

83 seed: int = 42, 

84 name: str | None = None, 

85) -> optuna.Study: 

86 """Create and run an Optuna study, returning the optuna.Study.""" 

87 optuna.logging.set_verbosity(optuna.logging.WARNING) 

88 s = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=seed), study_name=name) 

89 s.optimize(objective, n_trials=n_trials, show_progress_bar=False) 

90 return s 

91 

92 

93def _build_objective( 

94 suggest_portfolio_fn: Callable[[optuna.Trial], Portfolio], 

95) -> Callable[[optuna.Trial], float]: 

96 """Objective factory: wraps a portfolio-returning function with Sharpe scoring.""" 

97 

98 def objective(trial: optuna.Trial) -> float: 

99 """Call suggest_portfolio_fn and return the Sharpe ratio.""" 

100 return _sharpe(suggest_portfolio_fn(trial)) 

101 

102 return objective 

103 

104 

105def optimize( 

106 suggest_portfolio_fn: Callable[[optuna.Trial], Portfolio], 

107 n_trials: int = 100, 

108 seed: int = 42, 

109) -> Study: 

110 """Build objective, run study, log the summary and return a frozen Study.""" 

111 s = _run_study(_build_objective(suggest_portfolio_fn), n_trials=n_trials, seed=seed) 

112 study = Study.from_optuna(s) 

113 logger.info(str(study)) 

114 return study