Coverage for src/local_deep_research/benchmarks/optimization/optuna_optimizer.py: 79%

1""" 

2Optuna-based parameter optimizer for Local Deep Research. 

3 

4This module provides the core optimization functionality using Optuna 

5to find optimal parameters for the research system, balancing quality 

6and performance metrics. 

7""" 

8 

9import os 

10from pathlib import Path 

11import time 

12from datetime import datetime, UTC 

13from functools import partial 

14from typing import Any, Callable, Dict, List, Optional, Tuple 

15 

16import joblib 

17import numpy as np 

18import optuna 

19 

20# (matplotlib / optuna.visualization are imported lazily — see 

21# _ensure_plotting_loaded below and #4431.) 

22 

23from local_deep_research.benchmarks.efficiency.speed_profiler import ( 

24 SpeedProfiler, 

25) 

26from local_deep_research.security import sanitize_data 

27from loguru import logger 

28 

29from local_deep_research.benchmarks.evaluators import ( 

30 CompositeBenchmarkEvaluator, 

31) 

32 

33# Import benchmark evaluator components 

34 

35# Visualization libraries (matplotlib + optuna.visualization) are imported 

36# LAZILY — see _ensure_plotting_loaded below and #4431. find_spec only 

37# probes availability; it does NOT execute/import the module, so importing 

38# this file (and therefore `local_deep_research.benchmarks`) never pays 

39# matplotlib's ~60s cold-import cost. 

40import importlib.util 

41 

42PLOTTING_AVAILABLE = importlib.util.find_spec("matplotlib") is not None 

43if not PLOTTING_AVAILABLE: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 logger.warning("Matplotlib not available, visualization will be limited") 

45 

46 

47# Module-level placeholders so tests can @patch these names and so the 

48# loader below can fill them in place. None until first real visualization. 

49plt = None 

50Line2D = None 

51plot_contour = None 

52plot_optimization_history = None 

53plot_param_importances = None 

54plot_slice = None 

55 

56 

57def _ensure_plotting_loaded(): 

58 """Import matplotlib + optuna.visualization into module globals on first 

59 use. 

60 

61 Deferred from module load on purpose: matplotlib's import is heavy and, 

62 under the 2-core CI runner's GIL/CPU starvation, stretched to ~60s while 

63 holding the import lock. Because benchmarks/__init__.py imports this 

64 module, a module-level matplotlib import froze the whole werkzeug request 

65 pipeline the first time any request touched `local_deep_research.benchmarks` 

66 — the flaky UI-shard navigation timeouts (#4431). The 

67 _create_*_visualizations methods call this after their PLOTTING_AVAILABLE 

68 guard; visualization only happens in benchmark-optimization runs, never 

69 on a request path. Early-returns if already loaded (or a test has patched 

70 plt) so it never clobbers mocks. 

71 """ 

72 global plt, Line2D 

73 global plot_contour, plot_optimization_history 

74 global plot_param_importances, plot_slice 

75 if plt is not None: 

76 return 

77 import matplotlib.pyplot as plt 

78 from matplotlib.lines import Line2D 

79 from optuna.visualization import ( 

80 plot_contour, 

81 plot_optimization_history, 

82 plot_param_importances, 

83 plot_slice, 

84 ) 

85 

86 

87class OptunaOptimizer: 

88 """ 

89 Optimize parameters for Local Deep Research using Optuna. 

90 

91 This class provides functionality to: 

92 1. Define search spaces for parameter optimization 

93 2. Evaluate parameter combinations using objective functions 

94 3. Find optimal parameters via Optuna 

95 4. Visualize and analyze optimization results 

96 """ 

97 

98 def __init__( 

99 self, 

100 base_query: str, 

101 output_dir: str = "optimization_results", 

102 model_name: Optional[str] = None, 

103 provider: Optional[str] = None, 

104 search_tool: Optional[str] = None, 

105 temperature: float = 0.7, 

106 n_trials: int = 30, 

107 timeout: Optional[int] = None, 

108 n_jobs: int = 1, 

109 study_name: Optional[str] = None, 

110 optimization_metrics: Optional[List[str]] = None, 

111 metric_weights: Optional[Dict[str, float]] = None, 

112 progress_callback: Optional[Callable[[int, int, Dict], None]] = None, 

113 benchmark_weights: Optional[Dict[str, float]] = None, 

114 ): 

