1use crate::error::{StatsError, StatsResult};
4use crate::regression::stat_tests::{f_test_p_value, t_test_p_value};
5use crate::regression::utils::*;
6use crate::regression::RegressionResults;
7use scirs2_core::ndarray::{s, Array1, Array2, ArrayView1, ArrayView2};
8use scirs2_core::numeric::Float;
9use scirs2_linalg::lstsq;
10use std::collections::HashSet;
11
12#[derive(Debug, Clone, Copy, PartialEq)]
14pub enum StepwiseDirection {
15 Forward,
17 Backward,
19 Both,
21}
22
23#[derive(Debug, Clone, Copy)]
25pub enum StepwiseCriterion {
26 AIC,
28 BIC,
30 AdjR2,
32 F,
34 T,
36}
37
38pub struct StepwiseResults<F>
40where
41 F: Float + std::fmt::Debug + std::fmt::Display + 'static,
42{
43 pub final_model: RegressionResults<F>,
45
46 pub selected_indices: Vec<usize>,
48
49 pub sequence: Vec<(usize, bool)>, pub criteria_values: Vec<F>,
54}
55
56impl<F> StepwiseResults<F>
57where
58 F: Float + std::fmt::Debug + std::fmt::Display + 'static,
59{
60 pub fn summary(&self) -> String {
62 let mut summary = String::new();
63
64 summary.push_str("=== Stepwise Regression Results ===\n\n");
65
66 summary.push_str("Selected variables: ");
68 for (i, &idx) in self.selected_indices.iter().enumerate() {
69 if i > 0 {
70 summary.push_str(", ");
71 }
72 summary.push_str(&format!("X{}", idx));
73 }
74 summary.push_str("\n\n");
75
76 summary.push_str("Sequence of variable entry/exit:\n");
78 for (i, &(idx, is_entry)) in self.sequence.iter().enumerate() {
79 summary.push_str(&format!(
80 "Step {}: {} X{} (criterion value: {})\n",
81 i + 1,
82 if is_entry { "Added" } else { "Removed" },
83 idx,
84 self.criteria_values[i]
85 ));
86 }
87 summary.push('\n');
88
89 summary.push_str("Final Model:\n");
91 summary.push_str(&self.final_model.summary());
92
93 summary
94 }
95}
96
97#[allow(clippy::too_many_arguments)]
155#[allow(dead_code)]
156pub fn stepwise_regression<F>(
157 x: &ArrayView2<F>,
158 y: &ArrayView1<F>,
159 direction: StepwiseDirection,
160 criterion: StepwiseCriterion,
161 p_enter: Option<F>,
162 p_remove: Option<F>,
163 max_steps: Option<usize>,
164 include_intercept: bool,
165) -> StatsResult<StepwiseResults<F>>
166where
167 F: Float
168 + std::iter::Sum<F>
169 + std::ops::Div<Output = F>
170 + std::fmt::Debug
171 + std::fmt::Display
172 + 'static
173 + scirs2_core::numeric::NumAssign
174 + scirs2_core::numeric::One
175 + scirs2_core::ndarray::ScalarOperand
176 + Send
177 + Sync,
178{
179 if x.nrows() != y.len() {
181 return Err(StatsError::DimensionMismatch(format!(
182 "Input x has {} rows but y has length {}",
183 x.nrows(),
184 y.len()
185 )));
186 }
187
188 let n = x.nrows();
189 let p = x.ncols();
190
191 if n < 3 {
193 return Err(StatsError::InvalidArgument(
194 "At least 3 observations required for stepwise regression".to_string(),
195 ));
196 }
197
198 let p_enter =
200 p_enter.unwrap_or_else(|| F::from(0.05).expect("Failed to convert constant to float"));
201 let p_remove =
202 p_remove.unwrap_or_else(|| F::from(0.1).expect("Failed to convert constant to float"));
203
204 let max_steps = max_steps.unwrap_or(p * 2);
206
207 let mut selected_indices = match direction {
209 StepwiseDirection::Forward => HashSet::new(),
210 StepwiseDirection::Backward | StepwiseDirection::Both => {
211 let mut indices = HashSet::new();
213 for i in 0..p {
214 indices.insert(i);
215 }
216 indices
217 }
218 };
219
220 let mut sequence = Vec::new();
222 let mut criteria_values = Vec::new();
223
224 let mut current_x = match direction {
226 StepwiseDirection::Forward => {
227 if include_intercept {
229 Array2::<F>::ones((n, 1))
230 } else {
231 Array2::<F>::zeros((n, 0))
232 }
233 }
234 StepwiseDirection::Backward | StepwiseDirection::Both => {
235 if include_intercept {
237 let mut x_full = Array2::<F>::zeros((n, p + 1));
238 x_full.slice_mut(s![.., 0]).fill(F::one());
239 for i in 0..p {
240 x_full.slice_mut(s![.., i + 1]).assign(&x.slice(s![.., i]));
241 }
242 x_full
243 } else {
244 x.to_owned()
245 }
246 }
247 };
248
249 let mut step = 0;
251 let mut criterion_improved = true;
252
253 while step < max_steps && criterion_improved {
254 criterion_improved = false;
255
256 if direction == StepwiseDirection::Forward || direction == StepwiseDirection::Both {
258 let mut best_var = None;
260 let mut best_criterion = F::infinity();
261
262 for i in 0..p {
263 if selected_indices.contains(&i) {
265 continue;
266 }
267
268 let mut test_x = create_model_matrix(x, &selected_indices, include_intercept);
270 let var_col = x.slice(s![.., i]).to_owned();
271 test_x
272 .push_column(var_col.view())
273 .expect("Failed to push column");
274
275 if let Ok(model) = linear_regression(&test_x.view(), y) {
277 let crit_value =
278 calculate_criterion(&model, n, model.coefficients.len(), criterion);
279
280 if is_criterion_better(crit_value, best_criterion, criterion) {
281 best_var = Some(i);
282 best_criterion = crit_value;
283 }
284 }
285 }
286
287 if let Some(var_idx) = best_var {
289 let mut test_x = create_model_matrix(x, &selected_indices, include_intercept);
290 let var_col = x.slice(s![.., var_idx]).to_owned();
291 test_x
292 .push_column(var_col.view())
293 .expect("Failed to push column");
294
295 if let Ok(model) = linear_regression(&test_x.view(), y) {
296 let var_pos = test_x.ncols() - 1;
297 let _t_value = model.t_values[var_pos];
298 let p_value = model.p_values[var_pos];
299
300 if p_value <= p_enter {
301 selected_indices.insert(var_idx);
302 current_x = test_x;
303 sequence.push((var_idx, true));
304 criteria_values.push(best_criterion);
305 criterion_improved = true;
306 }
307 }
308 }
309 }
310
311 if (direction == StepwiseDirection::Backward || direction == StepwiseDirection::Both)
313 && !criterion_improved
314 && !selected_indices.is_empty()
315 {
316 let mut worst_var = None;
318 let mut worst_criterion = F::infinity();
319
320 for &var_idx in &selected_indices {
321 let mut test_indices = selected_indices.clone();
323 test_indices.remove(&var_idx);
324
325 let test_x = create_model_matrix(x, &test_indices, include_intercept);
326
327 if let Ok(model) = linear_regression(&test_x.view(), y) {
329 let crit_value =
330 calculate_criterion(&model, n, model.coefficients.len(), criterion);
331
332 if is_criterion_better(crit_value, worst_criterion, criterion) {
333 worst_var = Some(var_idx);
334 worst_criterion = crit_value;
335 }
336 }
337 }
338
339 if let Some(var_idx) = worst_var {
341 let var_pos = find_var_position(¤t_x, x, var_idx, include_intercept);
342
343 if let Ok(model) = linear_regression(¤t_x.view(), y) {
344 let p_value = model.p_values[var_pos];
345
346 if p_value > p_remove {
347 selected_indices.remove(&var_idx);
348 current_x = create_model_matrix(x, &selected_indices, include_intercept);
349 sequence.push((var_idx, false));
350 criteria_values.push(worst_criterion);
351 criterion_improved = true;
352 }
353 }
354 }
355 }
356
357 step += 1;
358 }
359
360 let final_model = linear_regression(¤t_x.view(), y)?;
362
363 let selected_indices = selected_indices.into_iter().collect();
365
366 Ok(StepwiseResults {
367 final_model,
368 selected_indices,
369 sequence,
370 criteria_values,
371 })
372}
373
374#[allow(dead_code)]
376fn create_model_matrix<F>(
377 x: &ArrayView2<F>,
378 indices: &HashSet<usize>,
379 include_intercept: bool,
380) -> Array2<F>
381where
382 F: Float + 'static + std::iter::Sum<F> + std::fmt::Display,
383{
384 let n = x.nrows();
385 let p = indices.len();
386
387 let cols = if include_intercept { p + 1 } else { p };
388 let mut x_model = Array2::<F>::zeros((n, cols));
389
390 if include_intercept {
391 x_model.slice_mut(s![.., 0]).fill(F::one());
392 }
393
394 let offset = if include_intercept { 1 } else { 0 };
395
396 for (i, &idx) in indices.iter().enumerate() {
397 x_model
398 .slice_mut(s![.., i + offset])
399 .assign(&x.slice(s![.., idx]));
400 }
401
402 x_model
403}
404
405#[allow(dead_code)]
406fn find_var_position<F>(
407 current_x: &Array2<F>,
408 x: &ArrayView2<F>,
409 var_idx: usize,
410 include_intercept: bool,
411) -> usize
412where
413 F: Float + 'static + std::iter::Sum<F> + std::fmt::Display,
414{
415 let offset = if include_intercept { 1 } else { 0 };
416
417 for i in offset..current_x.ncols() {
418 let col = current_x.slice(s![.., i]);
419 let x_col = x.slice(s![.., var_idx]);
420
421 if col
422 .iter()
423 .zip(x_col.iter())
424 .all(|(&a, &b)| (a - b).abs() < F::epsilon())
425 {
426 return i;
427 }
428 }
429
430 current_x.ncols() - 1
432}
433
434#[allow(dead_code)]
435fn calculate_criterion<F>(
436 model: &RegressionResults<F>,
437 n: usize,
438 p: usize,
439 criterion: StepwiseCriterion,
440) -> F
441where
442 F: Float + 'static + std::iter::Sum<F> + std::fmt::Debug + std::fmt::Display,
443{
444 match criterion {
445 StepwiseCriterion::AIC => {
446 let rss: F = model
447 .residuals
448 .iter()
449 .map(|&r| scirs2_core::numeric::Float::powi(r, 2))
450 .sum();
451 let n_f = F::from(n).expect("Failed to convert to float");
452 let k_f = F::from(p).expect("Failed to convert to float");
453 n_f * scirs2_core::numeric::Float::ln(rss / n_f)
454 + F::from(2.0).expect("Failed to convert constant to float") * k_f
455 }
456 StepwiseCriterion::BIC => {
457 let rss: F = model
458 .residuals
459 .iter()
460 .map(|&r| scirs2_core::numeric::Float::powi(r, 2))
461 .sum();
462 let n_f = F::from(n).expect("Failed to convert to float");
463 let k_f = F::from(p).expect("Failed to convert to float");
464 n_f * scirs2_core::numeric::Float::ln(rss / n_f)
465 + k_f * scirs2_core::numeric::Float::ln(n_f)
466 }
467 StepwiseCriterion::AdjR2 => {
468 -model.adj_r_squared }
470 StepwiseCriterion::F => {
471 -model.f_statistic }
473 StepwiseCriterion::T => {
474 let min_t = model
476 .t_values
477 .iter()
478 .map(|&t| t.abs())
479 .fold(F::infinity(), |a, b| a.min(b));
480 -min_t }
482 }
483}
484
485#[allow(dead_code)]
486fn is_criterion_better<F>(_new_value: F, oldvalue: F, criterion: StepwiseCriterion) -> bool
487where
488 F: Float + std::fmt::Display,
489{
490 match criterion {
491 StepwiseCriterion::AIC | StepwiseCriterion::BIC => _new_value < oldvalue,
493
494 StepwiseCriterion::AdjR2 | StepwiseCriterion::F | StepwiseCriterion::T => {
496 _new_value < oldvalue
497 }
498 }
499}
500
501#[allow(dead_code)]
503fn linear_regression<F>(x: &ArrayView2<F>, y: &ArrayView1<F>) -> StatsResult<RegressionResults<F>>
504where
505 F: Float
506 + std::iter::Sum<F>
507 + std::ops::Div<Output = F>
508 + std::fmt::Debug
509 + std::fmt::Display
510 + 'static
511 + scirs2_core::numeric::NumAssign
512 + scirs2_core::numeric::One
513 + scirs2_core::ndarray::ScalarOperand
514 + Send
515 + Sync,
516{
517 let n = x.nrows();
518 let p = x.ncols();
519
520 if n <= p {
522 return Err(StatsError::InvalidArgument(format!(
523 "Number of observations ({}) must be greater than number of predictors ({})",
524 n, p
525 )));
526 }
527
528 let coefficients = match lstsq(x, y, None) {
530 Ok(result) => result.x,
531 Err(e) => {
532 return Err(StatsError::ComputationError(format!(
533 "Least squares computation failed: {:?}",
534 e
535 )))
536 }
537 };
538
539 let fitted_values = x.dot(&coefficients);
541 let residuals = y.to_owned() - &fitted_values;
542
543 let df_model = p - 1; let df_residuals = n - p;
546
547 let (_y_mean, ss_total, ss_residual, ss_explained) =
549 calculate_sum_of_squares(y, &residuals.view());
550
551 let r_squared = ss_explained / ss_total;
553 let adj_r_squared = F::one()
554 - (F::one() - r_squared) * F::from(n - 1).expect("Failed to convert to float")
555 / F::from(df_residuals).expect("Failed to convert to float");
556
557 let mse = ss_residual / F::from(df_residuals).expect("Failed to convert to float");
559 let residual_std_error = scirs2_core::numeric::Float::sqrt(mse);
560
561 let std_errors = match calculate_std_errors(x, &residuals.view(), df_residuals) {
563 Ok(se) => se,
564 Err(_) => Array1::<F>::zeros(p),
565 };
566
567 let t_values = calculate_t_values(&coefficients, &std_errors);
569
570 let p_values = t_values.mapv(|t| t_test_p_value(t, df_residuals));
573
574 let mut conf_intervals = Array2::<F>::zeros((p, 2));
576 for i in 0..p {
577 let margin = std_errors[i] * F::from(1.96).expect("Failed to convert constant to float"); conf_intervals[[i, 0]] = coefficients[i] - margin;
579 conf_intervals[[i, 1]] = coefficients[i] + margin;
580 }
581
582 let f_statistic = if df_model > 0 && df_residuals > 0 {
584 (ss_explained / F::from(df_model).expect("Failed to convert to float"))
585 / (ss_residual / F::from(df_residuals).expect("Failed to convert to float"))
586 } else {
587 F::infinity()
588 };
589
590 let f_p_value = f_test_p_value(f_statistic, df_model, df_residuals);
593
594 Ok(RegressionResults {
596 coefficients,
597 std_errors,
598 t_values,
599 p_values,
600 conf_intervals,
601 r_squared,
602 adj_r_squared,
603 f_statistic,
604 f_p_value,
605 residual_std_error,
606 df_residuals,
607 residuals,
608 fitted_values,
609 inlier_mask: vec![true; n], })
611}
612
613#[cfg(test)]
630mod f_p_value_fix_tests {
631 use super::*;
632 use approx::assert_relative_eq;
633 use scirs2_core::ndarray::array;
634
635 fn fixture_x1() -> Vec<f64> {
636 vec![
637 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
638 17.0, 18.0, 19.0, 20.0,
639 ]
640 }
641
642 fn fixture_x2() -> Vec<f64> {
643 vec![
644 5.0, 3.0, 8.0, 2.0, 9.0, 4.0, 7.0, 1.0, 6.0, 10.0, 2.0, 8.0, 3.0, 9.0, 1.0, 7.0, 4.0,
645 10.0, 5.0, 6.0,
646 ]
647 }
648
649 fn fixture_x_with_intercept() -> Array2<f64> {
652 let x1 = fixture_x1();
653 let x2 = fixture_x2();
654 let n = x1.len();
655 let mut x = Array2::<f64>::zeros((n, 3));
656 for i in 0..n {
657 x[[i, 0]] = 1.0;
658 x[[i, 1]] = x1[i];
659 x[[i, 2]] = x2[i];
660 }
661 x
662 }
663
664 fn fixture_x_no_intercept() -> Array2<f64> {
667 let x1 = fixture_x1();
668 let x2 = fixture_x2();
669 let n = x1.len();
670 let mut x = Array2::<f64>::zeros((n, 2));
671 for i in 0..n {
672 x[[i, 0]] = x1[i];
673 x[[i, 1]] = x2[i];
674 }
675 x
676 }
677
678 fn fixture_y_strong() -> Array1<f64> {
679 array![
680 -2.2, 3.3, -0.9, 10.6, 3.7, 14.05, 12.35, 24.75, 19.9, 17.12, 31.95, 26.18, 36.28,
681 30.58, 45.2, 39.64, 46.92, 41.22, 51.32, 53.1
682 ]
683 }
684
685 fn fixture_y_noise() -> Array1<f64> {
686 array![
687 3.0, 7.0, 2.0, 9.0, 4.0, 8.0, 1.0, 6.0, 5.0, 10.0, 2.5, 7.5, 3.5, 9.5, 1.5, 6.5, 4.5,
688 10.5, 5.5, 8.5
689 ]
690 }
691
692 #[test]
697 fn test_internal_linear_regression_f_p_value_matches_scipy() {
698 let x = fixture_x_with_intercept();
699
700 let strong =
701 linear_regression(&x.view(), &fixture_y_strong().view()).expect("regression ok");
702 assert_relative_eq!(strong.f_statistic, 97408.17758838173, max_relative = 1e-4);
703 assert!(
704 strong.f_p_value < 1e-12,
705 "expected ~0 (near-perfect fit), got {}",
706 strong.f_p_value
707 );
708
709 let noise = linear_regression(&x.view(), &fixture_y_noise().view()).expect("regression ok");
710 assert_relative_eq!(noise.f_statistic, 1.6747197597833423, max_relative = 1e-4);
711 assert_relative_eq!(
712 noise.f_p_value,
713 0.21683030932143513,
714 max_relative = 1e-3,
715 epsilon = 1e-6
716 );
717 assert!(
722 noise.f_p_value > 0.05,
723 "expected a large, non-significant p-value, got {}",
724 noise.f_p_value
725 );
726 }
727
728 #[test]
733 fn test_stepwise_regression_final_model_f_p_value_distinguishes_signal_from_noise() {
734 let x = fixture_x_no_intercept();
735
736 let strong = stepwise_regression(
741 &x.view(),
742 &fixture_y_strong().view(),
743 StepwiseDirection::Backward,
744 StepwiseCriterion::AdjR2,
745 None,
746 None,
747 None,
748 true,
749 )
750 .expect("stepwise regression should succeed");
751 assert!((0.0..=1.0).contains(&strong.final_model.f_p_value));
752 assert!(
753 strong.final_model.f_p_value < 0.01,
754 "strong-signal final model should be highly significant, got {}",
755 strong.final_model.f_p_value
756 );
757
758 let noise = stepwise_regression(
762 &x.view(),
763 &fixture_y_noise().view(),
764 StepwiseDirection::Forward,
765 StepwiseCriterion::F,
766 None,
767 None,
768 None,
769 true,
770 )
771 .expect("stepwise regression should succeed");
772 assert!((0.0..=1.0).contains(&noise.final_model.f_p_value));
773 assert!(
780 noise.final_model.f_p_value > 0.05,
781 "weak-signal final model should not look significant, got {}",
782 noise.final_model.f_p_value
783 );
784 }
785
786 #[test]
805 fn test_internal_linear_regression_p_values_matches_scipy() {
806 let x = fixture_x_with_intercept();
807
808 let strong =
809 linear_regression(&x.view(), &fixture_y_strong().view()).expect("regression ok");
810 for &p in strong.p_values.iter() {
811 assert!((0.0..=1.0).contains(&p), "p-value out of range: {p}");
812 }
813 assert_relative_eq!(strong.p_values[0], 2.77999845e-12, epsilon = 1e-9);
814 assert!(
815 strong.p_values[1] < 1e-9 && strong.p_values[2] < 1e-9,
816 "expected near-zero p-values for x1/x2, got {:?}",
817 strong.p_values
818 );
819
820 let noise = linear_regression(&x.view(), &fixture_y_noise().view()).expect("regression ok");
821 for &p in noise.p_values.iter() {
822 assert!((0.0..=1.0).contains(&p), "p-value out of range: {p}");
823 }
824 assert_relative_eq!(noise.p_values[0], 0.11240478, max_relative = 1e-3);
825 assert_relative_eq!(noise.p_values[1], 0.36838016, max_relative = 1e-3);
826 assert_relative_eq!(noise.p_values[2], 0.16437859, max_relative = 1e-3);
827 assert!(
831 noise.p_values[1] > 0.05 && noise.p_values[2] > 0.05,
832 "expected non-significant p-values for noise-only x1/x2, got {:?}",
833 noise.p_values
834 );
835 }
836}