Source code for arcade.classifieroptimization.tuner

"""Hyperparameter tuning optimizer."""

from __future__ import annotations

import itertools
import logging
from typing import Any, Callable

import numpy as np

from arcade._logging import get_logger, tqdm
from .base import BaseOptimizer, OptimizationResult, register_optimizer

log = get_logger(__name__)


def _compute_metric(preds: np.ndarray, labels: np.ndarray, metric: str) -> float:
    """Compute the requested metric from predictions and labels."""
    if metric in ("accuracy", "val_accuracy"):
        return float((preds == labels).mean())
    if metric == "f1":
        unique = np.unique(labels)
        if len(unique) <= 2:
            pos = unique[-1]
            tp = float(((preds == pos) & (labels == pos)).sum())
            fp = float(((preds == pos) & (labels != pos)).sum())
            fn = float(((preds != pos) & (labels == pos)).sum())
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
            recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
            return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
        f1s = []
        for c in unique:
            tp = float(((preds == c) & (labels == c)).sum())
            fp = float(((preds == c) & (labels != c)).sum())
            fn = float(((preds != c) & (labels == c)).sum())
            prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0
            rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0
            f1s.append(2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0.0)
        return float(np.mean(f1s))
    return float((preds == labels).mean())


[docs] @register_optimizer("tuning") class HyperparameterTuner(BaseOptimizer): """Hyperparameter tuning via Optuna or grid search.""" def __init__( self, method: str = "optuna", n_trials: int = 100, metric: str = "accuracy", search_space: dict[str, Any] | None = None, ) -> None: self.method = method self.n_trials = n_trials self.metric = metric self.search_space = search_space or {}
[docs] def optimize( self, model: Any, train_data: tuple[Any, Any], val_data: tuple[Any, Any] | None = None, **kwargs: Any, ) -> OptimizationResult: train_features, train_labels = train_data if val_data is not None: val_features, val_labels = val_data else: val_features, val_labels = train_features, train_labels classifier_factory = kwargs.get("classifier_factory") if classifier_factory is None: def classifier_factory(params: dict) -> Any: # type: ignore[misc] return model if self.method == "optuna": result = self._tune_optuna( classifier_factory, train_features, train_labels, val_features, val_labels, **kwargs, ) elif self.method == "grid": result = self._tune_grid( classifier_factory, train_features, train_labels, val_features, val_labels, **kwargs, ) else: raise ValueError(f"Unknown tuning method: {self.method!r}") best_model = classifier_factory(result["best_params"]) best_model.fit(train_features, train_labels, val_features=val_features, val_labels=val_labels) return OptimizationResult( model=best_model, accuracy_after=result["best_score"], metadata=result, )
def _tune_optuna( self, factory: Callable, train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, **kwargs: Any, ) -> dict[str, Any]: import optuna optuna.logging.set_verbosity(optuna.logging.WARNING) search_space = self.search_space metric = self.metric def objective(trial: optuna.Trial) -> float: params: dict[str, Any] = {} for name, spec in search_space.items(): if isinstance(spec, list): params[name] = trial.suggest_categorical(name, spec) elif isinstance(spec, dict): low, high = spec["low"], spec["high"] use_log = spec.get("log", False) if isinstance(low, float) or isinstance(high, float): params[name] = trial.suggest_float(name, low, high, log=use_log) else: params[name] = trial.suggest_int(name, low, high, log=use_log) else: params[name] = spec clf = factory(params) clf.fit(train_features, train_labels, val_features=val_features, val_labels=val_labels) preds = clf.predict(val_features) return _compute_metric(preds, val_labels, metric) study = optuna.create_study(direction="maximize") study.optimize(objective, n_trials=self.n_trials) return { "best_params": study.best_params, "best_score": study.best_value, "study": study, } def _tune_grid( self, factory: Callable, train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, **kwargs: Any, ) -> dict[str, Any]: param_names = list(self.search_space.keys()) param_values = [v if isinstance(v, list) else [v] for v in self.search_space.values()] best_score = -1.0 best_params: dict[str, Any] = {} all_results: list[dict[str, Any]] = [] for combo in itertools.product(*param_values): params = dict(zip(param_names, combo)) clf = factory(params) clf.fit(train_features, train_labels, val_features=val_features, val_labels=val_labels) preds = clf.predict(val_features) score = _compute_metric(preds, val_labels, self.metric) all_results.append({"params": params, "score": score}) if score > best_score: best_score = score best_params = params return { "best_params": best_params, "best_score": best_score, "all_results": all_results, }
[docs] def tune( classifier_factory: Callable[..., Any], train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, *, method: str = "optuna", n_trials: int = 50, search_space: dict | None = None, metric: str = "val_accuracy", **kwargs: Any, ) -> dict[str, Any]: """Run hyperparameter optimization for a classifier. Args: classifier_factory: Callable that accepts a dict of hyperparameters and returns a fresh classifier instance. train_features: Training features. train_labels: Training labels. val_features: Validation features. val_labels: Validation labels. method: ``"optuna"`` or ``"grid"``. n_trials: Number of optimization trials (Optuna only). search_space: Per-parameter search ranges. For Optuna, values can be lists (categorical) or dicts with ``low``/``high`` (and optional ``log``). For grid, all values must be lists. metric: Metric to optimize (e.g. ``"val_accuracy"``). **kwargs: Forwarded to the classifier's ``fit()`` method. Returns: Dict with ``"best_params"``, ``"best_score"``, and ``"study"`` (if Optuna) or ``"all_results"`` (if grid). """ if search_space is None: search_space = {} if method == "optuna": return _tune_optuna( classifier_factory, train_features, train_labels, val_features, val_labels, n_trials=n_trials, search_space=search_space, metric=metric, **kwargs, ) if method == "grid": return _tune_grid( classifier_factory, train_features, train_labels, val_features, val_labels, search_space=search_space, metric=metric, **kwargs, ) raise ValueError(f"Unknown tuning method: {method!r}")
def _tune_optuna( factory: Callable, train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, *, n_trials: int, search_space: dict, metric: str, **kwargs: Any, ) -> dict[str, Any]: try: import optuna except ImportError: raise ImportError( "Optuna is required for hyperparameter tuning. " "Install with: pip install arcade-readout[tuning]" ) optuna.logging.set_verbosity(optuna.logging.WARNING) def objective(trial: optuna.Trial) -> float: params: dict[str, Any] = {} for name, spec in search_space.items(): if isinstance(spec, list): params[name] = trial.suggest_categorical(name, spec) elif isinstance(spec, dict): low, high = spec["low"], spec["high"] use_log = spec.get("log", False) if isinstance(low, float) or isinstance(high, float): params[name] = trial.suggest_float(name, low, high, log=use_log) else: params[name] = trial.suggest_int(name, low, high, log=use_log) else: params[name] = spec clf = factory(params) clf.fit( train_features, train_labels, val_features=val_features, val_labels=val_labels, **kwargs, ) preds = clf.predict(val_features) return float((preds == val_labels).mean()) study = optuna.create_study(direction="maximize") study.optimize(objective, n_trials=n_trials) log.info( "Optuna tuning: best_score=%.4f, best_params=%s", study.best_value, study.best_params, ) return { "best_params": study.best_params, "best_score": study.best_value, "study": study, } def _tune_grid( factory: Callable, train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, *, search_space: dict, metric: str, **kwargs: Any, ) -> dict[str, Any]: param_names = list(search_space.keys()) param_values = [v if isinstance(v, list) else [v] for v in search_space.values()] best_score = -1.0 best_params: dict[str, Any] = {} all_results: list[dict[str, Any]] = [] combos = list(itertools.product(*param_values)) for combo in tqdm(combos, desc="Grid search", leave=False): params = dict(zip(param_names, combo)) clf = factory(params) clf.fit( train_features, train_labels, val_features=val_features, val_labels=val_labels, **kwargs, ) preds = clf.predict(val_features) score = float((preds == val_labels).mean()) all_results.append({"params": params, "score": score}) if score > best_score: best_score = score best_params = params log.info( "Grid search: best_score=%.4f, best_params=%s (%d combos)", best_score, best_params, len(combos), ) return { "best_params": best_params, "best_score": best_score, "all_results": all_results, }
[docs] def maybe_tune( cfg: Any, clf_name: str, clf_cls: type, train_features: np.ndarray, train_labels: np.ndarray, val_features: np.ndarray, val_labels: np.ndarray, filter_set: Any, ) -> dict | None: """Conditionally run tuning if search_space is configured (moved from stages.py).""" search_space = cfg.tuning.search_space.get(clf_name, {}) if not search_space: return None def factory(**params: Any) -> Any: kw = dict(params) if clf_name == "threshold" and filter_set is not None: kw["filter_set"] = filter_set return clf_cls.from_config(kw) log.info( " Tuning %s (%s, %d trials)", clf_name, cfg.tuning.method, cfg.tuning.n_trials, ) result = tune( factory, train_features, train_labels, val_features, val_labels, method=cfg.tuning.method, n_trials=cfg.tuning.n_trials, metric=cfg.tuning.metric, search_space=search_space, ) best_params = result.get("best_params", result) if isinstance(result, dict) else result log.info(" Best params: %s", best_params) return best_params