115 """ 

116 Initialize the optimizer. 

117 

118 Args: 

119 base_query: The research query to use for all experiments 

120 output_dir: Directory to save optimization results 

121 model_name: Name of the LLM model to use 

122 provider: LLM provider 

123 search_tool: Search engine to use 

124 temperature: LLM temperature 

125 n_trials: Number of parameter combinations to try 

126 timeout: Maximum seconds to run optimization (None for no limit) 

127 n_jobs: Number of parallel jobs for optimization 

128 study_name: Name of the Optuna study 

129 optimization_metrics: List of metrics to optimize (default: ["quality", "speed"]) 

130 metric_weights: Dictionary of weights for each metric (e.g., {"quality": 0.6, "speed": 0.4}) 

131 progress_callback: Optional callback for progress updates 

132 benchmark_weights: Dictionary mapping benchmark types to weights 

133 (e.g., {"simpleqa": 0.6, "browsecomp": 0.4}) 

134 If None, only SimpleQA is used with weight 1.0 

135 """ 

136 self.base_query = base_query 

137 self.output_dir = output_dir 

138 self.model_name = model_name 

139 self.provider = provider 

140 self.search_tool = search_tool 

141 self.temperature = temperature 

142 self.n_trials = n_trials 

143 self.timeout = timeout 

144 self.n_jobs = n_jobs 

145 self.optimization_metrics = optimization_metrics or ["quality", "speed"] 

146 self.metric_weights = metric_weights or {"quality": 0.6, "speed": 0.4} 

147 self.progress_callback = progress_callback 

148 

149 # Initialize benchmark evaluator with weights 

150 self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0} 

151 self.benchmark_evaluator = CompositeBenchmarkEvaluator( 

152 self.benchmark_weights 

153 ) 

154 

155 # Normalize weights to sum to 1.0 

156 total_weight = sum(self.metric_weights.values()) 

157 if total_weight > 0: 

158 self.metric_weights = { 

159 k: v / total_weight for k, v in self.metric_weights.items() 

160 } 

161 

162 # Generate a unique study name if not provided 

163 timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") 

164 self.study_name = study_name or f"ldr_opt_{timestamp}" 

165 

166 # Create output directory 

167 os.makedirs(output_dir, exist_ok=True) 

168 

169 # Store the trial history for analysis 

170 self.trials_history: List[Dict[str, Any]] = [] 

171 

172 # Storage for the best parameters and study 

173 self.best_params: Optional[Dict[str, Any]] = None 

174 self.study: Optional[optuna.Study] = None 

175 

176 def optimize( 

177 self, param_space: Optional[Dict[str, Any]] = None 

178 ) -> Tuple[Dict[str, Any], float]: 

179 """ 

180 Run the optimization process using Optuna. 

181 

182 Args: 

183 param_space: Dictionary defining parameter search spaces 

184 (if None, use default spaces) 

185 

186 Returns: 

187 Tuple containing (best_parameters, best_score) 

188 """ 

189 param_space = param_space or self._get_default_param_space() 

190 

191 # Create a study object 

192 storage_name = f"sqlite:///{self.output_dir}/{self.study_name}.db" 

193 self.study = optuna.create_study( 

194 study_name=self.study_name, 

195 storage=storage_name, 

196 load_if_exists=True, 

197 direction="maximize", 

198 sampler=optuna.samplers.TPESampler(seed=42), 

199 ) 

200 

201 # Create partial function with param_space 

202 objective = partial(self._objective, param_space=param_space) 

203 

204 # Log optimization start 

205 logger.info( 

206 f"Starting optimization with {self.n_trials} trials, {self.n_jobs} parallel jobs" 

207 ) 

208 logger.info(f"Parameter space: {param_space}") 

209 logger.info(f"Metric weights: {self.metric_weights}") 

210 logger.info(f"Benchmark weights: {self.benchmark_weights}") 

211 

212 # Initialize progress tracking 

213 if self.progress_callback: 

214 self.progress_callback( 

215 0, 

216 self.n_trials, 

217 { 

218 "status": "starting", 

219 "stage": "initialization", 

220 "trials_completed": 0, 

221 "total_trials": self.n_trials, 

222 }, 

223 ) 

224 

225 try: 

226 # Run optimization 

