optirs_nas/nas_engine/engine/
support.rs

1//! Supporting traits (SearchStrategy, MultiObjectiveOptimizer, ArchitectureController) and the real evaluator/predictor/progressive-search implementations they are built from.
2
3use crate::error::Result;
4use crate::multi_objective;
5use crate::nas_engine::config::*;
6use crate::nas_engine::results::*;
7use scirs2_core::numeric::Float;
8use std::collections::{HashMap, VecDeque};
9use std::fmt::Debug;
10use std::sync::{Arc, Mutex};
11
12/// Diversity metrics for population
13#[derive(Debug, Clone)]
14pub struct DiversityMetrics<T: Float + Debug + Send + Sync + 'static> {
15    pub crowding_distance: Vec<T>,
16    pub entropy: T,
17    pub average_distance: T,
18    pub min_distance: T,
19    pub max_distance: T,
20}
21/// Performance evaluator for architectures.
22///
23/// This is the engine-side handle onto the real evaluation subsystem in
24/// [`crate::evaluation`]: every call to [`PerformanceEvaluator::evaluate`]
25/// instantiates the concrete optimizer described by the candidate architecture
26/// and actually runs it on the registered benchmark test functions. Earlier
27/// releases carried a stub here that returned a constant score, which made the
28/// whole search a no-op; that stub is gone.
29#[derive(Debug)]
30pub struct PerformanceEvaluator<T: Float + Debug + Send + Sync + 'static> {
31    pub(super) config: EvaluationConfig<T>,
32    /// Live evaluator performing benchmark execution, caching and statistics.
33    pub(super) inner: crate::evaluation::PerformanceEvaluator<T>,
34    /// Results returned so far, shared so cheap clones observe the same view.
35    pub(super) evaluation_cache: Arc<Mutex<HashMap<String, EvaluationResults<T>>>>,
36    pub(super) evaluation_count: usize,
37}
38impl<T: Float + Debug + Default + Clone + Send + Sync + 'static + std::iter::Sum>
39    PerformanceEvaluator<T>
40{
41    pub fn new(config: EvaluationConfig<T>) -> Result<Self> {
42        let evaluation_config = crate::EvaluationConfig::from_engine_config(&config);
43        let mut inner = crate::evaluation::PerformanceEvaluator::<T>::new(evaluation_config)?;
44        inner.initialize()?;
45        Ok(Self {
46            config,
47            inner,
48            evaluation_cache: Arc::new(Mutex::new(HashMap::new())),
49            evaluation_count: 0,
50        })
51    }
52    /// Evaluate a candidate architecture by actually running it.
53    ///
54    /// Delegates to [`crate::evaluation::PerformanceEvaluator::evaluate_architecture`],
55    /// which builds the optimizer the architecture describes and minimizes each
56    /// registered benchmark function with it. The achieved objectives determine
57    /// the returned scores, so two different architectures receive different
58    /// scores and an identical architecture reproduces exactly.
59    pub fn evaluate(
60        &mut self,
61        architecture: &OptimizerArchitecture<T>,
62    ) -> Result<EvaluationResults<T>> {
63        let results = self.inner.evaluate_architecture(architecture)?;
64        self.evaluation_count += 1;
65        match self.evaluation_cache.lock() {
66            Ok(mut cache) => {
67                cache.insert(architecture.architecture_id.clone(), results.clone());
68            }
69            Err(poisoned) => {
70                let mut cache = poisoned.into_inner();
71                cache.insert(architecture.architecture_id.clone(), results.clone());
72            }
73        }
74        Ok(results)
75    }
76    /// Number of architectures evaluated by this evaluator.
77    pub fn evaluation_count(&self) -> usize {
78        self.evaluation_count
79    }
80    /// Engine-level evaluation configuration in force.
81    pub fn config(&self) -> &EvaluationConfig<T> {
82        &self.config
83    }
84}
85/// Performance prediction system.
86///
87/// Wraps the real learned predictor in [`crate::evaluation::PerformancePredictor`]
88/// (a ridge-regularised linear model over a deterministic architecture feature
89/// vector, trained online from observed evaluations). Earlier releases returned
90/// a constant `0.6` here.
91#[derive(Debug)]
92pub struct PerformancePredictor<T: Float + Debug + Send + Sync + 'static> {
93    pub(super) model_type: PredictorType,
94    /// Live predictor performing feature extraction, scoring and online updates.
95    pub(super) inner: crate::evaluation::PerformancePredictor<T>,
96    pub(super) training_data: Vec<(OptimizerArchitecture<T>, EvaluationResults<T>)>,
97    pub(super) prediction_accuracy: T,
98    pub(super) confidence_threshold: T,
99}
100impl<T: Float + Debug + Default + Send + Sync + 'static> PerformancePredictor<T> {
101    pub fn new(config: &EvaluationConfig<T>) -> Result<Self> {
102        let evaluation_config = crate::EvaluationConfig::from_engine_config(config);
103        let inner = crate::evaluation::PerformancePredictor::<T>::new(&evaluation_config)?;
104        Ok(Self {
105            model_type: PredictorType::LinearRegression,
106            inner,
107            training_data: Vec::new(),
108            prediction_accuracy: T::zero(),
109            confidence_threshold: scirs2_core::numeric::NumCast::from(0.7)
110                .unwrap_or_else(|| T::zero()),
111        })
112    }
113    /// Predict the performance of an architecture without running it.
114    ///
115    /// Delegates to the learned model, which extracts a deterministic feature
116    /// vector from the architecture and scores it; the returned confidence
117    /// interval widens when little training data has been seen.
118    pub fn predict(
119        &mut self,
120        architecture: &OptimizerArchitecture<T>,
121    ) -> Result<EvaluationResults<T>> {
122        self.inner.predict_performance(architecture)
123    }
124    /// Feed observed evaluations back into the model.
125    ///
126    /// Each result performs one online gradient step on the predictor's
127    /// weights, and the running prediction accuracy (1 - mean absolute error
128    /// over the retained history) is refreshed.
129    pub fn update_training_data(&mut self, results: &[SearchResult<T>]) -> Result<()> {
130        let evaluations: Vec<EvaluationResults<T>> = results
131            .iter()
132            .map(|r| r.evaluation_results.clone())
133            .collect();
134        self.inner.update_with_results(&evaluations)?;
135        for result in results {
136            self.training_data.push((
137                result.architecture.clone(),
138                result.evaluation_results.clone(),
139            ));
140        }
141        const ACCURACY_WINDOW: usize = 64;
142        let window: Vec<&(OptimizerArchitecture<T>, EvaluationResults<T>)> = self
143            .training_data
144            .iter()
145            .rev()
146            .take(ACCURACY_WINDOW)
147            .collect();
148        if !window.is_empty() {
149            let mut error_sum = T::zero();
150            let mut counted = 0usize;
151            for (architecture, observed) in &window {
152                if let Ok(prediction) = self.inner.predict_performance(architecture) {
153                    let diff = prediction.overall_score - observed.overall_score;
154                    error_sum = error_sum + diff.abs();
155                    counted += 1;
156                }
157            }
158            if counted > 0 {
159                let count: T =
160                    scirs2_core::numeric::NumCast::from(counted).unwrap_or_else(|| T::one());
161                let mean_abs_error = error_sum / count;
162                self.prediction_accuracy = (T::one() - mean_abs_error).max(T::zero());
163            }
164        }
165        Ok(())
166    }
167    /// Current estimated prediction accuracy in `[0, 1]`.
168    pub fn prediction_accuracy(&self) -> T {
169        self.prediction_accuracy
170    }
171    /// Minimum accuracy at which the predictor should be trusted in place of a
172    /// full evaluation.
173    pub fn confidence_threshold(&self) -> T {
174        self.confidence_threshold
175    }
176    /// Kind of model backing this predictor.
177    pub fn model_type(&self) -> &PredictorType {
178        &self.model_type
179    }
180    /// Number of `(architecture, evaluation)` pairs observed so far.
181    pub fn training_sample_count(&self) -> usize {
182        self.training_data.len()
183    }
184}
185/// Types of predictors available
186#[derive(Debug, Clone, PartialEq, Eq)]
187pub enum PredictorType {
188    /// Ridge-regularised linear model over the architecture feature vector.
189    /// This is the model actually implemented by
190    /// [`crate::evaluation::PerformancePredictor`].
191    LinearRegression,
192    NeuralNetwork,
193    GaussianProcess,
194    RandomForest,
195    Ensemble,
196}
197/// Engine-side progressive search: a staged complexity schedule that constrains
198/// which candidates may be evaluated in each phase of the run.
199///
200/// This complements [`crate::search_strategies::ProgressiveNAS`], which *generates*
201/// progressively more complex architectures: whatever strategy the engine is
202/// configured with, this filter enforces the schedule on the candidates that
203/// actually reach the evaluator, so `NASConfig::progressive_search` means
204/// something for every strategy.
205///
206/// It used to mean nothing at all. `new` built an empty `stages` vector and
207/// `filter_candidates` was `Ok(candidates)` — the input returned verbatim — while
208/// `current_stage` stayed `0` and `stage_history` was never written. Enabling
209/// `progressive_search` therefore changed no behaviour whatsoever, even though the
210/// engine reports it among the run's key hyperparameters.
211#[derive(Debug)]
212pub struct ProgressiveNAS<T: Float + Debug + Send + Sync + 'static> {
213    pub(super) stages: Vec<ProgressiveStage<T>>,
214    pub(super) current_stage: usize,
215    pub(super) stage_history: Vec<Vec<SearchResult<T>>>,
216    /// Generations allotted to each stage, derived from the search budget.
217    pub(super) generations_per_stage: usize,
218}
219impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> ProgressiveNAS<T> {
220    /// Build the stage schedule from the configuration.
221    ///
222    /// One stage per allowed component count between
223    /// `SearchSpaceConfig::min_components` and `max_components`, each stage holding
224    /// a narrowed copy of the search space (its own `max_components`) so the stage
225    /// configuration describes exactly what that stage may explore. The generation
226    /// budget is split evenly across the stages, which is
227    /// [`TimeBudgetAllocation::Uniform`] applied to the generation axis.
228    pub fn new(config: &NASConfig<T>) -> Result<Self> {
229        let min_components = config.search_space.min_components.max(1);
230        let max_components = config.search_space.max_components.max(min_components);
231        let stage_count = max_components - min_components + 1;
232
233        let total_seconds = config
234            .resource_constraints
235            .time_constraints
236            .max_search_time
237            .as_secs_f64();
238        let per_stage_hours = total_seconds / 3600.0 / stage_count as f64;
239
240        let mut stages = Vec::with_capacity(stage_count);
241        for (index, complexity) in (min_components..=max_components).enumerate() {
242            let mut stage_search_space = config.search_space.clone();
243            stage_search_space.max_components = complexity;
244            let mut stage_config = config.clone();
245            stage_config.search_space = stage_search_space.clone();
246            stages.push(ProgressiveStage {
247                name: format!("stage_{index}_upto_{complexity}_components"),
248                search_space: stage_search_space,
249                duration_hours: scirs2_core::numeric::NumCast::from(per_stage_hours)
250                    .unwrap_or_else(T::zero),
251                transfer_knowledge: config.enable_transfer_learning,
252                stage_config,
253            });
254        }
255
256        let generations_per_stage = (config.search_budget / stage_count.max(1)).max(1);
257        Ok(Self {
258            stages,
259            current_stage: 0,
260            stage_history: vec![Vec::new(); stage_count],
261            generations_per_stage,
262        })
263    }
264
265    /// The stage `generation` falls into, clamped to the last stage once the
266    /// schedule is exhausted (the search may legitimately run longer than the
267    /// budget the stages were sized from).
268    pub fn stage_for_generation(&self, generation: usize) -> usize {
269        if self.stages.is_empty() {
270            return 0;
271        }
272        (generation / self.generations_per_stage.max(1)).min(self.stages.len() - 1)
273    }
274
275    /// Maximum component count allowed in the stage `generation` falls into.
276    pub fn complexity_limit(&self, generation: usize) -> Option<usize> {
277        self.stages
278            .get(self.stage_for_generation(generation))
279            .map(|stage| stage.search_space.max_components)
280    }
281
282    /// The stage schedule.
283    pub fn stages(&self) -> &[ProgressiveStage<T>] {
284        &self.stages
285    }
286
287    /// The stage the most recent `filter_candidates` call was in.
288    pub fn current_stage(&self) -> usize {
289        self.current_stage
290    }
291
292    /// Results recorded for `stage`.
293    pub fn stage_results(&self, stage: usize) -> &[SearchResult<T>] {
294        self.stage_history
295            .get(stage)
296            .map(|results| results.as_slice())
297            .unwrap_or(&[])
298    }
299
300    /// Restrict `candidates` to the complexity the current stage allows.
301    ///
302    /// Candidates over the limit are held back rather than evaluated early — that
303    /// is the whole point of a progressive schedule. If *every* candidate exceeds
304    /// the limit the simplest ones are kept anyway: returning an empty list would
305    /// stall the search, and a stalled generation is a worse answer than evaluating
306    /// the closest candidates available.
307    pub fn filter_candidates(
308        &mut self,
309        candidates: Vec<OptimizerArchitecture<T>>,
310        generation: usize,
311    ) -> Result<Vec<OptimizerArchitecture<T>>> {
312        if self.stages.is_empty() || candidates.is_empty() {
313            return Ok(candidates);
314        }
315        let stage = self.stage_for_generation(generation);
316        if stage != self.current_stage {
317            log::info!(
318                "progressive search entering {} at generation {}",
319                self.stages
320                    .get(stage)
321                    .map(|stage| stage.name.as_str())
322                    .unwrap_or("<unknown stage>"),
323                generation
324            );
325            self.current_stage = stage;
326        }
327        let limit = self
328            .stages
329            .get(stage)
330            .map(|stage| stage.search_space.max_components)
331            .unwrap_or(usize::MAX);
332
333        let (allowed, deferred): (Vec<_>, Vec<_>) = candidates
334            .into_iter()
335            .partition(|candidate| candidate.components.len() <= limit);
336        if !allowed.is_empty() {
337            if !deferred.is_empty() {
338                log::debug!(
339                    "progressive search deferred {} candidate(s) over the {}-component \
340                     limit of stage {}",
341                    deferred.len(),
342                    limit,
343                    stage
344                );
345            }
346            return Ok(allowed);
347        }
348
349        // Nothing fitted: keep the simplest candidates so the generation is not
350        // empty, and say so.
351        let mut fallback = deferred;
352        fallback.sort_by_key(|candidate| candidate.components.len());
353        let smallest = fallback
354            .first()
355            .map(|candidate| candidate.components.len())
356            .unwrap_or(0);
357        fallback.retain(|candidate| candidate.components.len() == smallest);
358        log::warn!(
359            "progressive search stage {} allows at most {} component(s) but every \
360             candidate exceeded it; evaluating the {} simplest ({} components each)",
361            stage,
362            limit,
363            fallback.len(),
364            smallest
365        );
366        Ok(fallback)
367    }
368
369    /// Record `results` against the stage they were produced in, so the stage
370    /// history describes what each stage actually achieved.
371    pub fn record_stage_results(&mut self, generation: usize, results: &[SearchResult<T>]) {
372        if results.is_empty() {
373            return;
374        }
375        let stage = self.stage_for_generation(generation);
376        if let Some(history) = self.stage_history.get_mut(stage) {
377            history.extend(results.iter().cloned());
378        }
379    }
380}
381/// Progressive search stage
382#[derive(Debug, Clone)]
383pub struct ProgressiveStage<T: Float + Debug + Send + Sync + 'static> {
384    pub name: String,
385    pub search_space: SearchSpaceConfig,
386    pub duration_hours: T,
387    pub transfer_knowledge: bool,
388    pub stage_config: NASConfig<T>,
389}
390/// Core search strategy trait
391pub trait SearchStrategy<T: Float + Debug + Send + Sync + 'static>: Send + Sync {
392    /// Generate new candidate architectures
393    fn generate_candidates(
394        &mut self,
395        history: &VecDeque<SearchResult<T>>,
396    ) -> Result<Vec<OptimizerArchitecture<T>>>;
397
398    /// Update strategy based on search results
399    fn update_strategy(&mut self, results: &[SearchResult<T>]) -> Result<()>;
400
401    /// Check if strategy has converged
402    fn has_converged(&self) -> bool;
403
404    /// Get strategy name
405    fn strategy_name(&self) -> &str;
406}
407
408/// Multi-objective optimization trait
409pub trait MultiObjectiveOptimizer<T: Float + Debug + Send + Sync + 'static>: Send + Sync {
410    /// Update Pareto front with new results
411    fn update_pareto_front(
412        &mut self,
413        results: &[SearchResult<T>],
414    ) -> Result<multi_objective::ParetoFront<T>>;
415
416    /// Select next generation candidates
417    fn select_candidates(
418        &self,
419        candidates: &[SearchResult<T>],
420        population_size: usize,
421    ) -> Result<Vec<SearchResult<T>>>;
422
423    /// Calculate diversity metrics
424    fn calculate_diversity(&self, population: &[SearchResult<T>]) -> f64;
425}
426
427/// Architecture controller trait
428pub trait ArchitectureController<T: Float + Debug + Send + Sync + 'static>: Send + Sync {
429    /// Generate random architecture
430    fn generate_random(&mut self) -> Result<OptimizerArchitecture<T>>;
431
432    /// Mutate existing architecture
433    fn mutate(
434        &mut self,
435        architecture: &OptimizerArchitecture<T>,
436    ) -> Result<OptimizerArchitecture<T>>;
437
438    /// Crossover two architectures
439    fn crossover(
440        &mut self,
441        parent1: &OptimizerArchitecture<T>,
442        parent2: &OptimizerArchitecture<T>,
443    ) -> Result<OptimizerArchitecture<T>>;
444
445    /// Validate architecture
446    fn validate(&self, architecture: &OptimizerArchitecture<T>) -> Result<bool>;
447}
optirs_nas/nas_engine/engine/support.rs

optirs_nas/nas_engine/engine/
support.rs
