use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::numeric::{Float, NumCast};
use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::Mutex;
use std::time::Instant;
use super::types::*;
use crate::error::{OptimError, Result};
use crate::nas_engine::results::EvaluationResults;
use crate::{EvaluationConfig, EvaluationMetric, OptimizerArchitecture};
const FEATURE_DIM: usize = 12;
#[derive(Debug)]
pub struct PerformancePredictor<T: Float + Debug + Send + Sync + 'static> {
predictor_model: PredictorModel<T>,
training_data: PredictorTrainingData<T>,
confidence_level: T,
learning_rate: T,
ridge_lambda: T,
pending_features: Mutex<Vec<Array1<T>>>,
}
#[derive(Debug)]
pub struct PredictorModel<T: Float + Debug + Send + Sync + 'static> {
model_type: PredictorModelType,
parameters: ModelParameters<T>,
training_state: ModelTrainingState<T>,
}
#[derive(Debug)]
pub struct ModelParameters<T: Float + Debug + Send + Sync + 'static> {
weights: Vec<Array2<T>>,
}
#[derive(Debug)]
pub struct ModelTrainingState<T: Float + Debug + Send + Sync + 'static> {
current_epoch: usize,
loss_history: Vec<T>,
}
#[derive(Debug)]
pub struct PredictorTrainingData<T: Float + Debug + Send + Sync + 'static> {
architecture_features: Vec<Array1<T>>,
performance_targets: Vec<T>,
}
impl<T: Float + Debug + Default + Send + Sync> PerformancePredictor<T> {
const DEFAULT_CONFIDENCE_LEVEL: f64 = 0.95;
pub fn new(_config: &EvaluationConfig) -> Result<Self> {
Ok(Self {
predictor_model: PredictorModel::new(),
training_data: PredictorTrainingData::new(),
confidence_level: Self::scalar(Self::DEFAULT_CONFIDENCE_LEVEL),
learning_rate: Self::scalar(0.05),
ridge_lambda: Self::scalar(1e-4),
pending_features: Mutex::new(Vec::new()),
})
}
pub fn model_type(&self) -> &PredictorModelType {
&self.predictor_model.model_type
}
pub fn epochs_trained(&self) -> usize {
self.predictor_model.training_state.current_epoch
}
pub fn training_loss_history(&self) -> &[T] {
&self.predictor_model.training_state.loss_history
}
pub fn observation_count(&self) -> usize {
self.num_observations()
}
pub fn training_rmse(&self) -> Option<T> {
let targets = &self.training_data.performance_targets;
let features = &self.training_data.architecture_features;
let n = targets.len().min(features.len());
if n == 0 {
return None;
}
let mut sum_sq = T::zero();
for i in 0..n {
let error = Self::sigmoid(self.raw_score(&features[i])) - targets[i];
sum_sq = sum_sq + error * error;
}
let count: T = NumCast::from(n as f64)?;
Some((sum_sq / count).sqrt())
}
pub fn confidence_level(&self) -> T {
self.confidence_level
}
pub fn set_confidence_level(&mut self, level: T) -> Result<()> {
if !(level > T::zero() && level < T::one()) {
return Err(OptimError::InvalidParameter(format!(
"confidence level must lie strictly inside (0, 1), got {:?}",
level
)));
}
self.confidence_level = level;
Ok(())
}
fn scalar(value: f64) -> T {
NumCast::from(value).unwrap_or_else(T::zero)
}
fn sigmoid(raw: T) -> T {
let one = T::one();
if raw >= T::zero() {
let z = (-raw).exp();
one / (one + z)
} else {
let z = raw.exp();
z / (one + z)
}
}
fn clamp_unit(value: T) -> T {
value.max(T::zero()).min(T::one())
}
fn structure_signature(architecture: &OptimizerArchitecture<T>) -> T {
let mut entries: Vec<&str> = architecture.components.iter().map(|s| s.as_str()).collect();
entries.sort_unstable();
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
for entry in entries {
for byte in entry.as_bytes() {
hash ^= *byte as u64;
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
hash ^= 0x1f;
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
let bucket = (hash >> 40) as f64;
let max = ((1u64 << 24) - 1) as f64;
Self::scalar(bucket / max)
}
fn extract_features(&self, architecture: &OptimizerArchitecture<T>) -> Array1<T> {
let mut features = Array1::zeros(FEATURE_DIM);
features[0] = T::one();
let n_components = architecture.components.len();
features[1] = Self::saturating_count(n_components);
let mut distinct: Vec<&str> = architecture.components.iter().map(|s| s.as_str()).collect();
distinct.sort_unstable();
distinct.dedup();
features[2] = Self::saturating_count(distinct.len());
let numeric_values: Vec<f64> = architecture
.parameters
.values()
.chain(architecture.hyperparameters.values())
.map(|v| v.to_f64().unwrap_or(0.0))
.filter(|v| v.is_finite())
.collect();
features[3] = Self::saturating_count(numeric_values.len());
if numeric_values.is_empty() {
features[4] = T::zero();
features[5] = T::zero();
features[6] = T::zero();
features[7] = T::zero();
} else {
let mut sum = 0.0_f64;
let mut max = f64::NEG_INFINITY;
let mut min = f64::INFINITY;
for &v in &numeric_values {
sum += v;
if v > max {
max = v;
}
if v < min {
min = v;
}
}
let count = numeric_values.len() as f64;
features[4] = Self::scalar(sum / count);
features[5] = Self::scalar((sum.abs()).tanh()); features[6] = Self::scalar(max);
features[7] = Self::scalar(min);
}
features[8] = Self::structure_signature(architecture);
features[9] = Self::saturating_count(architecture.architecture_id.len());
if n_components == 0 {
features[10] = T::zero();
} else {
const KEYWORDS: [&str; 7] = [
"adam", "sgd", "momentum", "rmsprop", "adagrad", "adamw", "lamb",
];
let mut hits = 0usize;
for entry in &architecture.components {
let lower = entry.to_ascii_lowercase();
if KEYWORDS.iter().any(|kw| lower.contains(kw)) {
hits += 1;
}
}
features[10] = Self::scalar(hits as f64 / n_components as f64);
}
let lr = ["learning_rate", "lr"]
.iter()
.find_map(|key| {
architecture
.hyperparameters
.get(*key)
.or_else(|| architecture.parameters.get(*key))
})
.and_then(|v| v.to_f64())
.filter(|v| v.is_finite())
.unwrap_or(0.0);
features[11] = Self::scalar(lr);
features
}
fn saturating_count(n: usize) -> T {
let nf = n as f64;
Self::scalar(nf / (nf + 1.0))
}
fn raw_score(&self, features: &Array1<T>) -> T {
let weights = &self.predictor_model.parameters.weights[0];
let mut acc = T::zero();
for i in 0..FEATURE_DIM {
acc = acc + weights[[i, 0]] * features[i];
}
acc
}
fn num_observations(&self) -> usize {
self.training_data.performance_targets.len()
}
fn confidence_half_width(&self) -> T {
let n = self.num_observations();
let base = Self::scalar(0.25);
let denom = Self::scalar((1.0 + n as f64).sqrt());
let spread = base / denom;
let scale = T::one() + self.confidence_level;
spread * scale
}
pub fn predict_performance(
&self,
architecture: &OptimizerArchitecture<T>,
) -> Result<EvaluationResults<T>> {
let start_time = Instant::now();
let features = self.extract_features(architecture);
let raw = self.raw_score(&features);
let overall_score = Self::clamp_unit(Self::sigmoid(raw));
let half_width = self.confidence_half_width();
let lower = Self::clamp_unit(overall_score - half_width);
let upper = Self::clamp_unit(overall_score + half_width);
if let Ok(mut pending) = self.pending_features.lock() {
pending.push(features);
}
let mut metric_scores: HashMap<EvaluationMetric, T> = HashMap::new();
metric_scores.insert(EvaluationMetric::FinalPerformance, overall_score);
let mut confidence_intervals: HashMap<EvaluationMetric, (T, T)> = HashMap::new();
confidence_intervals.insert(EvaluationMetric::FinalPerformance, (lower, upper));
Ok(EvaluationResults {
metric_scores,
overall_score,
confidence_intervals,
evaluation_time: start_time.elapsed(),
success: true,
error_message: None,
cv_results: None,
benchmark_results: HashMap::new(),
training_trajectory: Vec::new(),
})
}
pub fn update_with_results(&mut self, results: &[EvaluationResults<T>]) -> Result<()> {
let mut pairs: Vec<(Array1<T>, T)> = Vec::new();
{
let mut pending = self.pending_features.lock().map_err(|_| {
OptimError::EvaluationError("predictor feature buffer poisoned".into())
})?;
for result in results {
if !result.success {
continue;
}
if pending.is_empty() {
break;
}
let features = pending.remove(0);
let target = Self::clamp_unit(result.overall_score);
pairs.push((features, target));
}
}
let lr = self.learning_rate;
let lambda = self.ridge_lambda;
for (features, target) in pairs {
let raw = self.raw_score(&features);
let prediction = Self::sigmoid(raw);
let error = prediction - target;
{
let weights = &mut self.predictor_model.parameters.weights[0];
for i in 0..FEATURE_DIM {
let grad = error * features[i] + lambda * weights[[i, 0]];
weights[[i, 0]] = weights[[i, 0]] - lr * grad;
}
}
let loss = error * error;
self.predictor_model.training_state.loss_history.push(loss);
self.predictor_model.training_state.current_epoch += 1;
self.training_data.architecture_features.push(features);
self.training_data.performance_targets.push(target);
}
Ok(())
}
}
impl<T: Float + Debug + Default + Send + Sync> PredictorModel<T> {
fn new() -> Self {
Self {
model_type: PredictorModelType::LinearRegression,
parameters: ModelParameters {
weights: vec![Array2::zeros((FEATURE_DIM, 1))],
},
training_state: ModelTrainingState {
current_epoch: 0,
loss_history: Vec::new(),
},
}
}
}
impl<T: Float + Debug + Default + Send + Sync> PredictorTrainingData<T> {
fn new() -> Self {
Self {
architecture_features: Vec::new(),
performance_targets: Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_config() -> EvaluationConfig {
EvaluationConfig {
epochs: 1,
batch_size: 8,
learning_rate: 0.01,
performance_prediction: true,
}
}
fn make_architecture(id: &str) -> OptimizerArchitecture<f64> {
let mut parameters = HashMap::new();
parameters.insert("learning_rate".to_string(), 0.01_f64);
parameters.insert("momentum".to_string(), 0.9_f64);
parameters.insert("weight_decay".to_string(), 0.0001_f64);
OptimizerArchitecture {
architecture_id: id.to_string(),
parameters,
components: vec![
"AdamW".to_string(),
"CosineSchedule".to_string(),
"GradientClipping".to_string(),
],
connections: Vec::new(),
metadata: HashMap::new(),
hyperparameters: HashMap::new(),
}
}
#[test]
fn test_feature_vector_is_fixed_length_and_deterministic() {
let predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-features");
let f1 = predictor.extract_features(&architecture);
let f2 = predictor.extract_features(&architecture);
assert_eq!(f1.len(), FEATURE_DIM);
for i in 0..FEATURE_DIM {
assert!((f1[i] - f2[i]).abs() < 1e-12);
}
assert!((f1[0] - 1.0).abs() < 1e-12);
for i in 0..FEATURE_DIM {
assert!(f1[i].is_finite());
}
}
#[test]
fn test_predict_returns_valid_results() {
let predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-predict");
let results = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
assert!(results.overall_score >= 0.0 && results.overall_score <= 1.0);
assert!(results.success);
assert!(results.error_message.is_none());
assert!(!results.metric_scores.is_empty());
let predicted = results
.metric_scores
.get(&EvaluationMetric::FinalPerformance)
.copied()
.expect("FinalPerformance metric should be present");
assert!((predicted - results.overall_score).abs() < 1e-12);
let (lo, hi) = results
.confidence_intervals
.get(&EvaluationMetric::FinalPerformance)
.copied()
.expect("confidence interval should be present");
assert!(lo >= 0.0 && hi <= 1.0);
assert!(lo <= hi);
assert!(results.cv_results.is_none());
assert!(results.benchmark_results.is_empty());
assert!(results.training_trajectory.is_empty());
}
#[test]
fn test_confidence_interval_narrows_with_more_data() {
let mut predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-uncertainty");
let wide = predictor.confidence_half_width();
for _ in 0..25 {
let _ = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
let result = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
predictor
.update_with_results(std::slice::from_ref(&result))
.expect("update should succeed");
}
let narrow = predictor.confidence_half_width();
assert!(narrow < wide);
}
#[test]
fn test_model_learns_toward_high_target() {
let mut predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-learn");
let baseline = predictor
.predict_performance(&architecture)
.expect("prediction should succeed")
.overall_score;
let high_target = 0.95_f64;
for _ in 0..200 {
let mut result = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
result.overall_score = high_target;
predictor
.update_with_results(std::slice::from_ref(&result))
.expect("update should succeed");
}
let learned = predictor
.predict_performance(&architecture)
.expect("prediction should succeed")
.overall_score;
assert!(
learned > baseline + 0.05,
"prediction did not increase: baseline={baseline}, learned={learned}"
);
assert!(
(high_target - learned).abs() < 0.25,
"prediction did not converge toward target: learned={learned}, target={high_target}"
);
assert!(predictor.num_observations() >= 200);
}
#[test]
fn test_update_skips_failed_results() {
let mut predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-failed");
let mut result = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
result.success = false;
result.error_message = Some("synthetic failure".to_string());
result.overall_score = 0.99;
predictor
.update_with_results(std::slice::from_ref(&result))
.expect("update should succeed");
assert_eq!(predictor.num_observations(), 0);
}
#[test]
fn test_recorded_training_state_is_readable_and_consistent() {
let mut predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
let architecture = make_architecture("arch-state");
assert_eq!(
predictor.model_type(),
&PredictorModelType::LinearRegression
);
assert_eq!(predictor.epochs_trained(), 0);
assert!(predictor.training_loss_history().is_empty());
assert_eq!(predictor.observation_count(), 0);
assert!(predictor.training_rmse().is_none());
for _ in 0..40 {
let mut result = predictor
.predict_performance(&architecture)
.expect("prediction should succeed");
result.overall_score = 0.9;
predictor
.update_with_results(std::slice::from_ref(&result))
.expect("update should succeed");
}
assert_eq!(predictor.epochs_trained(), 40);
assert_eq!(predictor.training_loss_history().len(), 40);
assert_eq!(predictor.observation_count(), 40);
let rmse = predictor
.training_rmse()
.expect("training RMSE is defined once pairs exist");
assert!(rmse.is_finite());
assert!((0.0..=1.0).contains(&rmse));
assert!(rmse < 0.4, "training did not reduce the fit error: {rmse}");
}
#[test]
fn test_confidence_level_is_validated_and_widens_the_interval() {
let mut predictor = PerformancePredictor::<f64>::new(&make_config())
.expect("predictor construction should succeed");
assert!((predictor.confidence_level() - 0.95).abs() < 1e-12);
let wide = predictor.confidence_half_width();
predictor
.set_confidence_level(0.5)
.expect("0.5 is a valid confidence level");
let narrow = predictor.confidence_half_width();
assert!(
narrow < wide,
"a lower confidence level must not widen the interval: {narrow} vs {wide}"
);
for bad in [0.0, 1.0, -0.5, 1.5, f64::NAN] {
assert!(
predictor.set_confidence_level(bad).is_err(),
"confidence level {bad} must be rejected"
);
}
assert!((predictor.confidence_level() - 0.5).abs() < 1e-12);
}
}