227 self.study.optimize( 

228 objective, 

229 n_trials=self.n_trials, 

230 timeout=self.timeout, 

231 n_jobs=self.n_jobs, 

232 callbacks=[self._optimization_callback], 

233 show_progress_bar=True, 

234 ) 

235 

236 # Store best parameters 

237 if self.study is None: 237 ↛ 238line 237 didn't jump to line 238 because the condition on line 237 was never true

238 raise RuntimeError("Study was not created") 

239 _completed_study = self.study 

240 self.best_params = _completed_study.best_params 

241 

242 # Save the results 

243 self._save_results() 

244 

245 # Create visualizations 

246 self._create_visualizations() 

247 

248 if self.best_params is None: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true

249 raise RuntimeError("No best parameters found") 

250 logger.info( 

251 f"Optimization complete. Best parameters: {self.best_params}" 

252 ) 

253 logger.info(f"Best value: {_completed_study.best_value}") 

254 

255 # Report completion 

256 if self.progress_callback: 

257 self.progress_callback( 

258 self.n_trials, 

259 self.n_trials, 

260 { 

261 "status": "completed", 

262 "stage": "finished", 

263 "trials_completed": len(_completed_study.trials), 

264 "total_trials": self.n_trials, 

265 "best_params": self.best_params, 

266 "best_value": _completed_study.best_value, 

267 }, 

268 ) 

269 

270 return self.best_params, _completed_study.best_value 

271 

272 except KeyboardInterrupt: 

273 logger.info("Optimization interrupted by user") 

274 # Still save what we have 

275 self._save_results() 

276 self._create_visualizations() 

277 

278 if self.study is None: 278 ↛ 279line 278 didn't jump to line 279 because the condition on line 278 was never true

279 raise RuntimeError("Study was not created") 

280 _interrupted_study = self.study 

281 # Report interruption 

282 if self.progress_callback: 

283 self.progress_callback( 

284 len(_interrupted_study.trials), 

285 self.n_trials, 

286 { 

287 "status": "interrupted", 

288 "stage": "interrupted", 

289 "trials_completed": len(_interrupted_study.trials), 

290 "total_trials": self.n_trials, 

291 "best_params": _interrupted_study.best_params, 

292 "best_value": _interrupted_study.best_value, 

293 }, 

294 ) 

295 

296 return _interrupted_study.best_params, _interrupted_study.best_value 

297 

298 def _get_default_param_space(self) -> Dict[str, Any]: 

299 """ 

300 Get default parameter search space. 

301 

302 Returns: 

303 Dictionary defining the default parameter search spaces 

304 """ 

305 return { 

306 "iterations": { 

307 "type": "int", 

308 "low": 1, 

309 "high": 5, 

310 "step": 1, 

311 }, 

312 "questions_per_iteration": { 

313 "type": "int", 

314 "low": 1, 

315 "high": 5, 

316 "step": 1, 

317 }, 

318 "search_strategy": { 

319 "type": "categorical", 

320 "choices": [ 

321 "source-based", 

322 "focused-iteration", 

323 "focused-iteration-standard", 

324 "topic-organization", 

325 ], 

326 }, 

327 "max_results": { 

328 "type": "int", 

329 "low": 10, 

330 "high": 100, 

331 "step": 10, 

332 }, 

333 } 

334 

335 def _objective( 

336 self, trial: optuna.Trial, param_space: Dict[str, Any] 

337 ) -> float: 

338 """ 

339 Objective function for Optuna optimization. 

340 

341 Args: 

342 trial: Optuna trial object 

343 param_space: Dictionary defining parameter search spaces 

344 

345 Returns: 

346 Score to maximize 

347 """ 

348 # Generate parameters for this trial 

349 params: Dict[str, Any] = {} 

350 for param_name, param_config in param_space.items(): 

351 param_type = param_config["type"] 

352 

353 if param_type == "int": 

354 params[param_name] = trial.suggest_int( 

355 param_name, 

356 param_config["low"], 

357 param_config["high"], 

358 step=param_config.get("step", 1), 

359 ) 

360 elif param_type == "float": 

361 params[param_name] = trial.suggest_float( 

362 param_name, 

363 param_config["low"], 

364 param_config["high"], 

365 step=param_config.get("step"), 

366 log=param_config.get("log", False), 

367 ) 

368 elif param_type == "categorical": 

369 params[param_name] = trial.suggest_categorical( 

370 param_name, param_config["choices"] 

371 ) 

372 

