1use scirs2_core::ndarray::{Array1, Array2, ArrayView2};
8use sklears_core::{
9 error::{Result as SklResult, SklearsError},
10 traits::{Estimator, Fit, Predict, Untrained},
11 types::Float,
12};
13use std::collections::HashMap;
14
15#[derive(Debug, Clone)]
21pub struct EarlyStoppingConfig {
22 pub min_delta: Float,
24 pub patience: usize,
26 pub monitor: String,
28 pub mode_max: bool,
30 pub restore_best_weights: bool,
32}
33
34impl Default for EarlyStoppingConfig {
35 fn default() -> Self {
36 Self {
37 min_delta: 1e-4,
38 patience: 10,
39 monitor: "loss".to_string(),
40 mode_max: false,
41 restore_best_weights: true,
42 }
43 }
44}
45
46#[derive(Debug, Clone)]
48pub struct EarlyStopping {
49 config: EarlyStoppingConfig,
50 best_value: Option<Float>,
51 best_iteration: usize,
52 wait_count: usize,
53 should_stop: bool,
54}
55
56impl EarlyStopping {
57 pub fn new(config: EarlyStoppingConfig) -> Self {
59 Self {
60 config,
61 best_value: None,
62 best_iteration: 0,
63 wait_count: 0,
64 should_stop: false,
65 }
66 }
67
68 pub fn update(&mut self, value: Float, iteration: usize) -> bool {
70 match self.best_value {
71 None => {
72 self.best_value = Some(value);
73 self.best_iteration = iteration;
74 false
75 }
76 Some(best) => {
77 let is_improvement = if self.config.mode_max {
78 value > best + self.config.min_delta
79 } else {
80 value < best - self.config.min_delta
81 };
82
83 if is_improvement {
84 self.best_value = Some(value);
85 self.best_iteration = iteration;
86 self.wait_count = 0;
87 false
88 } else {
89 self.wait_count += 1;
90 if self.wait_count >= self.config.patience {
91 self.should_stop = true;
92 true
93 } else {
94 false
95 }
96 }
97 }
98 }
99 }
100
101 pub fn should_stop(&self) -> bool {
103 self.should_stop
104 }
105
106 pub fn best_value(&self) -> Option<Float> {
108 self.best_value
109 }
110
111 pub fn best_iteration(&self) -> usize {
113 self.best_iteration
114 }
115}
116
117#[derive(Debug, Clone)]
123pub struct WarmStartRegressorConfig {
124 pub max_iter: usize,
126 pub learning_rate: Float,
128 pub alpha: Float,
130 pub tol: Float,
132 pub early_stopping: Option<EarlyStoppingConfig>,
134 pub verbose: bool,
136}
137
138impl Default for WarmStartRegressorConfig {
139 fn default() -> Self {
140 Self {
141 max_iter: 1000,
142 learning_rate: 0.01,
143 alpha: 0.0001,
144 tol: 1e-4,
145 early_stopping: Some(EarlyStoppingConfig::default()),
146 verbose: false,
147 }
148 }
149}
150
151#[derive(Debug, Clone)]
180pub struct WarmStartRegressor<S = Untrained> {
181 state: S,
182 config: WarmStartRegressorConfig,
183}
184
185#[derive(Debug, Clone)]
187pub struct WarmStartRegressorTrained {
188 pub coef: Array2<Float>,
190 pub intercept: Array1<Float>,
192 pub n_features: usize,
194 pub n_outputs: usize,
196 pub n_iter: usize,
198 pub loss_history: Vec<Float>,
200 pub best_loss: Float,
202 pub best_iter: usize,
204 pub best_coef: Option<Array2<Float>>,
206 pub best_intercept: Option<Array1<Float>>,
208 pub converged: bool,
210 pub config: WarmStartRegressorConfig,
212}
213
214impl WarmStartRegressor<Untrained> {
215 pub fn new() -> Self {
217 Self {
218 state: Untrained,
219 config: WarmStartRegressorConfig::default(),
220 }
221 }
222
223 pub fn config(mut self, config: WarmStartRegressorConfig) -> Self {
225 self.config = config;
226 self
227 }
228
229 pub fn max_iter(mut self, max_iter: usize) -> Self {
231 self.config.max_iter = max_iter;
232 self
233 }
234
235 pub fn learning_rate(mut self, lr: Float) -> Self {
237 self.config.learning_rate = lr;
238 self
239 }
240
241 pub fn early_stopping(mut self, config: EarlyStoppingConfig) -> Self {
243 self.config.early_stopping = Some(config);
244 self
245 }
246}
247
248impl Default for WarmStartRegressor<Untrained> {
249 fn default() -> Self {
250 Self::new()
251 }
252}
253
254impl Fit<ArrayView2<'_, Float>, ArrayView2<'_, Float>> for WarmStartRegressor<Untrained> {
255 type Fitted = WarmStartRegressor<WarmStartRegressorTrained>;
256
257 #[allow(non_snake_case)] fn fit(self, X: &ArrayView2<Float>, y: &ArrayView2<Float>) -> SklResult<Self::Fitted> {
259 if X.nrows() != y.nrows() {
260 return Err(SklearsError::InvalidInput(
261 "Number of samples in X and y must match".to_string(),
262 ));
263 }
264
265 let n_samples = X.nrows();
266 let n_features = X.ncols();
267 let n_outputs = y.ncols();
268
269 let mut coef = Array2::zeros((n_features, n_outputs));
271 let mut intercept = Array1::zeros(n_outputs);
272
273 let mut loss_history = Vec::new();
274 let mut best_loss = Float::INFINITY;
275 let mut best_iter = 0;
276 let mut best_coef = None;
277 let mut best_intercept = None;
278
279 let mut early_stopping = self
280 .config
281 .early_stopping
282 .as_ref()
283 .map(|cfg| EarlyStopping::new(cfg.clone()));
284
285 let mut converged = false;
286
287 for iter in 0..self.config.max_iter {
289 let mut total_loss = 0.0;
290
291 for i in 0..n_samples {
293 let x_i = X.row(i);
294 let y_i = y.row(i);
295
296 let pred = coef.t().dot(&x_i) + &intercept;
298
299 let error = &y_i - &pred;
301 total_loss += error.mapv(|x| x.powi(2)).sum();
302
303 for j in 0..n_features {
305 for k in 0..n_outputs {
306 let gradient = -error[k] * x_i[j] + self.config.alpha * coef[[j, k]];
307 coef[[j, k]] -= self.config.learning_rate * gradient;
308 }
309 }
310
311 for k in 0..n_outputs {
313 intercept[k] += self.config.learning_rate * error[k];
314 }
315 }
316
317 let avg_loss = total_loss / (n_samples as Float * n_outputs as Float);
319 loss_history.push(avg_loss);
320
321 if avg_loss < best_loss {
323 best_loss = avg_loss;
324 best_iter = iter;
325 if self.config.early_stopping.is_some() {
326 best_coef = Some(coef.clone());
327 best_intercept = Some(intercept.clone());
328 }
329 }
330
331 if iter > 0 && (loss_history[iter - 1] - avg_loss).abs() < self.config.tol {
333 converged = true;
334 if self.config.verbose {
335 println!("Converged at iteration {}", iter);
336 }
337 break;
338 }
339
340 if let Some(ref mut es) = early_stopping {
342 if es.update(avg_loss, iter) {
343 if self.config.verbose {
344 println!("Early stopping at iteration {}", iter);
345 }
346 break;
347 }
348 }
349
350 if self.config.verbose && iter % 100 == 0 {
351 println!("Iteration {}: loss = {:.6}", iter, avg_loss);
352 }
353 }
354
355 if let Some(cfg) = &self.config.early_stopping {
357 if cfg.restore_best_weights {
358 if let Some(ref best_c) = best_coef {
359 coef = best_c.clone();
360 }
361 if let Some(ref best_i) = best_intercept {
362 intercept = best_i.clone();
363 }
364 }
365 }
366
367 Ok(WarmStartRegressor {
368 state: WarmStartRegressorTrained {
369 coef,
370 intercept,
371 n_features,
372 n_outputs,
373 n_iter: loss_history.len(),
374 loss_history,
375 best_loss,
376 best_iter,
377 best_coef,
378 best_intercept,
379 converged,
380 config: self.config,
381 },
382 config: WarmStartRegressorConfig::default(),
383 })
384 }
385}
386
387impl WarmStartRegressor<WarmStartRegressorTrained> {
388 #[allow(non_snake_case)] pub fn continue_training(
391 mut self,
392 X: &ArrayView2<Float>,
393 y: &ArrayView2<Float>,
394 additional_iterations: usize,
395 ) -> SklResult<Self> {
396 if X.nrows() != y.nrows() {
397 return Err(SklearsError::InvalidInput(
398 "Number of samples in X and y must match".to_string(),
399 ));
400 }
401
402 if X.ncols() != self.state.n_features || y.ncols() != self.state.n_outputs {
403 return Err(SklearsError::InvalidInput(
404 "Feature or output dimensions do not match".to_string(),
405 ));
406 }
407
408 let n_samples = X.nrows();
409
410 let mut early_stopping = self
411 .state
412 .config
413 .early_stopping
414 .as_ref()
415 .map(|cfg| EarlyStopping::new(cfg.clone()));
416
417 for iter in 0..additional_iterations {
419 let mut total_loss = 0.0;
420
421 for i in 0..n_samples {
423 let x_i = X.row(i);
424 let y_i = y.row(i);
425
426 let pred = self.state.coef.t().dot(&x_i) + &self.state.intercept;
427 let error = &y_i - &pred;
428 total_loss += error.mapv(|x| x.powi(2)).sum();
429
430 for j in 0..self.state.n_features {
432 for k in 0..self.state.n_outputs {
433 let gradient =
434 -error[k] * x_i[j] + self.state.config.alpha * self.state.coef[[j, k]];
435 self.state.coef[[j, k]] -= self.state.config.learning_rate * gradient;
436 }
437 }
438
439 for k in 0..self.state.n_outputs {
441 self.state.intercept[k] += self.state.config.learning_rate * error[k];
442 }
443 }
444
445 let avg_loss = total_loss / (n_samples as Float * self.state.n_outputs as Float);
446 self.state.loss_history.push(avg_loss);
447
448 if avg_loss < self.state.best_loss {
450 self.state.best_loss = avg_loss;
451 self.state.best_iter = self.state.n_iter + iter;
452 if self.state.config.early_stopping.is_some() {
453 self.state.best_coef = Some(self.state.coef.clone());
454 self.state.best_intercept = Some(self.state.intercept.clone());
455 }
456 }
457
458 let loss_len = self.state.loss_history.len();
460 if loss_len > 1 {
461 let prev_loss = self.state.loss_history[loss_len - 2];
462 if (prev_loss - avg_loss).abs() < self.state.config.tol {
463 self.state.converged = true;
464 break;
465 }
466 }
467
468 if let Some(ref mut es) = early_stopping {
470 if es.update(avg_loss, self.state.n_iter + iter) {
471 break;
472 }
473 }
474 }
475
476 self.state.n_iter += additional_iterations;
477 Ok(self)
478 }
479
480 pub fn loss_history(&self) -> &[Float] {
482 &self.state.loss_history
483 }
484
485 pub fn best_loss(&self) -> Float {
487 self.state.best_loss
488 }
489
490 pub fn converged(&self) -> bool {
492 self.state.converged
493 }
494
495 pub fn coef(&self) -> &Array2<Float> {
497 &self.state.coef
498 }
499
500 pub fn n_iter(&self) -> usize {
502 self.state.n_iter
503 }
504}
505
506impl Predict<ArrayView2<'_, Float>, Array2<Float>>
507 for WarmStartRegressor<WarmStartRegressorTrained>
508{
509 #[allow(non_snake_case)] fn predict(&self, X: &ArrayView2<Float>) -> SklResult<Array2<Float>> {
511 if X.ncols() != self.state.n_features {
512 return Err(SklearsError::InvalidInput(format!(
513 "Expected {} features, got {}",
514 self.state.n_features,
515 X.ncols()
516 )));
517 }
518
519 let n_samples = X.nrows();
520 let mut predictions = Array2::zeros((n_samples, self.state.n_outputs));
521
522 for i in 0..n_samples {
523 let x_i = X.row(i);
524 let pred = self.state.coef.t().dot(&x_i) + &self.state.intercept;
525 predictions.row_mut(i).assign(&pred);
526 }
527
528 Ok(predictions)
529 }
530}
531
532impl Estimator for WarmStartRegressor<Untrained> {
533 type Config = WarmStartRegressorConfig;
534 type Error = SklearsError;
535 type Float = Float;
536
537 fn config(&self) -> &Self::Config {
538 &self.config
539 }
540}
541
542impl Estimator for WarmStartRegressor<WarmStartRegressorTrained> {
543 type Config = WarmStartRegressorConfig;
544 type Error = SklearsError;
545 type Float = Float;
546
547 fn config(&self) -> &Self::Config {
548 &self.state.config
549 }
550}
551
552#[derive(Debug, Clone)]
558pub struct PredictionCache {
559 cache: HashMap<u64, Array2<Float>>,
561 max_size: usize,
563 hits: usize,
565 misses: usize,
567}
568
569impl PredictionCache {
570 pub fn new(max_size: usize) -> Self {
572 Self {
573 cache: HashMap::new(),
574 max_size,
575 hits: 0,
576 misses: 0,
577 }
578 }
579
580 #[allow(non_snake_case)] pub fn get(&mut self, X: &ArrayView2<Float>) -> Option<Array2<Float>> {
583 let hash = self.hash_input(X);
584 if let Some(pred) = self.cache.get(&hash) {
585 self.hits += 1;
586 Some(pred.clone())
587 } else {
588 self.misses += 1;
589 None
590 }
591 }
592
593 #[allow(non_snake_case)] pub fn put(&mut self, X: &ArrayView2<Float>, prediction: Array2<Float>) {
596 if self.cache.len() >= self.max_size {
597 if let Some(first_key) = self.cache.keys().next().copied() {
599 self.cache.remove(&first_key);
600 }
601 }
602 let hash = self.hash_input(X);
603 self.cache.insert(hash, prediction);
604 }
605
606 pub fn clear(&mut self) {
608 self.cache.clear();
609 }
610
611 pub fn stats(&self) -> (usize, usize, Float) {
613 let total = self.hits + self.misses;
614 let hit_rate = if total > 0 {
615 self.hits as Float / total as Float
616 } else {
617 0.0
618 };
619 (self.hits, self.misses, hit_rate)
620 }
621
622 #[allow(non_snake_case)] fn hash_input(&self, X: &ArrayView2<Float>) -> u64 {
625 use std::collections::hash_map::DefaultHasher;
626 use std::hash::{Hash, Hasher};
627
628 let mut hasher = DefaultHasher::new();
629 for &val in X.iter() {
630 val.to_bits().hash(&mut hasher);
631 }
632 hasher.finish()
633 }
634}
635
636#[cfg(test)]
641#[allow(non_snake_case)] mod tests {
643 use super::*;
644 use approx::assert_abs_diff_eq;
645 use scirs2_core::ndarray::array;
647
648 #[test]
649 fn test_early_stopping_basic() {
650 let config = EarlyStoppingConfig {
651 min_delta: 0.1,
652 patience: 3,
653 mode_max: false,
654 ..Default::default()
655 };
656
657 let mut es = EarlyStopping::new(config);
658
659 assert!(!es.update(1.0, 0));
660 assert!(!es.update(0.8, 1)); assert!(!es.update(0.79, 2)); assert!(!es.update(0.78, 3)); assert!(es.update(0.77, 4)); }
665
666 #[test]
667 fn test_early_stopping_mode_max() {
668 let config = EarlyStoppingConfig {
669 min_delta: 0.01,
670 patience: 2,
671 mode_max: true,
672 ..Default::default()
673 };
674
675 let mut es = EarlyStopping::new(config);
676
677 assert!(!es.update(0.5, 0));
678 assert!(!es.update(0.6, 1)); assert!(!es.update(0.59, 2)); assert!(es.update(0.58, 3)); }
682
683 #[test]
684 #[allow(non_snake_case)]
685 fn test_warm_start_regressor_basic() {
686 let X = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];
687 let y = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];
688
689 let model = WarmStartRegressor::new().max_iter(100).learning_rate(0.1);
690
691 let trained = model
692 .fit(&X.view(), &y.view())
693 .expect("model fitting should succeed");
694 let predictions = trained
695 .predict(&X.view())
696 .expect("prediction should succeed");
697
698 assert_eq!(predictions.dim(), (3, 2));
699 assert!(trained.n_iter() > 0);
700 }
701
702 #[test]
703 #[allow(non_snake_case)]
704 fn test_warm_start_continue_training() {
705 let X = array![[1.0, 2.0], [2.0, 3.0]];
706 let y = array![[1.0, 2.0], [2.0, 3.0]];
707
708 let model = WarmStartRegressor::new().max_iter(10).learning_rate(0.1);
709
710 let trained = model
711 .fit(&X.view(), &y.view())
712 .expect("model fitting should succeed");
713 let initial_iter = trained.n_iter();
714 let initial_loss = trained
715 .loss_history()
716 .last()
717 .copied()
718 .expect("collection should not be empty");
719
720 let continued = trained
722 .continue_training(&X.view(), &y.view(), 20)
723 .expect("operation should succeed");
724 let final_loss = continued
725 .loss_history()
726 .last()
727 .copied()
728 .expect("collection should not be empty");
729
730 assert!(continued.n_iter() > initial_iter);
731 assert!(final_loss <= initial_loss + 1.0); }
734
735 #[test]
736 #[allow(non_snake_case)]
737 fn test_warm_start_with_early_stopping() {
738 let X = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];
739 let y = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];
740
741 let es_config = EarlyStoppingConfig {
742 patience: 5,
743 min_delta: 1e-6,
744 ..Default::default()
745 };
746
747 let model = WarmStartRegressor::new()
748 .max_iter(1000)
749 .early_stopping(es_config)
750 .learning_rate(0.1);
751
752 let trained = model
753 .fit(&X.view(), &y.view())
754 .expect("model fitting should succeed");
755
756 assert!(trained.n_iter() < 1000);
758 assert!(trained.best_loss() < Float::INFINITY);
759 }
760
761 #[test]
762 fn test_prediction_cache_basic() {
763 let mut cache = PredictionCache::new(10);
764
765 let X = array![[1.0, 2.0], [2.0, 3.0]];
766 let pred = array![[1.0, 2.0], [2.0, 3.0]];
767
768 assert!(cache.get(&X.view()).is_none());
770
771 cache.put(&X.view(), pred.clone());
773 let cached = cache.get(&X.view()).expect("index should be valid");
774
775 assert_eq!(cached.dim(), pred.dim());
776 assert_eq!(cache.stats().0, 1); assert_eq!(cache.stats().1, 1); }
779
780 #[test]
781 fn test_prediction_cache_eviction() {
782 let mut cache = PredictionCache::new(2);
783
784 let X1 = array![[1.0, 2.0]];
785 let X2 = array![[2.0, 3.0]];
786 let X3 = array![[3.0, 4.0]];
787 let pred = array![[1.0, 2.0]];
788
789 cache.put(&X1.view(), pred.clone());
790 cache.put(&X2.view(), pred.clone());
791 cache.put(&X3.view(), pred.clone()); assert_eq!(cache.cache.len(), 2);
794 }
795
796 #[test]
797 fn test_cache_stats() {
798 let mut cache = PredictionCache::new(10);
799
800 let X = array![[1.0, 2.0]];
801 let pred = array![[1.0, 2.0]];
802
803 cache.get(&X.view()); cache.put(&X.view(), pred);
805 cache.get(&X.view()); cache.get(&X.view()); let (hits, misses, hit_rate) = cache.stats();
809 assert_eq!(hits, 2);
810 assert_eq!(misses, 1);
811 assert_abs_diff_eq!(hit_rate, 2.0 / 3.0, epsilon = 1e-6);
812 }
813}