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 _viz_dir_path.mkdir(parents=True, exist_ok=True) 

642 viz_dir = str(_viz_dir_path) 

643 

644 # Create Optuna visualizations 

645 self._create_optuna_visualizations(viz_dir) 

646 

647 # Create custom visualizations 

648 self._create_custom_visualizations(viz_dir) 

649 

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

651 

652 def _create_quick_visualizations(self): 

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

654 if ( 

655 not PLOTTING_AVAILABLE 

656 or not self.study 

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

658 ): 

659 return 

660 _ensure_plotting_loaded() 

661 

662 # Create directory for visualizations 

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

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

665 viz_dir = str(_quick_viz_dir_path) 

666 

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

668 try: 

669 fig = plot_optimization_history(self.study) 

670 fig.write_image( 

671 str( 

672 Path(viz_dir) 

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

674 ) 

675 ) 

676 except Exception: 

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

678 

679 def _create_optuna_visualizations(self, viz_dir: str): 

680 """ 

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

682 

683 Args: 

684 viz_dir: Directory to save visualizations 

685 """ 

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

687 return 

688 _ensure_plotting_loaded() 

689 study = self.study 

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

691 

692 # 1. Optimization history 

693 try: 

694 fig = plot_optimization_history(study) 

695 fig.write_image( 

696 str( 

697 Path(viz_dir) 

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

699 ) 

700 ) 

701 except Exception: 

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

703 

704 # 2. Parameter importances 

705 try: 

706 fig = plot_param_importances(study) 

707 fig.write_image( 

708 str( 

709 Path(viz_dir) 

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

711 ) 

712 ) 

713 except Exception: 

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

715 

716 # 3. Slice plot for each parameter 

717 try: 

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

719 fig = plot_slice(study, [param_name]) 

720 fig.write_image( 

721 str( 

722 Path(viz_dir) 

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

724 ) 

725 ) 

726 except Exception: 

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

728 

729 # 4. Contour plots for important parameter pairs 

730 try: 

731 # Get all parameter names 

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

733 

734 # Create contour plots for each pair 

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

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

737 try: 

738 fig = plot_contour( 

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

740 ) 

741 fig.write_image( 

742 str( 

743 Path(viz_dir) 

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

745 ) 

746 ) 

747 except Exception: 

748 logger.warning( 

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

750 ) 

751 except Exception: 

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

753 

754 def _create_custom_visualizations(self, viz_dir: str): 

755 """ 

756 Create custom visualizations based on trial history. 

757 

758 Args: 

759 viz_dir: Directory to save visualizations 

760 """ 

761 if not self.trials_history: 

762 return 

763 

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

765 

766 # Create quality vs speed plot 

767 self._create_quality_vs_speed_plot(viz_dir, timestamp) 

768 

769 # Create parameter evolution plots 

770 self._create_parameter_evolution_plots(viz_dir, timestamp) 

771 

772 # Create trial duration vs score plot 

773 self._create_duration_vs_score_plot(viz_dir, timestamp) 

774 

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

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

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

778 return 

779 

780 # Extract data from successful trials 

781 successful_trials = [ 

782 t 

783 for t in self.trials_history 

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

785 ] 

786 

787 if not successful_trials: 

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

789 return 

790 

791 try: 

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

793 

794 # Extract metrics 

795 quality_scores = [] 

796 speed_scores = [] 

797 labels = [] 

798 iterations_values = [] 

799 questions_values = [] 

800 

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

802 result = trial["result"] 

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

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

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

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

807 

808 quality_scores.append(quality) 

809 speed_scores.append(speed) 

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

811 iterations_values.append(iterations) 

812 questions_values.append(questions) 

813 

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

815 sizes = [ 

816 i * q * 5 

817 for i, q in zip( 

818 iterations_values, questions_values, strict=False 

819 ) 

820 ] 

821 scatter = plt.scatter( 

822 quality_scores, 

823 speed_scores, 

824 s=sizes, 

825 alpha=0.7, 

826 c=range(len(quality_scores)), 

827 cmap="viridis", 

828 ) 

829 

830 # Highlight best trial 

831 best_trial = max( 

832 successful_trials, 

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

834 ) 

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

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

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

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

839 "questions_per_iteration", 0 

840 ) 

841 

842 plt.scatter( 

843 [best_quality], 

844 [best_speed], 

845 s=200, 

846 facecolors="none", 

847 edgecolors="red", 

848 linewidth=2, 

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

850 ) 

851 

852 # Add annotations for key points 

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

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

855 ): 

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

857 plt.annotate( 

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

859 (q, s), 

860 xytext=(5, 5), 

861 textcoords="offset points", 

862 ) 

863 

864 # Add colorbar and labels 

865 cbar = plt.colorbar(scatter) 

866 cbar.set_label("Trial Progression") 

867 

868 # Add benchmark weight information 

869 weights_str = ", ".join( 

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

871 ) 

872 plt.title( 

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

874 ) 

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

876 plt.ylabel("Speed Score") 

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

878 

879 # Add legend explaining size