373 # Log the trial parameters 

374 logger.info(f"Trial {trial.number}: {params}") 

375 

376 # Update progress callback if available 

377 if self.progress_callback: 

378 self.progress_callback( 

379 trial.number, 

380 self.n_trials, 

381 { 

382 "status": "running", 

383 "stage": "trial_started", 

384 "trial_number": trial.number, 

385 "params": params, 

386 "trials_completed": trial.number, 

387 "total_trials": self.n_trials, 

388 }, 

389 ) 

390 

391 # Run an experiment with these parameters 

392 try: 

393 start_time = time.time() 

394 result = self._run_experiment(params) 

395 duration = time.time() - start_time 

396 

397 # Store details about the trial 

398 trial_info = { 

399 "trial_number": trial.number, 

400 "params": params, 

401 "result": result, 

402 "score": result.get("score", 0), 

403 "duration": duration, 

404 "timestamp": datetime.now(UTC).isoformat(), 

405 } 

406 self.trials_history.append(trial_info) 

407 

408 # Update callback with results 

409 if self.progress_callback: 

410 self.progress_callback( 

411 trial.number, 

412 self.n_trials, 

413 { 

414 "status": "completed", 

415 "stage": "trial_completed", 

416 "trial_number": trial.number, 

417 "params": params, 

418 "score": result.get("score", 0), 

419 "trials_completed": trial.number + 1, 

420 "total_trials": self.n_trials, 

421 }, 

422 ) 

423 

424 logger.info( 

425 f"Trial {trial.number} completed: {params}, score: {result['score']:.4f}" 

426 ) 

427 

428 return float(result["score"]) 

429 except Exception as e: 

430 logger.exception(f"Error in trial {trial.number}") 

431 

432 # Update callback with error 

433 if self.progress_callback: 

434 self.progress_callback( 

435 trial.number, 

436 self.n_trials, 

437 { 

438 "status": "error", 

439 "stage": "trial_error", 

440 "trial_number": trial.number, 

441 "params": params, 

442 "error": str(e), 

443 "trials_completed": trial.number, 

444 "total_trials": self.n_trials, 

445 }, 

446 ) 

447 

448 return float("-inf") # Return a very low score for failed trials 

449 

450 def _run_experiment(self, params: Dict[str, Any]) -> Dict[str, Any]: 

451 """ 

452 Run a single experiment with the given parameters. 

453 

454 Args: 

455 params: Dictionary of parameters to test 

456 

457 Returns: 

458 Results dictionary with metrics and score 

459 """ 

460 # Extract parameters 

461 iterations = params.get("iterations", 2) 

462 questions_per_iteration = params.get("questions_per_iteration", 2) 

463 search_strategy = params.get("search_strategy", "source-based") 

464 max_results = params.get("max_results", 50) 

465 

466 # Initialize profiling tools 

467 speed_profiler = SpeedProfiler() 

468 

469 # Start profiling 

470 speed_profiler.start() 

471 

472 try: 

473 # Create system configuration 

474 system_config = { 

475 "iterations": iterations, 

476 "questions_per_iteration": questions_per_iteration, 

477 "search_strategy": search_strategy, 

478 "search_tool": self.search_tool, 

479 "max_results": max_results, 

480 "model_name": self.model_name, 

481 "provider": self.provider, 

482 } 

483 

484 # Evaluate quality using composite benchmark evaluator 

485 # Use a small number of examples for efficiency 

486 benchmark_dir = str(Path(self.output_dir) / "benchmark_temp") 

487 quality_results = self.benchmark_evaluator.evaluate( 

488 system_config=system_config, 

489 num_examples=5, # Small number for optimization efficiency 

490 output_dir=benchmark_dir, 

491 ) 

492 

493 # Stop timing 

494 speed_profiler.stop() 

495 timing_results = speed_profiler.get_summary() 

496 

497 # Extract key metrics 

498 quality_score = quality_results.get("quality_score", 0.0) 

499 benchmark_results = quality_results.get("benchmark_results", {}) 

500 

501 # Speed score: convert duration to a 0-1 score where faster is better 

502 # Using a reasonable threshold (e.g., 180 seconds for 5 examples) 

503 # Below this threshold: high score, above it: declining score 

504 total_duration = timing_results.get("total_duration", 180) 

505 speed_score = max(0.0, min(1.0, 1.0 - (total_duration - 60) / 180)) 

