Coverage for src / tracekit / api / optimization.py: 99%
127 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-11 23:04 +0000
1"""Parameter optimization for signal analysis.
3This module provides parameter optimization utilities including
4grid search, parameter space definition, and optimization result tracking.
5"""
7from __future__ import annotations
9import itertools
10import logging
11import time
12from dataclasses import dataclass, field
13from typing import TYPE_CHECKING, Any
15import numpy as np
17if TYPE_CHECKING:
18 from collections.abc import Callable, Iterator
20logger = logging.getLogger(__name__)
22__all__ = [
23 "GridSearch",
24 "OptimizationResult",
25 "ParameterSpace",
26 "optimize_parameters",
27]
30@dataclass
31class ParameterSpace:
32 """Definition of parameter search space.
34 Attributes:
35 name: Parameter name
36 values: List of values to try
37 low: Low bound (for continuous params)
38 high: High bound (for continuous params)
39 log_scale: Use logarithmic scale
40 num_samples: Number of samples for continuous
42 Example:
43 >>> # Discrete parameter
44 >>> window = ParameterSpace("window", values=["hann", "hamming", "blackman"])
45 >>> # Continuous parameter
46 >>> cutoff = ParameterSpace("cutoff", low=1e3, high=1e6, num_samples=10)
48 References:
49 API-014: Parameter Optimization
50 """
52 name: str
53 values: list[Any] | None = None
54 low: float | None = None
55 high: float | None = None
56 log_scale: bool = False
57 num_samples: int = 10
59 def __post_init__(self) -> None:
60 """Generate values if continuous parameter."""
61 if self.values is None:
62 if self.low is not None and self.high is not None:
63 if self.log_scale:
64 self.values = list(
65 np.logspace(np.log10(self.low), np.log10(self.high), self.num_samples)
66 )
67 else:
68 self.values = list(np.linspace(self.low, self.high, self.num_samples))
69 else:
70 raise ValueError(
71 f"Parameter {self.name}: must specify either values or (low, high) bounds"
72 )
74 def __iter__(self) -> Iterator[Any]:
75 """Iterate over parameter values."""
76 return iter(self.values or [])
78 def __len__(self) -> int:
79 """Number of parameter values."""
80 return len(self.values or [])
83@dataclass
84class OptimizationResult:
85 """Result of parameter optimization.
87 Attributes:
88 best_params: Best parameter combination
89 best_score: Best objective score
90 all_results: All evaluated combinations
91 elapsed_time: Total optimization time
92 num_evaluations: Number of combinations evaluated
94 References:
95 API-014: Parameter Optimization
96 """
98 best_params: dict[str, Any]
99 best_score: float
100 all_results: list[tuple[dict[str, Any], float]] = field(default_factory=list)
101 elapsed_time: float = 0.0
102 num_evaluations: int = 0
104 def top_n(self, n: int = 5) -> list[tuple[dict[str, Any], float]]:
105 """Get top N parameter combinations.
107 Args:
108 n: Number of top results
110 Returns:
111 List of (params, score) tuples
112 """
113 sorted_results = sorted(self.all_results, key=lambda x: x[1], reverse=True)
114 return sorted_results[:n]
116 def to_dict(self) -> dict[str, Any]:
117 """Convert to dictionary."""
118 return {
119 "best_params": self.best_params,
120 "best_score": self.best_score,
121 "num_evaluations": self.num_evaluations,
122 "elapsed_time": self.elapsed_time,
123 }
126class GridSearch:
127 """Grid search optimization for parameter tuning.
129 Exhaustively searches all combinations of parameters to find
130 the best combination based on an objective function.
132 Example:
133 >>> def objective(params, data):
134 ... result = analyze(data, **params)
135 ... return result.snr
136 >>>
137 >>> search = GridSearch([
138 ... ParameterSpace("nfft", values=[1024, 2048, 4096, 8192]),
139 ... ParameterSpace("window", values=["hann", "hamming"]),
140 ... ParameterSpace("overlap", low=0.25, high=0.75, num_samples=5)
141 ... ])
142 >>>
143 >>> result = search.fit(objective, data)
144 >>> print(f"Best params: {result.best_params}")
146 References:
147 API-014: Parameter Optimization
148 """
150 def __init__(self, param_spaces: list[ParameterSpace], verbose: bool = True):
151 """Initialize grid search.
153 Args:
154 param_spaces: List of parameter spaces
155 verbose: Print progress
156 """
157 self.param_spaces = param_spaces
158 self.verbose = verbose
159 self._progress_callback: Callable[[int, int], None] | None = None
161 @property
162 def num_combinations(self) -> int:
163 """Total number of parameter combinations."""
164 total = 1
165 for space in self.param_spaces:
166 total *= len(space)
167 return total
169 def on_progress(self, callback: Callable[[int, int], None]) -> GridSearch:
170 """Set progress callback.
172 Args:
173 callback: Function called with (current, total)
175 Returns:
176 Self (for chaining)
177 """
178 self._progress_callback = callback
179 return self
181 def fit(
182 self,
183 objective: Callable[[dict[str, Any], Any], float],
184 data: Any,
185 *,
186 maximize: bool = True,
187 early_stop: float | None = None,
188 ) -> OptimizationResult:
189 """Run grid search optimization.
191 Args:
192 objective: Objective function (params, data) -> score
193 data: Data to pass to objective
194 maximize: If True, maximize score; if False, minimize
195 early_stop: Stop if score reaches this threshold
197 Returns:
198 Optimization result
199 """
200 start_time = time.time()
201 all_results: list[tuple[dict[str, Any], float]] = []
202 best_params: dict[str, Any] = {}
203 best_score = float("-inf") if maximize else float("inf")
205 # Generate all combinations
206 param_names = [s.name for s in self.param_spaces]
207 param_values = [list(s) for s in self.param_spaces]
209 total = self.num_combinations
210 if self.verbose:
211 logger.info(f"Grid search: {total} combinations")
213 for i, values in enumerate(itertools.product(*param_values)):
214 params = dict(zip(param_names, values, strict=False))
216 try:
217 score = objective(params, data)
218 except Exception as e:
219 logger.warning(f"Objective failed for {params}: {e}")
220 score = float("-inf") if maximize else float("inf")
222 all_results.append((params, score))
224 # Update best
225 if maximize:
226 if score > best_score:
227 best_score = score
228 best_params = params.copy()
229 elif score < best_score:
230 best_score = score
231 best_params = params.copy()
233 # Progress
234 if self._progress_callback:
235 self._progress_callback(i + 1, total)
237 # Early stopping
238 if early_stop is not None and (
239 (maximize and score >= early_stop) or (not maximize and score <= early_stop)
240 ):
241 if self.verbose:
242 logger.info(f"Early stop at {i + 1}/{total}")
243 break
245 elapsed = time.time() - start_time
247 if self.verbose:
248 logger.info(f"Completed: best_score={best_score:.4f}, time={elapsed:.2f}s")
250 return OptimizationResult(
251 best_params=best_params,
252 best_score=best_score,
253 all_results=all_results,
254 elapsed_time=elapsed,
255 num_evaluations=len(all_results),
256 )
259class RandomSearch:
260 """Random search optimization.
262 Samples random combinations from parameter space.
264 References:
265 API-014: Parameter Optimization
266 """
268 def __init__(
269 self,
270 param_spaces: list[ParameterSpace],
271 n_iterations: int = 100,
272 random_state: int | None = None,
273 ):
274 """Initialize random search.
276 Args:
277 param_spaces: Parameter spaces
278 n_iterations: Number of random samples
279 random_state: Random seed
280 """
281 self.param_spaces = param_spaces
282 self.n_iterations = n_iterations
283 self.random_state = random_state
285 def fit(
286 self,
287 objective: Callable[[dict[str, Any], Any], float],
288 data: Any,
289 *,
290 maximize: bool = True,
291 ) -> OptimizationResult:
292 """Run random search.
294 Args:
295 objective: Objective function
296 data: Data for objective
297 maximize: Maximize or minimize
299 Returns:
300 Optimization result
301 """
302 rng = np.random.default_rng(self.random_state)
303 start_time = time.time()
304 all_results: list[tuple[dict[str, Any], float]] = []
305 best_params: dict[str, Any] = {}
306 best_score = float("-inf") if maximize else float("inf")
308 for _ in range(self.n_iterations):
309 # Sample random parameters
310 params = {}
311 for space in self.param_spaces:
312 if space.values: 312 ↛ 311line 312 didn't jump to line 311 because the condition on line 312 was always true
313 params[space.name] = rng.choice(space.values)
315 try:
316 score = objective(params, data)
317 except Exception:
318 score = float("-inf") if maximize else float("inf")
320 all_results.append((params, score))
322 if maximize:
323 if score > best_score:
324 best_score = score
325 best_params = params.copy()
326 elif score < best_score:
327 best_score = score
328 best_params = params.copy()
330 return OptimizationResult(
331 best_params=best_params,
332 best_score=best_score,
333 all_results=all_results,
334 elapsed_time=time.time() - start_time,
335 num_evaluations=len(all_results),
336 )
339def optimize_parameters(
340 objective: Callable[[dict[str, Any], Any], float],
341 data: Any,
342 param_spaces: list[ParameterSpace] | dict[str, list[Any]],
343 *,
344 method: str = "grid",
345 maximize: bool = True,
346 **kwargs: Any,
347) -> OptimizationResult:
348 """Optimize parameters for objective function.
350 Convenience function for parameter optimization.
352 Args:
353 objective: Objective function (params, data) -> score
354 data: Data to pass to objective
355 param_spaces: Parameter spaces (list or dict)
356 method: Optimization method ("grid", "random")
357 maximize: Maximize or minimize
358 **kwargs: Additional arguments for optimizer
360 Returns:
361 Optimization result
363 Raises:
364 ValueError: If method is not one of the supported types.
366 Example:
367 >>> result = optimize_parameters(
368 ... objective=lambda p, d: analyze(d, **p).snr,
369 ... data=trace,
370 ... param_spaces={
371 ... "nfft": [1024, 2048, 4096],
372 ... "window": ["hann", "hamming"]
373 ... }
374 ... )
376 References:
377 API-014: Parameter Optimization
378 """
379 # Convert dict to ParameterSpace list
380 if isinstance(param_spaces, dict):
381 param_spaces = [
382 ParameterSpace(name, values=values) for name, values in param_spaces.items()
383 ]
385 if method == "grid":
386 optimizer = GridSearch(param_spaces, **kwargs)
387 elif method == "random":
388 optimizer = RandomSearch(param_spaces, **kwargs) # type: ignore[assignment]
389 else:
390 raise ValueError(f"Unknown optimization method: {method}")
392 return optimizer.fit(objective, data, maximize=maximize)