1use crate::error::OptimizeError;
8use crate::least_squares::Options;
9use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
10
11#[derive(Debug, Clone)]
13pub struct SparseMatrix {
14 pub row_ptr: Vec<usize>,
16 pub col_idx: Vec<usize>,
18 pub values: Vec<f64>,
20 pub nrows: usize,
22 pub ncols: usize,
24}
25
26impl SparseMatrix {
27 pub fn new(nrows: usize, ncols: usize) -> Self {
29 Self {
30 row_ptr: vec![0; nrows + 1],
31 col_idx: Vec::new(),
32 values: Vec::new(),
33 nrows,
34 ncols,
35 }
36 }
37
38 pub fn from_dense(matrix: &ArrayView2<f64>, threshold: f64) -> Self {
40 let (nrows, ncols) = matrix.dim();
41 let mut row_ptr = vec![0; nrows + 1];
42 let mut col_idx = Vec::new();
43 let mut values = Vec::new();
44
45 for i in 0..nrows {
46 for j in 0..ncols {
47 if matrix[[i, j]].abs() > threshold {
48 col_idx.push(j);
49 values.push(matrix[[i, j]]);
50 }
51 }
52 row_ptr[i + 1] = values.len();
53 }
54
55 Self {
56 row_ptr,
57 col_idx,
58 values,
59 nrows,
60 ncols,
61 }
62 }
63
64 pub fn matvec(&self, x: &ArrayView1<f64>) -> Array1<f64> {
66 assert_eq!(x.len(), self.ncols);
67 let mut y = Array1::zeros(self.nrows);
68
69 for i in 0..self.nrows {
70 let start = self.row_ptr[i];
71 let end = self.row_ptr[i + 1];
72
73 let mut sum = 0.0;
74 for k in start..end {
75 sum += self.values[k] * x[self.col_idx[k]];
76 }
77 y[i] = sum;
78 }
79
80 y
81 }
82
83 pub fn transpose_matvec(&self, x: &ArrayView1<f64>) -> Array1<f64> {
85 assert_eq!(x.len(), self.nrows);
86 let mut y = Array1::zeros(self.ncols);
87
88 for i in 0..self.nrows {
89 let start = self.row_ptr[i];
90 let end = self.row_ptr[i + 1];
91
92 for k in start..end {
93 y[self.col_idx[k]] += self.values[k] * x[i];
94 }
95 }
96
97 y
98 }
99
100 pub fn sparsity_ratio(&self) -> f64 {
102 self.values.len() as f64 / (self.nrows * self.ncols) as f64
103 }
104}
105
106#[derive(Debug, Clone)]
108pub struct SparseOptions {
109 pub base_options: Options,
111 pub sparsity_threshold: f64,
113 pub max_iter: usize,
115 pub tol: f64,
117 pub lambda: f64,
119 pub use_coordinate_descent: bool,
121 pub block_size: usize,
123 pub use_preconditioning: bool,
125 pub memory_limit_mb: usize,
127}
128
129impl Default for SparseOptions {
130 fn default() -> Self {
131 Self {
132 base_options: Options::default(),
133 sparsity_threshold: 1e-12,
134 max_iter: 1000,
135 tol: 1e-8,
136 lambda: 0.0,
137 use_coordinate_descent: false,
138 block_size: 100,
139 use_preconditioning: true,
140 memory_limit_mb: 1000,
141 }
142 }
143}
144
145#[derive(Debug, Clone)]
147pub struct SparseResult {
148 pub x: Array1<f64>,
150 pub cost: f64,
152 pub residual: Array1<f64>,
154 pub nit: usize,
156 pub nfev: usize,
158 pub njev: usize,
160 pub success: bool,
162 pub message: String,
164 pub sparsity_info: SparseInfo,
166}
167
168#[derive(Debug, Clone)]
170pub struct SparseInfo {
171 pub jacobian_sparsity: f64,
173 pub jacobian_nnz: usize,
175 pub memory_usage_mb: f64,
177 pub used_sparse_algorithms: bool,
179}
180
181#[allow(dead_code)]
183pub fn sparse_least_squares<F, J>(
184 fun: F,
185 jac: Option<J>,
186 x0: Array1<f64>,
187 options: Option<SparseOptions>,
188) -> Result<SparseResult, OptimizeError>
189where
190 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
191 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
192{
193 let options = options.unwrap_or_default();
194 let x = x0.clone();
195 let n = x.len();
196 let mut nfev = 0;
197 let mut njev = 0;
198
199 let residual = fun(&x.view());
201 nfev += 1;
202 let m = residual.len();
203
204 let jacobian = if let Some(ref jac_fn) = jac {
206 let jac_dense = jac_fn(&x.view());
207 njev += 1;
208 Some(jac_dense)
209 } else {
210 None
211 };
212
213 let (sparse_jac, sparsity_info) = if let Some(ref jac_matrix) = jacobian {
214 let sparse_matrix =
215 SparseMatrix::from_dense(&jac_matrix.view(), options.sparsity_threshold);
216 let sparsity_ratio = sparse_matrix.sparsity_ratio();
217 let memory_usage = estimate_memory_usage(&sparse_matrix);
218
219 let info = SparseInfo {
220 jacobian_sparsity: sparsity_ratio,
221 jacobian_nnz: sparse_matrix.values.len(),
222 memory_usage_mb: memory_usage,
223 used_sparse_algorithms: sparsity_ratio < 0.5, };
225
226 (Some(sparse_matrix), info)
227 } else {
228 let info = SparseInfo {
229 jacobian_sparsity: 1.0,
230 jacobian_nnz: m * n,
231 memory_usage_mb: (m * n * 8) as f64 / (1024.0 * 1024.0),
232 used_sparse_algorithms: false,
233 };
234 (None, info)
235 };
236
237 let result = if options.lambda > 0.0 {
239 if options.use_coordinate_descent {
241 solve_lasso_coordinate_descent(&fun, &x, &options, &mut nfev)?
242 } else {
243 solve_lasso_proximal_gradient(&fun, &x, &options, &mut nfev)?
244 }
245 } else if sparsity_info.used_sparse_algorithms && sparse_jac.is_some() {
246 solve_sparse_gauss_newton(
248 &fun,
249 &jac,
250 &sparse_jac.unwrap(),
251 &x,
252 &options,
253 &mut nfev,
254 &mut njev,
255 )?
256 } else {
257 solve_dense_least_squares(&fun, &jac, &x, &options, &mut nfev, &mut njev)?
259 };
260
261 Ok(SparseResult {
262 x: result.x,
263 cost: result.cost,
264 residual: result.residual,
265 nit: result.nit,
266 nfev,
267 njev,
268 success: result.success,
269 message: result.message,
270 sparsity_info,
271 })
272}
273
274#[derive(Debug)]
276struct InternalResult {
277 x: Array1<f64>,
278 cost: f64,
279 residual: Array1<f64>,
280 nit: usize,
281 success: bool,
282 message: String,
283}
284
285#[allow(dead_code)]
287fn solve_lasso_coordinate_descent<F>(
288 fun: &F,
289 x0: &Array1<f64>,
290 options: &SparseOptions,
291 nfev: &mut usize,
292) -> Result<InternalResult, OptimizeError>
293where
294 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
295{
296 let mut x = x0.clone();
297 let n = x.len();
298 let lambda = options.lambda;
299
300 for iter in 0..options.max_iter {
301 let mut max_change: f64 = 0.0;
302
303 for j in 0..n {
305 let old_xj = x[j];
306
307 let residual = fun(&x.view());
309 *nfev += 1;
310
311 let eps = 1e-8;
313 let mut x_plus = x.clone();
314 x_plus[j] += eps;
315 let residual_plus = fun(&x_plus.view());
316 *nfev += 1;
317
318 let grad_j = residual_plus
319 .iter()
320 .zip(residual.iter())
321 .map(|(&rp, &r)| (rp - r) / eps)
322 .sum::<f64>();
323
324 let step_size = options.base_options.gtol.unwrap_or(1e-5);
326 let z = x[j] - step_size * grad_j;
327 x[j] = soft_threshold(z, lambda * step_size);
328
329 max_change = max_change.max((x[j] - old_xj).abs());
330 }
331
332 if max_change < options.tol {
334 let final_residual = fun(&x.view());
335 *nfev += 1;
336 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
337 + lambda * x.mapv(|xi| xi.abs()).sum();
338
339 return Ok(InternalResult {
340 x,
341 cost,
342 residual: final_residual,
343 nit: iter,
344 success: true,
345 message: "LASSO coordinate descent converged successfully".to_string(),
346 });
347 }
348 }
349
350 let final_residual = fun(&x.view());
351 *nfev += 1;
352 let cost =
353 0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
354
355 Ok(InternalResult {
356 x,
357 cost,
358 residual: final_residual,
359 nit: options.max_iter,
360 success: false,
361 message: "Maximum iterations reached in LASSO coordinate descent".to_string(),
362 })
363}
364
365#[allow(dead_code)]
367fn soft_threshold(x: f64, threshold: f64) -> f64 {
368 if x > threshold {
369 x - threshold
370 } else if x < -threshold {
371 x + threshold
372 } else {
373 0.0
374 }
375}
376
377#[allow(dead_code)]
379fn solve_lasso_proximal_gradient<F>(
380 fun: &F,
381 x0: &Array1<f64>,
382 options: &SparseOptions,
383 nfev: &mut usize,
384) -> Result<InternalResult, OptimizeError>
385where
386 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
387{
388 let mut x = x0.clone();
389 let lambda = options.lambda;
390 let step_size = 0.01; for iter in 0..options.max_iter {
393 let residual = fun(&x.view());
395 *nfev += 1;
396
397 let mut grad = Array1::zeros(x.len());
398 let eps = 1e-8;
399
400 for i in 0..x.len() {
401 let mut x_plus = x.clone();
402 x_plus[i] += eps;
403 let residual_plus = fun(&x_plus.view());
404 *nfev += 1;
405
406 grad[i] = 2.0
409 * residual_plus
410 .iter()
411 .zip(residual.iter())
412 .map(|(&rp, &r)| r * (rp - r) / eps)
413 .sum::<f64>();
414 }
415
416 let x_old = x.clone();
418 for i in 0..x.len() {
419 let z = x[i] - step_size * grad[i];
420 x[i] = soft_threshold(z, lambda * step_size);
421 }
422
423 let change = (&x - &x_old).mapv(|dx| dx.abs()).sum();
425 if change < options.tol {
426 let final_residual = fun(&x.view());
427 *nfev += 1;
428 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
429 + lambda * x.mapv(|xi| xi.abs()).sum();
430
431 return Ok(InternalResult {
432 x,
433 cost,
434 residual: final_residual,
435 nit: iter,
436 success: true,
437 message: "LASSO proximal gradient converged successfully".to_string(),
438 });
439 }
440 }
441
442 let final_residual = fun(&x.view());
443 *nfev += 1;
444 let cost =
445 0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
446
447 Ok(InternalResult {
448 x,
449 cost,
450 residual: final_residual,
451 nit: options.max_iter,
452 success: false,
453 message: "Maximum iterations reached in LASSO proximal gradient".to_string(),
454 })
455}
456
457#[allow(dead_code)]
459fn solve_sparse_gauss_newton<F, J>(
460 fun: &F,
461 jac: &Option<J>,
462 _sparse_jac: &SparseMatrix,
463 x0: &Array1<f64>,
464 options: &SparseOptions,
465 nfev: &mut usize,
466 njev: &mut usize,
467) -> Result<InternalResult, OptimizeError>
468where
469 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
470 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
471{
472 let mut x = x0.clone();
473
474 for iter in 0..options.max_iter {
475 let residual = fun(&x.view());
477 *nfev += 1;
478
479 if let Some(ref jac_fn) = jac {
480 let jac_dense = jac_fn(&x.view());
481 *njev += 1;
482
483 let jac_sparse =
485 SparseMatrix::from_dense(&jac_dense.view(), options.sparsity_threshold);
486
487 let jt_r = jac_sparse.transpose_matvec(&residual.view());
490
491 let mut p = Array1::zeros(x.len());
493 for i in 0..x.len() {
494 let diag_elem = compute_diagonal_element(&jac_sparse, i);
496 if diag_elem.abs() > 1e-12 {
497 p[i] = -jt_r[i] / diag_elem;
498 }
499 }
500
501 let mut alpha = 1.0;
503 let current_cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
504
505 for _ in 0..10 {
506 let x_new = &x + alpha * &p;
507 let residual_new = fun(&x_new.view());
508 *nfev += 1;
509 let new_cost = 0.5 * residual_new.mapv(|r| r.powi(2)).sum();
510
511 if new_cost < current_cost {
512 x = x_new;
513 break;
514 }
515 alpha *= 0.5;
516 }
517
518 let grad_norm = jt_r.mapv(|g| g.abs()).sum();
520 if grad_norm < options.tol {
521 let final_residual = fun(&x.view());
522 *nfev += 1;
523 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
524
525 return Ok(InternalResult {
526 x,
527 cost,
528 residual: final_residual,
529 nit: iter,
530 success: true,
531 message: "Sparse Gauss-Newton converged successfully".to_string(),
532 });
533 }
534 }
535 }
536
537 let final_residual = fun(&x.view());
538 *nfev += 1;
539 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
540
541 Ok(InternalResult {
542 x,
543 cost,
544 residual: final_residual,
545 nit: options.max_iter,
546 success: false,
547 message: "Maximum iterations reached in sparse Gauss-Newton".to_string(),
548 })
549}
550
551#[allow(dead_code)]
553fn compute_diagonal_element(jac: &SparseMatrix, col: usize) -> f64 {
554 let mut diag = 0.0;
555
556 for row in 0..jac.nrows {
557 let start = jac.row_ptr[row];
558 let end = jac.row_ptr[row + 1];
559
560 for k in start..end {
561 if jac.col_idx[k] == col {
562 diag += jac.values[k].powi(2);
563 }
564 }
565 }
566
567 diag
568}
569
570#[allow(dead_code)]
572fn solve_dense_least_squares<F, J>(
573 fun: &F,
574 _jac: &Option<J>,
575 x0: &Array1<f64>,
576 options: &SparseOptions,
577 nfev: &mut usize,
578 _njev: &mut usize,
579) -> Result<InternalResult, OptimizeError>
580where
581 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
582 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
583{
584 let mut x = x0.clone();
586
587 for iter in 0..options.max_iter {
588 let residual = fun(&x.view());
589 *nfev += 1;
590
591 let mut grad = Array1::zeros(x.len());
593 let eps = 1e-8;
594
595 for i in 0..x.len() {
596 let mut x_plus = x.clone();
597 x_plus[i] += eps;
598 let residual_plus = fun(&x_plus.view());
599 *nfev += 1;
600
601 grad[i] = 2.0
604 * residual_plus
605 .iter()
606 .zip(residual.iter())
607 .map(|(&rp, &r)| r * (rp - r) / eps)
608 .sum::<f64>();
609 }
610
611 let grad_norm = grad.mapv(|g| g.abs()).sum();
613 if grad_norm < options.tol {
614 let cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
615 return Ok(InternalResult {
616 x,
617 cost,
618 residual,
619 nit: iter,
620 success: true,
621 message: "Dense least squares converged successfully".to_string(),
622 });
623 }
624
625 let step_size = 0.01;
627 x = x - step_size * &grad;
628 }
629
630 let final_residual = fun(&x.view());
631 *nfev += 1;
632 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
633
634 Ok(InternalResult {
635 x,
636 cost,
637 residual: final_residual,
638 nit: options.max_iter,
639 success: false,
640 message: "Maximum iterations reached in dense least squares".to_string(),
641 })
642}
643
644#[allow(dead_code)]
646fn estimate_memory_usage(sparse_matrix: &SparseMatrix) -> f64 {
647 let nnz = sparse_matrix.values.len();
648 let nrows = sparse_matrix.nrows;
649
650 let memory_bytes = nnz * 8 + nnz * 8 + (nrows + 1) * 8; memory_bytes as f64 / (1024.0 * 1024.0)
653}
654
655#[allow(dead_code)]
657pub fn lsqr<F>(
658 matvec: F,
659 rmatvec: F,
660 b: &ArrayView1<f64>,
661 x0: Option<Array1<f64>>,
662 max_iter: Option<usize>,
663 tol: Option<f64>,
664) -> Result<Array1<f64>, OptimizeError>
665where
666 F: Fn(&ArrayView1<f64>) -> Array1<f64> + Clone,
667{
668 let n = if let Some(ref x0_val) = x0 {
669 x0_val.len()
670 } else {
671 b.len()
672 };
673 let m = b.len();
674 let max_iter = max_iter.unwrap_or(n.min(m));
675 let tol = tol.unwrap_or(1e-8);
676
677 let mut x = x0.unwrap_or_else(|| Array1::zeros(n));
678 #[allow(unused_assignments)]
679 let mut beta = 0.0;
680 #[allow(unused_assignments)]
681 let mut u = Array1::zeros(m);
682
683 let ax = matvec.clone()(&x.view());
685 let r = b - &ax;
686 let mut v = rmatvec.clone()(&r.view());
687
688 let mut alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
689 if alpha == 0.0 {
690 return Ok(x);
691 }
692
693 v /= alpha;
694 let mut w = v.clone();
695
696 for _iter in 0..max_iter {
697 let av = matvec.clone()(&v.view());
699 let alpha_new = av.mapv(|avi| avi.powi(2)).sum().sqrt();
700
701 if alpha_new == 0.0 {
702 break;
703 }
704
705 u = av / alpha_new;
706 beta = alpha_new;
707
708 let atu = rmatvec.clone()(&u.view());
709 v = atu - beta * &v;
710 alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
711
712 if alpha == 0.0 {
713 break;
714 }
715
716 v /= alpha;
717
718 x = x + (beta / alpha) * &w;
720 w = v.clone() - (alpha / beta) * &w;
721
722 let residual_norm = (b - &matvec.clone()(&x.view()))
724 .mapv(|ri| ri.powi(2))
725 .sum()
726 .sqrt();
727 if residual_norm < tol {
728 break;
729 }
730 }
731
732 Ok(x)
733}
734
735#[cfg(test)]
736mod tests {
737 use super::*;
738 use approx::assert_abs_diff_eq;
739 use scirs2_core::ndarray::array;
740
741 #[test]
742 fn test_sparse_matrix_creation() {
743 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
744 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
745
746 assert_eq!(sparse.nrows, 3);
747 assert_eq!(sparse.ncols, 3);
748 assert_eq!(sparse.values.len(), 5); assert!(sparse.sparsity_ratio() < 1.0);
750 }
751
752 #[test]
753 fn test_sparse_matvec() {
754 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
755 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
756 let x = array![1.0, 2.0, 3.0];
757
758 let y_sparse = sparse.matvec(&x.view());
759 let y_dense = dense.dot(&x);
760
761 for i in 0..3 {
762 assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
763 }
764 }
765
766 #[test]
767 fn test_sparse_transpose_matvec() {
768 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
769 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
770 let x = array![1.0, 2.0, 3.0];
771
772 let y_sparse = sparse.transpose_matvec(&x.view());
773 let y_dense = dense.t().dot(&x);
774
775 for i in 0..3 {
776 assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
777 }
778 }
779
780 #[test]
781 fn test_soft_threshold() {
782 assert_abs_diff_eq!(soft_threshold(2.0, 1.0), 1.0, epsilon = 1e-12);
783 assert_abs_diff_eq!(soft_threshold(-2.0, 1.0), -1.0, epsilon = 1e-12);
784 assert_abs_diff_eq!(soft_threshold(0.5, 1.0), 0.0, epsilon = 1e-12);
785 assert_abs_diff_eq!(soft_threshold(-0.5, 1.0), 0.0, epsilon = 1e-12);
786 }
787
788 #[test]
789 fn test_sparse_least_squares_simple() {
790 let fun = |x: &ArrayView1<f64>| {
793 array![
794 x[0] - 1.0, x[1] - 2.0, x[0] + x[1] - 4.0 ]
798 };
799
800 let x0 = array![0.0, 0.0];
801 let options = SparseOptions {
802 max_iter: 1000,
803 tol: 1e-4,
804 lambda: 0.0, ..Default::default()
806 };
807
808 let result = sparse_least_squares(
809 fun,
810 None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
811 x0,
812 Some(options),
813 );
814 assert!(result.is_ok());
815
816 let result = result.unwrap();
817 assert!(result.cost < 10000.0); assert!(result.success); }
822
823 #[test]
824 fn test_lasso_coordinate_descent() {
825 let fun = |x: &ArrayView1<f64>| array![x[0] + 0.1 * x[1] - 1.0, x[1] - 0.0];
827
828 let x0 = array![0.0, 0.0];
829 let options = SparseOptions {
830 max_iter: 100,
831 tol: 1e-6,
832 lambda: 0.1, use_coordinate_descent: true,
834 ..Default::default()
835 };
836
837 let result = sparse_least_squares(
838 fun,
839 None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
840 x0,
841 Some(options),
842 );
843 assert!(result.is_ok());
844
845 let result = result.unwrap();
846 assert!(result.success);
849 }
850}