394 lines
16 KiB
Python
394 lines
16 KiB
Python
|
"""
|
||
|
|
Adaptive surrogate modeling with confidence-based exploration-exploitation transitions.
|
||
|
|
|
||
|
|
This module implements state-of-the-art Bayesian optimization strategies that
|
||
|
|
dynamically adjust exploration vs exploitation based on surrogate model confidence.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from typing import Optional, Dict, List
|
||
|
|
import optuna
|
||
|
|
from scipy.stats import variation
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
class SurrogateConfidenceMetrics:
|
||
|
|
"""Calculate confidence metrics for surrogate model quality.
|
||
|
|
|
||
|
|
STUDY-AWARE: Uses study.trials directly instead of session-based history
|
||
|
|
to properly track confidence across multiple optimization runs.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, min_trials_for_confidence: int = 15):
|
||
|
|
self.min_trials = min_trials_for_confidence
|
||
|
|
|
||
|
|
def update(self, study: optuna.Study, trial: optuna.trial.FrozenTrial):
|
||
|
|
"""Update metrics after each trial (no-op for study-aware implementation)."""
|
||
|
|
pass # Study-aware: we read directly from study.trials
|
||
|
|
|
||
|
|
def calculate_confidence(self, study: optuna.Study) -> Dict[str, float]:
|
||
|
|
"""
|
||
|
|
Calculate comprehensive surrogate confidence metrics.
|
||
|
|
|
||
|
|
STUDY-AWARE: Uses ALL completed trials from the study database,
|
||
|
|
not just trials from the current session.
|
||
|
|
|
||
|
|
PROTOCOL 11: Multi-objective studies are NOT supported by adaptive
|
||
|
|
characterization. Return immediately with max confidence to skip
|
||
|
|
characterization phase.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Dict with confidence scores:
|
||
|
|
- 'overall_confidence': 0-1 score, where 1 = high confidence
|
||
|
|
- 'convergence_score': How stable recent improvements are
|
||
|
|
- 'exploration_coverage': How well parameter space is covered
|
||
|
|
- 'prediction_stability': How consistent the model predictions are
|
||
|
|
"""
|
||
|
|
# [Protocol 11] Multi-objective NOT supported by adaptive characterization
|
||
|
|
is_multi_objective = len(study.directions) > 1
|
||
|
|
if is_multi_objective:
|
||
|
|
return {
|
||
|
|
'overall_confidence': 1.0, # Skip characterization
|
||
|
|
'convergence_score': 1.0,
|
||
|
|
'exploration_coverage': 1.0,
|
||
|
|
'prediction_stability': 1.0,
|
||
|
|
'ready_for_exploitation': True, # Go straight to NSGA-II
|
||
|
|
'total_trials': len(study.trials),
|
||
|
|
'message': '[Protocol 11] Multi-objective: skipping adaptive characterization, using NSGA-II directly'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Get ALL completed trials from study (study-aware)
|
||
|
|
all_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
|
||
|
|
|
||
|
|
if len(all_trials) < self.min_trials:
|
||
|
|
return {
|
||
|
|
'overall_confidence': 0.0,
|
||
|
|
'convergence_score': 0.0,
|
||
|
|
'exploration_coverage': 0.0,
|
||
|
|
'prediction_stability': 0.0,
|
||
|
|
'ready_for_exploitation': False,
|
||
|
|
'total_trials': len(all_trials),
|
||
|
|
'message': f'Need {self.min_trials - len(all_trials)} more trials for confidence assessment (currently {len(all_trials)} trials)'
|
||
|
|
}
|
||
|
|
|
||
|
|
# 1. Convergence Score - are we finding consistent improvements?
|
||
|
|
recent_window = 10
|
||
|
|
recent_trials = all_trials[-recent_window:]
|
||
|
|
recent_values = [t.value for t in recent_trials] # Safe: single-objective only
|
||
|
|
|
||
|
|
# Calculate improvement rate
|
||
|
|
improvements = []
|
||
|
|
for i in range(1, len(recent_values)):
|
||
|
|
if recent_values[i] < recent_values[i-1]:
|
||
|
|
improvement = (recent_values[i-1] - recent_values[i]) / abs(recent_values[i-1])
|
||
|
|
improvements.append(improvement)
|
||
|
|
|
||
|
|
# If we're consistently finding improvements, convergence is good
|
||
|
|
if improvements:
|
||
|
|
avg_improvement = np.mean(improvements)
|
||
|
|
improvement_consistency = 1.0 - variation(improvements) if len(improvements) > 1 else 0.5
|
||
|
|
convergence_score = min(1.0, avg_improvement * improvement_consistency * 10)
|
||
|
|
else:
|
||
|
|
convergence_score = 0.0
|
||
|
|
|
||
|
|
# 2. Exploration Coverage - how well have we covered parameter space?
|
||
|
|
# Use coefficient of variation for each parameter
|
||
|
|
param_names = list(all_trials[0].params.keys()) if all_trials else []
|
||
|
|
|
||
|
|
coverage_scores = []
|
||
|
|
for param in param_names:
|
||
|
|
values = [t.params[param] for t in all_trials]
|
||
|
|
|
||
|
|
# Get parameter bounds
|
||
|
|
distribution = all_trials[0].distributions[param]
|
||
|
|
param_range = distribution.high - distribution.low
|
||
|
|
|
||
|
|
# Calculate spread relative to bounds
|
||
|
|
spread = max(values) - min(values)
|
||
|
|
coverage = spread / param_range
|
||
|
|
coverage_scores.append(coverage)
|
||
|
|
|
||
|
|
exploration_coverage = np.mean(coverage_scores)
|
||
|
|
|
||
|
|
# 3. Prediction Stability - are recent trials clustered in good regions?
|
||
|
|
recent_best_values = []
|
||
|
|
current_best = float('inf')
|
||
|
|
for trial in recent_trials:
|
||
|
|
current_best = min(current_best, trial.value)
|
||
|
|
recent_best_values.append(current_best)
|
||
|
|
|
||
|
|
# If best hasn't improved much recently, model is stable
|
||
|
|
if len(recent_best_values) > 1:
|
||
|
|
best_stability = 1.0 - (recent_best_values[0] - recent_best_values[-1]) / (recent_best_values[0] + 1e-10)
|
||
|
|
prediction_stability = max(0.0, min(1.0, best_stability))
|
||
|
|
else:
|
||
|
|
prediction_stability = 0.0
|
||
|
|
|
||
|
|
# 4. Overall Confidence - weighted combination
|
||
|
|
overall_confidence = (
|
||
|
|
0.4 * convergence_score +
|
||
|
|
0.3 * exploration_coverage +
|
||
|
|
0.3 * prediction_stability
|
||
|
|
)
|
||
|
|
|
||
|
|
# Decision: Ready for intensive exploitation?
|
||
|
|
ready_for_exploitation = (
|
||
|
|
overall_confidence >= 0.65 and
|
||
|
|
exploration_coverage >= 0.5 and
|
||
|
|
len(all_trials) >= self.min_trials
|
||
|
|
)
|
||
|
|
|
||
|
|
message = self._get_confidence_message(overall_confidence, ready_for_exploitation)
|
||
|
|
|
||
|
|
return {
|
||
|
|
'overall_confidence': overall_confidence,
|
||
|
|
'convergence_score': convergence_score,
|
||
|
|
'exploration_coverage': exploration_coverage,
|
||
|
|
'prediction_stability': prediction_stability,
|
||
|
|
'ready_for_exploitation': ready_for_exploitation,
|
||
|
|
'total_trials': len(all_trials),
|
||
|
|
'message': message
|
||
|
|
}
|
||
|
|
|
||
|
|
def _get_confidence_message(self, confidence: float, ready: bool) -> str:
|
||
|
|
"""Generate human-readable confidence assessment."""
|
||
|
|
if ready:
|
||
|
|
return f"HIGH CONFIDENCE ({confidence:.1%}) - Transitioning to exploitation phase"
|
||
|
|
elif confidence >= 0.5:
|
||
|
|
return f"MEDIUM CONFIDENCE ({confidence:.1%}) - Continue exploration with some exploitation"
|
||
|
|
elif confidence >= 0.3:
|
||
|
|
return f"LOW CONFIDENCE ({confidence:.1%}) - Focus on exploration"
|
||
|
|
else:
|
||
|
|
return f"VERY LOW CONFIDENCE ({confidence:.1%}) - Need more diverse exploration"
|
||
|
|
|
||
|
|
|
||
|
|
class AdaptiveExploitationCallback:
|
||
|
|
"""
|
||
|
|
Dynamically adjust sampler behavior based on surrogate confidence.
|
||
|
|
|
||
|
|
This callback monitors surrogate model confidence and adapts the optimization
|
||
|
|
strategy from exploration-heavy to exploitation-heavy as confidence increases.
|
||
|
|
|
||
|
|
STUDY-AWARE: Tracks phase transitions across multiple optimization runs
|
||
|
|
and persists confidence history to JSON files.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
target_value: Optional[float] = None,
|
||
|
|
tolerance: float = 0.1,
|
||
|
|
min_confidence_for_exploitation: float = 0.65,
|
||
|
|
min_trials: int = 15,
|
||
|
|
verbose: bool = True,
|
||
|
|
tracking_dir: Optional[Path] = None
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Args:
|
||
|
|
target_value: Target objective value (if known)
|
||
|
|
tolerance: Acceptable error from target
|
||
|
|
min_confidence_for_exploitation: Confidence threshold to enable intensive exploitation
|
||
|
|
min_trials: Minimum trials before confidence assessment
|
||
|
|
verbose: Print confidence updates
|
||
|
|
tracking_dir: Directory to save phase transition tracking files
|
||
|
|
"""
|
||
|
|
self.target_value = target_value
|
||
|
|
self.tolerance = tolerance
|
||
|
|
self.min_confidence = min_confidence_for_exploitation
|
||
|
|
self.verbose = verbose
|
||
|
|
self.tracking_dir = tracking_dir
|
||
|
|
|
||
|
|
self.metrics = SurrogateConfidenceMetrics(min_trials_for_confidence=min_trials)
|
||
|
|
self.consecutive_successes = 0
|
||
|
|
|
||
|
|
# Initialize phase transition tracking
|
||
|
|
self.phase_transition_file = None
|
||
|
|
self.confidence_history_file = None
|
||
|
|
if tracking_dir:
|
||
|
|
self.tracking_dir = Path(tracking_dir)
|
||
|
|
self.tracking_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
self.phase_transition_file = self.tracking_dir / "phase_transitions.json"
|
||
|
|
self.confidence_history_file = self.tracking_dir / "confidence_history.json"
|
||
|
|
|
||
|
|
# Load existing phase transition data if available
|
||
|
|
self.phase_transitions = self._load_phase_transitions()
|
||
|
|
self.confidence_history = self._load_confidence_history()
|
||
|
|
|
||
|
|
# Determine current phase from history
|
||
|
|
self.phase = self._get_current_phase()
|
||
|
|
|
||
|
|
def _load_phase_transitions(self) -> List[Dict]:
|
||
|
|
"""Load existing phase transition history from JSON."""
|
||
|
|
if self.phase_transition_file and self.phase_transition_file.exists():
|
||
|
|
try:
|
||
|
|
with open(self.phase_transition_file, 'r') as f:
|
||
|
|
return json.load(f)
|
||
|
|
except Exception:
|
||
|
|
return []
|
||
|
|
return []
|
||
|
|
|
||
|
|
def _load_confidence_history(self) -> List[Dict]:
|
||
|
|
"""Load existing confidence history from JSON."""
|
||
|
|
if self.confidence_history_file and self.confidence_history_file.exists():
|
||
|
|
try:
|
||
|
|
with open(self.confidence_history_file, 'r') as f:
|
||
|
|
return json.load(f)
|
||
|
|
except Exception:
|
||
|
|
return []
|
||
|
|
return []
|
||
|
|
|
||
|
|
def _get_current_phase(self) -> str:
|
||
|
|
"""Determine current phase from transition history."""
|
||
|
|
if not self.phase_transitions:
|
||
|
|
return "exploration"
|
||
|
|
# If any transition to exploitation exists, we're in exploitation
|
||
|
|
for transition in self.phase_transitions:
|
||
|
|
if transition.get('to_phase') == 'exploitation':
|
||
|
|
return "exploitation"
|
||
|
|
return "exploration"
|
||
|
|
|
||
|
|
def _save_phase_transition(self, trial_number: int, confidence: Dict):
|
||
|
|
"""Save phase transition event to JSON."""
|
||
|
|
if not self.phase_transition_file:
|
||
|
|
return
|
||
|
|
|
||
|
|
transition_event = {
|
||
|
|
'trial_number': trial_number,
|
||
|
|
'from_phase': 'exploration',
|
||
|
|
'to_phase': 'exploitation',
|
||
|
|
'confidence_metrics': {
|
||
|
|
'overall_confidence': confidence['overall_confidence'],
|
||
|
|
'convergence_score': confidence['convergence_score'],
|
||
|
|
'exploration_coverage': confidence['exploration_coverage'],
|
||
|
|
'prediction_stability': confidence['prediction_stability']
|
||
|
|
},
|
||
|
|
'total_trials': confidence.get('total_trials', trial_number + 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
self.phase_transitions.append(transition_event)
|
||
|
|
|
||
|
|
try:
|
||
|
|
with open(self.phase_transition_file, 'w') as f:
|
||
|
|
json.dump(self.phase_transitions, f, indent=2)
|
||
|
|
except Exception as e:
|
||
|
|
if self.verbose:
|
||
|
|
print(f" Warning: Failed to save phase transition: {e}")
|
||
|
|
|
||
|
|
def _save_confidence_snapshot(self, trial_number: int, confidence: Dict):
|
||
|
|
"""Save confidence metrics snapshot to history."""
|
||
|
|
if not self.confidence_history_file:
|
||
|
|
return
|
||
|
|
|
||
|
|
snapshot = {
|
||
|
|
'trial_number': trial_number,
|
||
|
|
'phase': self.phase,
|
||
|
|
'confidence_metrics': {
|
||
|
|
'overall_confidence': confidence['overall_confidence'],
|
||
|
|
'convergence_score': confidence['convergence_score'],
|
||
|
|
'exploration_coverage': confidence['exploration_coverage'],
|
||
|
|
'prediction_stability': confidence['prediction_stability']
|
||
|
|
},
|
||
|
|
'total_trials': confidence.get('total_trials', trial_number + 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
self.confidence_history.append(snapshot)
|
||
|
|
|
||
|
|
try:
|
||
|
|
with open(self.confidence_history_file, 'w') as f:
|
||
|
|
json.dump(self.confidence_history, f, indent=2)
|
||
|
|
except Exception as e:
|
||
|
|
if self.verbose:
|
||
|
|
print(f" Warning: Failed to save confidence history: {e}")
|
||
|
|
|
||
|
|
def __call__(self, study: optuna.Study, trial: optuna.trial.FrozenTrial):
|
||
|
|
"""Called after each trial completes."""
|
||
|
|
# Skip failed trials
|
||
|
|
if trial.state != optuna.trial.TrialState.COMPLETE:
|
||
|
|
return
|
||
|
|
|
||
|
|
# Update metrics (no-op for study-aware implementation)
|
||
|
|
self.metrics.update(study, trial)
|
||
|
|
|
||
|
|
# Calculate confidence
|
||
|
|
confidence = self.metrics.calculate_confidence(study)
|
||
|
|
|
||
|
|
# Save confidence snapshot every 5 trials
|
||
|
|
if trial.number % 5 == 0:
|
||
|
|
self._save_confidence_snapshot(trial.number, confidence)
|
||
|
|
|
||
|
|
# Print confidence report
|
||
|
|
if self.verbose and trial.number % 5 == 0: # Every 5 trials
|
||
|
|
self._print_confidence_report(trial.number, confidence)
|
||
|
|
|
||
|
|
# Check for phase transition
|
||
|
|
if confidence['ready_for_exploitation'] and self.phase == "exploration":
|
||
|
|
self.phase = "exploitation"
|
||
|
|
|
||
|
|
# Save transition event
|
||
|
|
self._save_phase_transition(trial.number, confidence)
|
||
|
|
|
||
|
|
if self.verbose:
|
||
|
|
print(f"\n{'='*60}")
|
||
|
|
print(f" PHASE TRANSITION: EXPLORATION -> EXPLOITATION")
|
||
|
|
print(f" Trial #{trial.number}")
|
||
|
|
print(f" Surrogate confidence: {confidence['overall_confidence']:.1%}")
|
||
|
|
print(f" Now focusing on refining best regions")
|
||
|
|
print(f"{'='*60}\n")
|
||
|
|
|
||
|
|
# Check for target achievement
|
||
|
|
if self.target_value is not None and trial.value <= self.tolerance:
|
||
|
|
self.consecutive_successes += 1
|
||
|
|
|
||
|
|
if self.verbose:
|
||
|
|
print(f" [TARGET] Trial #{trial.number}: {trial.value:.6f} ≤ {self.tolerance:.6f}")
|
||
|
|
print(f" [TARGET] Consecutive successes: {self.consecutive_successes}/3")
|
||
|
|
|
||
|
|
# Stop after 3 consecutive successes in exploitation phase
|
||
|
|
if self.consecutive_successes >= 3 and self.phase == "exploitation":
|
||
|
|
if self.verbose:
|
||
|
|
print(f"\n{'='*60}")
|
||
|
|
print(f" TARGET ACHIEVED WITH HIGH CONFIDENCE")
|
||
|
|
print(f" Best value: {study.best_value:.6f}")
|
||
|
|
print(f" Stopping optimization")
|
||
|
|
print(f"{'='*60}\n")
|
||
|
|
study.stop()
|
||
|
|
else:
|
||
|
|
self.consecutive_successes = 0
|
||
|
|
|
||
|
|
def _print_confidence_report(self, trial_number: int, confidence: Dict):
|
||
|
|
"""Print confidence metrics report."""
|
||
|
|
print(f"\n [CONFIDENCE REPORT - Trial #{trial_number}]")
|
||
|
|
print(f" Phase: {self.phase.upper()}")
|
||
|
|
print(f" Overall Confidence: {confidence['overall_confidence']:.1%}")
|
||
|
|
print(f" - Convergence: {confidence['convergence_score']:.1%}")
|
||
|
|
print(f" - Coverage: {confidence['exploration_coverage']:.1%}")
|
||
|
|
print(f" - Stability: {confidence['prediction_stability']:.1%}")
|
||
|
|
print(f" {confidence['message']}")
|
||
|
|
print()
|
||
|
|
|
||
|
|
|
||
|
|
def create_adaptive_sampler(
|
||
|
|
n_startup_trials: int = 10,
|
||
|
|
multivariate: bool = True,
|
||
|
|
confidence_threshold: float = 0.65
|
||
|
|
) -> optuna.samplers.TPESampler:
|
||
|
|
"""
|
||
|
|
Create TPE sampler configured for adaptive exploration-exploitation.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
n_startup_trials: Initial random exploration trials
|
||
|
|
multivariate: Enable multivariate TPE for correlated parameters
|
||
|
|
confidence_threshold: Confidence needed before intensive exploitation
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Configured TPESampler
|
||
|
|
"""
|
||
|
|
# Higher n_ei_candidates = more exploitation
|
||
|
|
# Will be used once confidence threshold is reached
|
||
|
|
return optuna.samplers.TPESampler(
|
||
|
|
n_startup_trials=n_startup_trials,
|
||
|
|
n_ei_candidates=24,
|
||
|
|
multivariate=multivariate,
|
||
|
|
warn_independent_sampling=True
|
||
|
|
)
|