1use ferrolearn_core::error::FerroError;
52use ferrolearn_core::introspection::HasCoefficients;
53use ferrolearn_core::pipeline::{FittedPipelineEstimator, PipelineEstimator};
54use ferrolearn_core::traits::{Fit, Predict};
55use ndarray::{Array1, Array2, Axis, ScalarOperand};
56use num_traits::{Float, FromPrimitive};
57
58#[derive(Debug, Clone)]
73pub struct OrthogonalMatchingPursuit<F> {
74 pub n_nonzero_coefs: Option<usize>,
77 pub tol: Option<F>,
80 pub fit_intercept: bool,
82}
83
84impl<F: Float> OrthogonalMatchingPursuit<F> {
85 #[must_use]
90 pub fn new() -> Self {
91 Self {
92 n_nonzero_coefs: None,
93 tol: None,
94 fit_intercept: true,
95 }
96 }
97
98 #[must_use]
100 pub fn with_n_nonzero_coefs(mut self, n: usize) -> Self {
101 self.n_nonzero_coefs = Some(n);
102 self
103 }
104
105 #[must_use]
107 pub fn with_tol(mut self, tol: F) -> Self {
108 self.tol = Some(tol);
109 self
110 }
111
112 #[must_use]
114 pub fn with_fit_intercept(mut self, fit_intercept: bool) -> Self {
115 self.fit_intercept = fit_intercept;
116 self
117 }
118}
119
120impl<F: Float> Default for OrthogonalMatchingPursuit<F> {
121 fn default() -> Self {
122 Self::new()
123 }
124}
125
126#[derive(Debug, Clone)]
130pub struct FittedOMP<F> {
131 coefficients: Array1<F>,
133 intercept: F,
135}
136
137fn cholesky_solve<F: Float>(a: &Array2<F>, b: &Array1<F>) -> Result<Array1<F>, FerroError> {
143 let n = a.nrows();
144 let mut l = Array2::<F>::zeros((n, n));
145
146 for i in 0..n {
147 for j in 0..=i {
148 let mut s = a[[i, j]];
149 for k in 0..j {
150 s = s - l[[i, k]] * l[[j, k]];
151 }
152 if i == j {
153 if s <= F::zero() {
154 return Err(FerroError::NumericalInstability {
155 message: "Cholesky: matrix not positive definite".into(),
156 });
157 }
158 l[[i, j]] = s.sqrt();
159 } else {
160 l[[i, j]] = s / l[[j, j]];
161 }
162 }
163 }
164
165 let mut z = Array1::<F>::zeros(n);
166 for i in 0..n {
167 let mut s = b[i];
168 for k in 0..i {
169 s = s - l[[i, k]] * z[k];
170 }
171 z[i] = s / l[[i, i]];
172 }
173
174 let mut x_sol = Array1::<F>::zeros(n);
175 for i in (0..n).rev() {
176 let mut s = z[i];
177 for k in (i + 1)..n {
178 s = s - l[[k, i]] * x_sol[k];
179 }
180 x_sol[i] = s / l[[i, i]];
181 }
182
183 Ok(x_sol)
184}
185
186fn gaussian_solve<F: Float>(
188 n: usize,
189 a: &Array2<F>,
190 b: &Array1<F>,
191) -> Result<Array1<F>, FerroError> {
192 let mut aug = Array2::<F>::zeros((n, n + 1));
193 for i in 0..n {
194 for j in 0..n {
195 aug[[i, j]] = a[[i, j]];
196 }
197 aug[[i, n]] = b[i];
198 }
199
200 for col in 0..n {
201 let mut max_val = aug[[col, col]].abs();
202 let mut max_row = col;
203 for row in (col + 1)..n {
204 let v = aug[[row, col]].abs();
205 if v > max_val {
206 max_val = v;
207 max_row = row;
208 }
209 }
210
211 if max_val < F::from(1e-12).unwrap_or_else(F::epsilon) {
212 return Err(FerroError::NumericalInstability {
213 message: "singular matrix in Gaussian elimination".into(),
214 });
215 }
216
217 if max_row != col {
218 for j in 0..=n {
219 let tmp = aug[[col, j]];
220 aug[[col, j]] = aug[[max_row, j]];
221 aug[[max_row, j]] = tmp;
222 }
223 }
224
225 let pivot = aug[[col, col]];
226 for row in (col + 1)..n {
227 let factor = aug[[row, col]] / pivot;
228 for j in col..=n {
229 let above = aug[[col, j]];
230 aug[[row, j]] = aug[[row, j]] - factor * above;
231 }
232 }
233 }
234
235 let mut x_sol = Array1::<F>::zeros(n);
236 for i in (0..n).rev() {
237 let mut s = aug[[i, n]];
238 for j in (i + 1)..n {
239 s = s - aug[[i, j]] * x_sol[j];
240 }
241 if aug[[i, i]].abs() < F::from(1e-12).unwrap_or_else(F::epsilon) {
242 return Err(FerroError::NumericalInstability {
243 message: "near-zero pivot in back substitution".into(),
244 });
245 }
246 x_sol[i] = s / aug[[i, i]];
247 }
248
249 Ok(x_sol)
250}
251
252fn ols_active<F: Float + FromPrimitive + 'static>(
254 x: &Array2<F>,
255 y: &Array1<F>,
256 support: &[usize],
257 n_features: usize,
258) -> Result<Array1<F>, FerroError> {
259 let n_samples = x.nrows();
260 let k = support.len();
261
262 let mut xa = Array2::<F>::zeros((n_samples, k));
263 for (col_idx, &j) in support.iter().enumerate() {
264 for i in 0..n_samples {
265 xa[[i, col_idx]] = x[[i, j]];
266 }
267 }
268
269 let xat = xa.t();
270 let xtx = xat.dot(&xa);
271 let xty = xat.dot(y);
272
273 let w_active = cholesky_solve(&xtx, &xty).or_else(|_| gaussian_solve(k, &xtx, &xty))?;
274
275 let mut w = Array1::<F>::zeros(n_features);
276 for (col_idx, &j) in support.iter().enumerate() {
277 w[j] = w_active[col_idx];
278 }
279 Ok(w)
280}
281
282impl<F: Float + Send + Sync + ScalarOperand + FromPrimitive + 'static> Fit<Array2<F>, Array1<F>>
287 for OrthogonalMatchingPursuit<F>
288{
289 type Fitted = FittedOMP<F>;
290 type Error = FerroError;
291
292 fn fit(&self, x: &Array2<F>, y: &Array1<F>) -> Result<FittedOMP<F>, FerroError> {
304 let (n_samples, n_features) = x.dim();
305
306 if n_samples != y.len() {
307 return Err(FerroError::ShapeMismatch {
308 expected: vec![n_samples],
309 actual: vec![y.len()],
310 context: "y length must match number of samples in X".into(),
311 });
312 }
313
314 if n_samples == 0 {
315 return Err(FerroError::InsufficientSamples {
316 required: 1,
317 actual: 0,
318 context: "OMP requires at least one sample".into(),
319 });
320 }
321
322 if x.iter().any(|v| !v.is_finite()) {
332 return Err(FerroError::InvalidParameter {
333 name: "X".into(),
334 reason: "Input X contains NaN or infinity.".into(),
335 });
336 }
337 if y.iter().any(|v| !v.is_finite()) {
338 return Err(FerroError::InvalidParameter {
339 name: "y".into(),
340 reason: "Input y contains NaN or infinity.".into(),
341 });
342 }
343
344 let effective_n_nonzero = if self.n_nonzero_coefs.is_none() && self.tol.is_none() {
348 Some(((n_features as f64 * 0.1) as usize).max(1))
349 } else {
350 self.n_nonzero_coefs
351 };
352
353 let max_k = effective_n_nonzero.unwrap_or(n_features).min(n_features);
354
355 if let Some(n) = self.n_nonzero_coefs
356 && n > n_features
357 {
358 return Err(FerroError::InvalidParameter {
359 name: "n_nonzero_coefs".into(),
360 reason: format!("cannot exceed number of features ({n_features})"),
361 });
362 }
363
364 let (x_work, y_work, x_mean, y_mean) = if self.fit_intercept {
366 let x_mean = x
367 .mean_axis(Axis(0))
368 .ok_or_else(|| FerroError::NumericalInstability {
369 message: "failed to compute column means".into(),
370 })?;
371 let y_mean = y.mean().ok_or_else(|| FerroError::NumericalInstability {
372 message: "failed to compute target mean".into(),
373 })?;
374 let x_c = x - &x_mean;
375 let y_c = y - y_mean;
376 (x_c, y_c, Some(x_mean), Some(y_mean))
377 } else {
378 (x.clone(), y.clone(), None, None)
379 };
380
381 let mut support: Vec<usize> = Vec::with_capacity(max_k);
382 let mut in_support = vec![false; n_features];
383 let mut w = Array1::<F>::zeros(n_features);
384 let mut residual = y_work.clone();
385
386 for _step in 0..max_k {
387 if let Some(tol_val) = self.tol {
389 let res_norm_sq = residual.dot(&residual);
390 if res_norm_sq < tol_val {
391 break;
392 }
393 }
394
395 let mut best_j = None;
397 let mut best_corr = F::zero();
398 for (j, &is_in_support) in in_support.iter().enumerate() {
399 if is_in_support {
400 continue;
401 }
402 let corr = x_work.column(j).dot(&residual).abs();
403 if corr > best_corr {
404 best_corr = corr;
405 best_j = Some(j);
406 }
407 }
408
409 let j = match best_j {
410 Some(j) => j,
411 None => break,
412 };
413
414 support.push(j);
415 in_support[j] = true;
416
417 w = ols_active(&x_work, &y_work, &support, n_features)?;
419
420 residual = &y_work - x_work.dot(&w);
422 }
423
424 let intercept = if let (Some(xm), Some(ym)) = (&x_mean, &y_mean) {
425 *ym - xm.dot(&w)
426 } else {
427 F::zero()
428 };
429
430 Ok(FittedOMP {
431 coefficients: w,
432 intercept,
433 })
434 }
435}
436
437impl<F: Float + Send + Sync + ScalarOperand + 'static> Predict<Array2<F>> for FittedOMP<F> {
442 type Output = Array1<F>;
443 type Error = FerroError;
444
445 fn predict(&self, x: &Array2<F>) -> Result<Array1<F>, FerroError> {
454 if x.ncols() != self.coefficients.len() {
455 return Err(FerroError::ShapeMismatch {
456 expected: vec![self.coefficients.len()],
457 actual: vec![x.ncols()],
458 context: "number of features must match fitted model".into(),
459 });
460 }
461 Ok(x.dot(&self.coefficients) + self.intercept)
462 }
463}
464
465impl<F: Float + Send + Sync + ScalarOperand + 'static> HasCoefficients<F> for FittedOMP<F> {
466 fn coefficients(&self) -> &Array1<F> {
467 &self.coefficients
468 }
469
470 fn intercept(&self) -> F {
471 self.intercept
472 }
473}
474
475impl<F> PipelineEstimator<F> for OrthogonalMatchingPursuit<F>
476where
477 F: Float + FromPrimitive + ScalarOperand + Send + Sync + 'static,
478{
479 fn fit_pipeline(
480 &self,
481 x: &Array2<F>,
482 y: &Array1<F>,
483 ) -> Result<Box<dyn FittedPipelineEstimator<F>>, FerroError> {
484 let fitted = self.fit(x, y)?;
485 Ok(Box::new(fitted))
486 }
487}
488
489impl<F> FittedPipelineEstimator<F> for FittedOMP<F>
490where
491 F: Float + ScalarOperand + Send + Sync + 'static,
492{
493 fn predict_pipeline(&self, x: &Array2<F>) -> Result<Array1<F>, FerroError> {
494 self.predict(x)
495 }
496}
497
498#[cfg(test)]
503mod tests {
504 use super::*;
505 use approx::assert_relative_eq;
506 use ndarray::array;
507
508 #[test]
509 fn test_defaults() {
510 let m = OrthogonalMatchingPursuit::<f64>::new();
511 assert!(m.n_nonzero_coefs.is_none());
512 assert!(m.tol.is_none());
513 assert!(m.fit_intercept);
514 }
515
516 #[test]
517 fn test_builder() {
518 let m = OrthogonalMatchingPursuit::<f64>::new()
519 .with_n_nonzero_coefs(3)
520 .with_tol(1e-4)
521 .with_fit_intercept(false);
522 assert_eq!(m.n_nonzero_coefs, Some(3));
523 assert_relative_eq!(m.tol.unwrap(), 1e-4);
524 assert!(!m.fit_intercept);
525 }
526
527 #[test]
528 fn test_shape_mismatch() {
529 let x = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
530 let y = array![1.0, 2.0];
531 assert!(
532 OrthogonalMatchingPursuit::<f64>::new()
533 .with_n_nonzero_coefs(1)
534 .fit(&x, &y)
535 .is_err()
536 );
537 }
538
539 #[test]
540 fn test_default_n_nonzero_fits() {
541 let x = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]);
545 let y = array![1.0, 2.0, 3.0];
546 assert!(x.is_ok(), "valid shape");
547 let Ok(x) = x else { return };
548 let result = OrthogonalMatchingPursuit::<f64>::new().fit(&x, &y);
549 assert!(result.is_ok(), "default OMP must fit, not error");
550 let Ok(fitted) = result else { return };
551 let nonzero = fitted
552 .coefficients()
553 .iter()
554 .filter(|&&c| c.abs() > 1e-10)
555 .count();
556 assert_eq!(nonzero, 1);
557 }
558
559 #[test]
560 fn test_n_nonzero_exceeds_features() {
561 let x = Array2::from_shape_vec((3, 2), vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]).unwrap();
562 let y = array![1.0, 2.0, 3.0];
563 assert!(
564 OrthogonalMatchingPursuit::<f64>::new()
565 .with_n_nonzero_coefs(5)
566 .fit(&x, &y)
567 .is_err()
568 );
569 }
570
571 #[test]
572 fn test_simple_linear() {
573 let x = Array2::from_shape_vec((5, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
574 let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
575
576 let fitted = OrthogonalMatchingPursuit::<f64>::new()
577 .with_n_nonzero_coefs(1)
578 .fit(&x, &y)
579 .unwrap();
580 assert_relative_eq!(fitted.coefficients()[0], 2.0, epsilon = 1e-6);
581 assert_relative_eq!(fitted.intercept(), 1.0, epsilon = 1e-6);
582 }
583
584 #[test]
585 fn test_sparsity() {
586 let x = Array2::from_shape_vec(
588 (10, 3),
589 vec![
590 1.0, 0.1, 0.01, 2.0, 0.2, 0.02, 3.0, 0.3, 0.03, 4.0, 0.4, 0.04, 5.0, 0.5, 0.05,
591 6.0, 0.6, 0.06, 7.0, 0.7, 0.07, 8.0, 0.8, 0.08, 9.0, 0.9, 0.09, 10.0, 1.0, 0.10,
592 ],
593 )
594 .unwrap();
595 let y = array![2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0, 20.0];
596
597 let fitted = OrthogonalMatchingPursuit::<f64>::new()
598 .with_n_nonzero_coefs(1)
599 .fit(&x, &y)
600 .unwrap();
601 let nonzero = fitted
602 .coefficients()
603 .iter()
604 .filter(|&&c| c.abs() > 1e-10)
605 .count();
606 assert_eq!(nonzero, 1);
607 }
608
609 #[test]
610 fn test_tol_stopping() {
611 let x = Array2::from_shape_vec((5, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
612 let y = array![2.0, 4.0, 6.0, 8.0, 10.0]; let fitted = OrthogonalMatchingPursuit::<f64>::new()
615 .with_tol(1e-10)
616 .fit(&x, &y)
617 .unwrap();
618 let preds = fitted.predict(&x).unwrap();
620 for (pred, actual) in preds.iter().zip(y.iter()) {
621 assert_relative_eq!(pred, actual, epsilon = 1e-4);
622 }
623 }
624
625 #[test]
626 fn test_predict() {
627 let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
628 let y = array![2.0, 4.0, 6.0, 8.0];
629
630 let fitted = OrthogonalMatchingPursuit::<f64>::new()
631 .with_n_nonzero_coefs(1)
632 .fit(&x, &y)
633 .unwrap();
634 let preds = fitted.predict(&x).unwrap();
635 assert_eq!(preds.len(), 4);
636 }
637
638 #[test]
639 fn test_predict_feature_mismatch() {
640 let x = Array2::from_shape_vec((3, 2), vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]).unwrap();
641 let y = array![1.0, 2.0, 3.0];
642 let fitted = OrthogonalMatchingPursuit::<f64>::new()
643 .with_n_nonzero_coefs(1)
644 .fit(&x, &y)
645 .unwrap();
646 let x_bad = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
647 assert!(fitted.predict(&x_bad).is_err());
648 }
649
650 #[test]
651 fn test_has_coefficients() {
652 let x = Array2::from_shape_vec((3, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
653 let y = array![1.0, 2.0, 3.0];
654 let fitted = OrthogonalMatchingPursuit::<f64>::new()
655 .with_n_nonzero_coefs(2)
656 .fit(&x, &y)
657 .unwrap();
658 assert_eq!(fitted.coefficients().len(), 2);
659 }
660
661 #[test]
662 fn test_no_intercept() {
663 let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
664 let y = array![2.0, 4.0, 6.0, 8.0];
665
666 let fitted = OrthogonalMatchingPursuit::<f64>::new()
667 .with_n_nonzero_coefs(1)
668 .with_fit_intercept(false)
669 .fit(&x, &y)
670 .unwrap();
671 assert_relative_eq!(fitted.intercept(), 0.0, epsilon = 1e-10);
672 }
673
674 #[test]
675 fn test_pipeline() {
676 let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
677 let y = array![3.0, 5.0, 7.0, 9.0];
678 let model = OrthogonalMatchingPursuit::<f64>::new().with_n_nonzero_coefs(1);
679 let fitted = model.fit_pipeline(&x, &y).unwrap();
680 let preds = fitted.predict_pipeline(&x).unwrap();
681 assert_eq!(preds.len(), 4);
682 }
683
684 #[test]
685 fn test_multivariate_recovery() {
686 let x = Array2::from_shape_vec(
688 (5, 3),
689 vec![
690 1.0, 0.0, 0.5, 0.0, 1.0, 0.3, 1.0, 1.0, 0.1, 2.0, 0.0, 0.8, 0.0, 2.0, 0.4,
691 ],
692 )
693 .unwrap();
694 let y = array![1.0, 3.0, 4.0, 2.0, 6.0]; let fitted = OrthogonalMatchingPursuit::<f64>::new()
697 .with_n_nonzero_coefs(2)
698 .fit(&x, &y)
699 .unwrap();
700
701 assert!(
703 fitted.coefficients()[2].abs() < 0.5,
704 "irrelevant feature should have near-zero coefficient"
705 );
706 }
707}