506 

507 # Calculate combined score based on weights 

508 combined_score = ( 

509 self.metric_weights.get("quality", 0.6) * quality_score 

510 + self.metric_weights.get("speed", 0.4) * speed_score 

511 ) 

512 

513 # Return streamlined results 

514 return { 

515 "quality_score": quality_score, 

516 "benchmark_results": benchmark_results, 

517 "speed_score": speed_score, 

518 "total_duration": total_duration, 

519 "score": combined_score, 

520 "success": True, 

521 } 

522 

523 except Exception as e: 

524 # Stop profiling on error 

525 speed_profiler.stop() 

526 

527 # Log error 

528 logger.exception("Error in experiment") 

529 

530 # Return error information 

531 return {"error": str(e), "score": 0.0, "success": False} 

532 

533 def _optimization_callback(self, study: optuna.Study, trial: optuna.Trial): 

534 """ 

535 Callback for the Optuna optimization process. 

536 

537 Args: 

538 study: Optuna study object 

539 trial: Current trial 

540 """ 

541 # Save intermediate results periodically 

542 if trial.number % 10 == 0 and trial.number > 0: 

543 self._save_results() 

544 self._create_quick_visualizations() 

545 

546 def _save_results(self): 

547 """Save the optimization results to disk.""" 

548 # Create a timestamp for filenames 

549 timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") 

550 

551 # Save trial history 

552 from ...security.file_write_verifier import write_json_verified 

553 

554 history_file = str( 

555 Path(self.output_dir) / f"{self.study_name}_history.json" 

556 ) 

557 

558 # Convert numpy values to native Python types for JSON serialization 

559 clean_history: List[Dict[str, Any]] = [] 

560 for trial in self.trials_history: 

561 clean_trial: Dict[str, Any] = {} 

562 for k, v in trial.items(): 

563 if isinstance(v, dict): 

564 clean_trial[k] = { 

565 dk: (float(dv) if isinstance(dv, np.number) else dv) 

566 for dk, dv in v.items() 

567 } 

568 elif isinstance(v, np.number): 

569 clean_trial[k] = float(v) 

570 else: 

571 clean_trial[k] = v 

572 clean_history.append(clean_trial) 

573 

574 # Sanitize sensitive data before writing to disk 

575 sanitized_history = sanitize_data(clean_history) 

576 

577 write_json_verified( 

578 history_file, 

579 sanitized_history, 

580 "benchmark.allow_file_output", 

581 context="optimization history", 

582 ) 

583 

584 # Save current best parameters 

585 if ( 

586 self.study 

587 and hasattr(self.study, "best_params") 

588 and self.study.best_params 

589 ): 

590 best_params_file = str( 

591 Path(self.output_dir) / f"{self.study_name}_best_params.json" 

592 ) 

593 

594 best_params_data = { 

595 "best_params": self.study.best_params, 

596 "best_value": float(self.study.best_value), 

597 "n_trials": len(self.study.trials), 

598 "timestamp": timestamp, 

599 "base_query": self.base_query, 

600 "model_name": self.model_name, 

601 "provider": self.provider, 

602 "search_tool": self.search_tool, 

603 "metric_weights": self.metric_weights, 

604 "benchmark_weights": self.benchmark_weights, 

605 } 

606 

607 # Sanitize sensitive data before writing to disk 

608 sanitized_best_params = sanitize_data(best_params_data) 

609 

610 write_json_verified( 

611 best_params_file, 

612 sanitized_best_params, 

613 "benchmark.allow_file_output", 

614 context="optimization best params", 

615 ) 

616 

617 # Save the Optuna study 

618 if self.study: 

619 study_file = str( 

620 Path(self.output_dir) / f"{self.study_name}_study.pkl" 

621 ) 

622 joblib.dump(self.study, study_file) 

623 

624 logger.info(f"Results saved to {self.output_dir}") 

625 

626 def _create_visualizations(self): 

627 """Create and save comprehensive visualizations of the optimization results.""" 

628 if not PLOTTING_AVAILABLE: 

629 logger.warning( 

630 "Matplotlib not available, skipping visualization creation" 

631 ) 

632 return 

633 _ensure_plotting_loaded() 

634 

635 if not self.study or len(self.study.trials) < 2: 

636 logger.warning("Not enough trials to create visualizations") 

637 return 

638 

