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 Example: 

20 >>> import optuna 

21 >>> from tinycta.hyper import Study 

22 >>> optuna.logging.set_verbosity(optuna.logging.WARNING) 

23 >>> s = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=0)) 

24 >>> s.optimize(lambda trial: trial.suggest_float("x", 0.0, 1.0), n_trials=5) 

25 >>> study = Study.from_optuna(s) 

26 >>> study.n_trials, study.n_completed 

27 (5, 5) 

28 >>> sorted(study.best_params) 

29 ['x'] 

30 

31 ``str`` renders the best trial as a report block: 

32 

33 >>> print(study) # doctest: +ELLIPSIS 

34 === Best parameters === 

35 x = 0... 

36 Sharpe = 0... 

37 Completed = 5 / 5 trials 

38 

39 A study in which every trial was pruned (each scored a NaN Sharpe) is not an 

40 error — it reports no best parameters and a NaN best value: 

41 

42 >>> pruned = optuna.create_study(direction="maximize") 

43 >>> pruned.optimize( 

44 ... lambda trial: (_ for _ in ()).throw(optuna.exceptions.TrialPruned()), n_trials=2 

45 ... ) 

46 >>> empty = Study.from_optuna(pruned) 

47 >>> empty.n_completed, empty.best_params 

48 (0, {}) 

49 >>> print(empty) 

50 No completed trials — all returned NaN Sharpe. 

51 """ 

52 

53 best_params: dict[str, Any] 

54 best_value: float 

55 n_completed: int 

56 n_trials: int 

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

58 

59 def __str__(self) -> str: 

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

61 if self.n_completed == 0: 

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

63 lines = ["=== Best parameters ==="] 

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

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

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

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

68 return "\n".join(lines) 

69 

70 @classmethod 

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

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

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

74 if n_completed == 0: 

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

76 else: 

77 best_params, best_value = s.best_params, s.best_value 

78 return cls( 

79 best_params=best_params, 

80 best_value=best_value, 

81 n_completed=n_completed, 

82 n_trials=len(s.trials), 

83 optuna_study=s, 

84 ) 

85 

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

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

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

89 figures = { 

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

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

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

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

94 } 

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

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

97 try: 

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

99 except (ValueError, ImportError) as exc: 

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

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

102 

103 

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

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

106 result = portfolio.stats.sharpe() 

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

108 if sharpe is None or sharpe != sharpe: 

109 raise optuna.exceptions.TrialPruned() 

110 return sharpe 

111 

112 

113def _run_study( 

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

115 *, 

116 n_trials: int = 100, 

117 seed: int = 42, 

118 name: str | None = None, 

119) -> optuna.Study: 

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

121 optuna.logging.set_verbosity(optuna.logging.WARNING) 

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

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

124 return s 

125 

126 

127def _build_objective( 

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

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

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

131 

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

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

134 return _sharpe(suggest_portfolio_fn(trial)) 

135 

136 return objective 

137 

138 

139def optimize( 

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

141 n_trials: int = 100, 

142 seed: int = 42, 

143) -> Study: 

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

145 

146 ``suggest_portfolio_fn`` draws its parameters from the trial and returns a 

147 portfolio; the trial is then scored by that portfolio's Sharpe ratio, which 

148 the study maximises. A trial whose Sharpe is NaN is pruned rather than fatal. 

149 

150 Example: 

151 >>> from types import SimpleNamespace 

152 >>> from tinycta.hyper import optimize 

153 

154 Any object exposing ``.stats.sharpe()`` works here; in real use that is a 

155 ``jquantstats`` ``Portfolio`` built from the strategy's returns: 

156 

157 >>> def portfolio(sharpe): 

158 ... return SimpleNamespace(stats=SimpleNamespace(sharpe=lambda: sharpe)) 

159 

160 The objective is maximised, so the best value is the highest Sharpe seen — 

161 here the reward peaks where ``fast`` is largest: 

162 

163 >>> study = optimize(lambda trial: portfolio(trial.suggest_int("fast", 1, 8)), n_trials=12) 

164 >>> study.best_value 

165 8.0 

166 >>> study.best_params 

167 {'fast': 8} 

168 

169 Runs are seeded, so the same objective and seed reproduce the same result: 

170 

171 >>> repeat = optimize(lambda trial: portfolio(trial.suggest_int("fast", 1, 8)), n_trials=12) 

172 >>> repeat.best_params == study.best_params 

173 True 

174 """ 

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

176 study = Study.from_optuna(s) 

177 logger.info(str(study)) 

178 return study