639 # Create directory for visualizations 

640 _viz_dir_path = Path(self.output_dir) / "visualizations" 

641 # nosemgrep: semgrep.rules.path-traversal-risk, reason: CLI/Python optimizer output directory with a fixed visualizations child 

642 _viz_dir_path.mkdir(parents=True, exist_ok=True) 

643 viz_dir = str(_viz_dir_path) 

644 

645 # Create Optuna visualizations 

646 self._create_optuna_visualizations(viz_dir) 

647 

648 # Create custom visualizations 

649 self._create_custom_visualizations(viz_dir) 

650 

651 logger.info(f"Visualizations saved to {viz_dir}") 

652 

653 def _create_quick_visualizations(self): 

654 """Create a smaller set of visualizations for intermediate progress.""" 

655 if ( 

656 not PLOTTING_AVAILABLE 

657 or not self.study 

658 or len(self.study.trials) < 2 

659 ): 

660 return 

661 _ensure_plotting_loaded() 

662 

663 # Create directory for visualizations 

664 _quick_viz_dir_path = Path(self.output_dir) / "visualizations" 

665 # nosemgrep: semgrep.rules.path-traversal-risk, reason: CLI/Python optimizer output directory with a fixed visualizations child 

666 _quick_viz_dir_path.mkdir(parents=True, exist_ok=True) 

667 viz_dir = str(_quick_viz_dir_path) 

668 

669 # Create optimization history only (faster than full visualization) 

670 try: 

671 fig = plot_optimization_history(self.study) 

672 fig.write_image( 

673 str( 

674 Path(viz_dir) 

675 / f"{self.study_name}_optimization_history_current.png" 

676 ) 

677 ) 

678 except Exception: 

679 logger.exception("Error creating optimization history plot") 

680 

681 def _create_optuna_visualizations(self, viz_dir: str): 

682 """ 

683 Create and save Optuna's built-in visualizations. 

684 

685 Args: 

686 viz_dir: Directory to save visualizations 

687 """ 

688 if not self.study or not PLOTTING_AVAILABLE: 688 ↛ 689line 688 didn't jump to line 689 because the condition on line 688 was never true

689 return 

690 _ensure_plotting_loaded() 

691 study = self.study 

692 timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") 

693 

694 # 1. Optimization history 

695 try: 

696 fig = plot_optimization_history(study) 

697 fig.write_image( 

698 str( 

699 Path(viz_dir) 

700 / f"{self.study_name}_optimization_history_{timestamp}.png" 

701 ) 

702 ) 

703 except Exception: 

704 logger.exception("Error creating optimization history plot") 

705 

706 # 2. Parameter importances 

707 try: 

708 fig = plot_param_importances(study) 

709 fig.write_image( 

710 str( 

711 Path(viz_dir) 

712 / f"{self.study_name}_param_importances_{timestamp}.png" 

713 ) 

714 ) 

715 except Exception: 

716 logger.exception("Error creating parameter importances plot") 

717 

718 # 3. Slice plot for each parameter 

719 try: 

720 for param_name in study.best_params.keys(): 

721 fig = plot_slice(study, [param_name]) 

722 fig.write_image( 

723 str( 

724 Path(viz_dir) 

725 / f"{self.study_name}_slice_{param_name}_{timestamp}.png" 

726 ) 

727 ) 

728 except Exception: 

729 logger.exception("Error creating slice plots") 

730 

731 # 4. Contour plots for important parameter pairs 

732 try: 

733 # Get all parameter names 

734 param_names = list(study.best_params.keys()) 

735 

736 # Create contour plots for each pair 

737 for i in range(len(param_names)): 

738 for j in range(i + 1, len(param_names)): 738 ↛ 739line 738 didn't jump to line 739 because the loop on line 738 never started

739 try: 

740 fig = plot_contour( 

741 study, params=[param_names[i], param_names[j]] 

742 ) 

743 fig.write_image( 

744 str( 

745 Path(viz_dir) 

746 / f"{self.study_name}_contour_{param_names[i]}_{param_names[j]}_{timestamp}.png" 

747 ) 

748 ) 

749 except Exception: 

750 logger.warning( 

751 f"Error creating contour plot for {param_names[i]} vs {param_names[j]}" 

752 ) 

753 except Exception: 

754 logger.exception("Error creating contour plots") 

755 

756 def _create_custom_visualizations(self, viz_dir: str): 

757 """ 

758 Create custom visualizations based on trial history. 

759 

760 Args: 

761 viz_dir: Directory to save visualizations 

762 """ 

763 if not self.trials_history: 

764 return 

765 

766 timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") 

767 

768 # Create quality vs speed plot 

769 self._create_quality_vs_speed_plot(viz_dir, timestamp) 

770 

771 # Create parameter evolution plots 

772 self._create_parameter_evolution_plots(viz_dir, timestamp) 

773 

774 # Create trial duration vs score plot 

775 self._create_duration_vs_score_plot(viz_dir, timestamp) 

776 

777 def _create_quality_vs_speed_plot(self, viz_dir: str, timestamp: str): 

778 """Create a plot showing quality vs. speed trade-off.""" 

779 if not self.trials_history: 779 ↛ 780line 779 didn't jump to line 780 because the condition on line 779 was never true

780 return 

781 

782 # Extract data from successful trials 

783 successful_trials = [ 

784 t 

785 for t in self.trials_history 

786 if t.get("result", {}).get("success", False) 

787 ] 

788 

789 if not successful_trials: 

790 logger.warning("No successful trials for visualization") 

791 return 

792 

793 try: 

794 plt.figure(figsize=(10, 8)) 

795 

796 # Extract metrics 

797 quality_scores = [] 

798 speed_scores = [] 

799 labels = [] 

800 iterations_values = [] 

801 questions_values = [] 

802 

803 for trial in successful_trials: 803 ↛ 817line 803 didn't jump to line 817 because the loop on line 803 didn't complete

804 result = trial["result"] 

805 quality = result.get("quality_score", 0) 

806 speed = result.get("speed_score", 0) 

807 iterations = trial["params"].get("iterations", 0) 

808 questions = trial["params"].get("questions_per_iteration", 0) 

809 

810 quality_scores.append(quality) 

811 speed_scores.append(speed) 

812 labels.append(f"Trial {trial['trial_number']}") 

813 iterations_values.append(iterations) 

814 questions_values.append(questions) 

815 

816 # Create scatter plot with size based on iterations*questions 

817 sizes = [ 

818 i * q * 5 

819 for i, q in zip( 

820 iterations_values, questions_values, strict=False 

821 ) 

822 ] 

823 scatter = plt.scatter( 

824 quality_scores, 

825 speed_scores, 

826 s=sizes, 

827 alpha=0.7, 

828 c=range(len(quality_scores)), 

829 cmap="viridis", 

830 ) 

831 

832 # Highlight best trial 

833 best_trial = max( 

834 successful_trials, 

835 key=lambda x: x.get("result", {}).get("score", 0), 

836 ) 

837 best_quality = best_trial["result"].get("quality_score", 0) 

838 best_speed = best_trial["result"].get("speed_score", 0) 

839 best_iter = best_trial["params"].get("iterations", 0) 

840 best_questions = best_trial["params"].get( 

841 "questions_per_iteration", 0 

842 ) 

843 

844 plt.scatter( 

845 [best_quality], 

846 [best_speed], 

847 s=200, 

848 facecolors="none", 

849 edgecolors="red", 

850 linewidth=2, 

851 label=f"Best: {best_iter}×{best_questions}", 

852 ) 

853 

854 # Add annotations for key points 

855 for i, (q, s, label) in enumerate( 

856 zip(quality_scores, speed_scores, labels, strict=False) 

857 ): 

858 if i % max(1, len(quality_scores) // 5) == 0: # Label ~5 points 

859 plt.annotate( 

860 f"{iterations_values[i]}×{questions_values[i]}", 

861 (q, s), 

862 xytext=(5, 5), 

863 textcoords="offset points", 

864 ) 

865 

866 # Add colorbar and labels 

867 cbar = plt.colorbar(scatter) 

868 cbar.set_label("Trial Progression") 

869 

870 # Add benchmark weight information 

871 weights_str = ", ".join( 

872 [f"{k}:{v:.1f}" for k, v in self.benchmark_weights.items()] 

873 ) 

874 plt.title( 

875 f"Quality vs. Speed Trade-off\nBenchmark Weights: {weights_str}" 

876 ) 

877 plt.xlabel("Quality Score (Benchmark Accuracy)") 

878 plt.ylabel("Speed Score") 

879 plt.grid(True, linestyle="--", alpha=0.7)