1use crate::error::OptimizeError;
8use crate::least_squares::Options;
9use 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
181pub fn sparse_least_squares<F, J>(
183 fun: F,
184 jac: Option<J>,
185 x0: Array1<f64>,
186 options: Option<SparseOptions>,
187) -> Result<SparseResult, OptimizeError>
188where
189 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
190 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
191{
192 let options = options.unwrap_or_default();
193 let x = x0.clone();
194 let n = x.len();
195 let mut nfev = 0;
196 let mut njev = 0;
197
198 let residual = fun(&x.view());
200 nfev += 1;
201 let m = residual.len();
202
203 let jacobian = if let Some(ref jac_fn) = jac {
205 let jac_dense = jac_fn(&x.view());
206 njev += 1;
207 Some(jac_dense)
208 } else {
209 None
210 };
211
212 let (sparse_jac, sparsity_info) = if let Some(ref jac_matrix) = jacobian {
213 let sparse_matrix =
214 SparseMatrix::from_dense(&jac_matrix.view(), options.sparsity_threshold);
215 let sparsity_ratio = sparse_matrix.sparsity_ratio();
216 let memory_usage = estimate_memory_usage(&sparse_matrix);
217
218 let info = SparseInfo {
219 jacobian_sparsity: sparsity_ratio,
220 jacobian_nnz: sparse_matrix.values.len(),
221 memory_usage_mb: memory_usage,
222 used_sparse_algorithms: sparsity_ratio < 0.5, };
224
225 (Some(sparse_matrix), info)
226 } else {
227 let info = SparseInfo {
228 jacobian_sparsity: 1.0,
229 jacobian_nnz: m * n,
230 memory_usage_mb: (m * n * 8) as f64 / (1024.0 * 1024.0),
231 used_sparse_algorithms: false,
232 };
233 (None, info)
234 };
235
236 let result = if options.lambda > 0.0 {
238 if options.use_coordinate_descent {
240 solve_lasso_coordinate_descent(&fun, &x, &options, &mut nfev)?
241 } else {
242 solve_lasso_proximal_gradient(&fun, &x, &options, &mut nfev)?
243 }
244 } else if sparsity_info.used_sparse_algorithms && sparse_jac.is_some() {
245 solve_sparse_gauss_newton(
247 &fun,
248 &jac,
249 &sparse_jac.unwrap(),
250 &x,
251 &options,
252 &mut nfev,
253 &mut njev,
254 )?
255 } else {
256 solve_dense_least_squares(&fun, &jac, &x, &options, &mut nfev, &mut njev)?
258 };
259
260 Ok(SparseResult {
261 x: result.x,
262 cost: result.cost,
263 residual: result.residual,
264 nit: result.nit,
265 nfev,
266 njev,
267 success: result.success,
268 message: result.message,
269 sparsity_info,
270 })
271}
272
273#[derive(Debug)]
275struct InternalResult {
276 x: Array1<f64>,
277 cost: f64,
278 residual: Array1<f64>,
279 nit: usize,
280 success: bool,
281 message: String,
282}
283
284fn solve_lasso_coordinate_descent<F>(
286 fun: &F,
287 x0: &Array1<f64>,
288 options: &SparseOptions,
289 nfev: &mut usize,
290) -> Result<InternalResult, OptimizeError>
291where
292 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
293{
294 let mut x = x0.clone();
295 let n = x.len();
296 let lambda = options.lambda;
297
298 for iter in 0..options.max_iter {
299 let mut max_change: f64 = 0.0;
300
301 for j in 0..n {
303 let old_xj = x[j];
304
305 let residual = fun(&x.view());
307 *nfev += 1;
308
309 let eps = 1e-8;
311 let mut x_plus = x.clone();
312 x_plus[j] += eps;
313 let residual_plus = fun(&x_plus.view());
314 *nfev += 1;
315
316 let grad_j = residual_plus
317 .iter()
318 .zip(residual.iter())
319 .map(|(&rp, &r)| (rp - r) / eps)
320 .sum::<f64>();
321
322 let step_size = options.base_options.gtol.unwrap_or(1e-5);
324 let z = x[j] - step_size * grad_j;
325 x[j] = soft_threshold(z, lambda * step_size);
326
327 max_change = max_change.max((x[j] - old_xj).abs());
328 }
329
330 if max_change < options.tol {
332 let final_residual = fun(&x.view());
333 *nfev += 1;
334 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
335 + lambda * x.mapv(|xi| xi.abs()).sum();
336
337 return Ok(InternalResult {
338 x,
339 cost,
340 residual: final_residual,
341 nit: iter,
342 success: true,
343 message: "LASSO coordinate descent converged successfully".to_string(),
344 });
345 }
346 }
347
348 let final_residual = fun(&x.view());
349 *nfev += 1;
350 let cost =
351 0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
352
353 Ok(InternalResult {
354 x,
355 cost,
356 residual: final_residual,
357 nit: options.max_iter,
358 success: false,
359 message: "Maximum iterations reached in LASSO coordinate descent".to_string(),
360 })
361}
362
363fn soft_threshold(x: f64, threshold: f64) -> f64 {
365 if x > threshold {
366 x - threshold
367 } else if x < -threshold {
368 x + threshold
369 } else {
370 0.0
371 }
372}
373
374fn solve_lasso_proximal_gradient<F>(
376 fun: &F,
377 x0: &Array1<f64>,
378 options: &SparseOptions,
379 nfev: &mut usize,
380) -> Result<InternalResult, OptimizeError>
381where
382 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
383{
384 let mut x = x0.clone();
385 let lambda = options.lambda;
386 let step_size = 0.01; for iter in 0..options.max_iter {
389 let residual = fun(&x.view());
391 *nfev += 1;
392
393 let mut grad = Array1::zeros(x.len());
394 let eps = 1e-8;
395
396 for i in 0..x.len() {
397 let mut x_plus = x.clone();
398 x_plus[i] += eps;
399 let residual_plus = fun(&x_plus.view());
400 *nfev += 1;
401
402 grad[i] = 2.0
405 * residual_plus
406 .iter()
407 .zip(residual.iter())
408 .map(|(&rp, &r)| r * (rp - r) / eps)
409 .sum::<f64>();
410 }
411
412 let x_old = x.clone();
414 for i in 0..x.len() {
415 let z = x[i] - step_size * grad[i];
416 x[i] = soft_threshold(z, lambda * step_size);
417 }
418
419 let change = (&x - &x_old).mapv(|dx| dx.abs()).sum();
421 if change < options.tol {
422 let final_residual = fun(&x.view());
423 *nfev += 1;
424 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
425 + lambda * x.mapv(|xi| xi.abs()).sum();
426
427 return Ok(InternalResult {
428 x,
429 cost,
430 residual: final_residual,
431 nit: iter,
432 success: true,
433 message: "LASSO proximal gradient converged successfully".to_string(),
434 });
435 }
436 }
437
438 let final_residual = fun(&x.view());
439 *nfev += 1;
440 let cost =
441 0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
442
443 Ok(InternalResult {
444 x,
445 cost,
446 residual: final_residual,
447 nit: options.max_iter,
448 success: false,
449 message: "Maximum iterations reached in LASSO proximal gradient".to_string(),
450 })
451}
452
453fn solve_sparse_gauss_newton<F, J>(
455 fun: &F,
456 jac: &Option<J>,
457 _sparse_jac: &SparseMatrix,
458 x0: &Array1<f64>,
459 options: &SparseOptions,
460 nfev: &mut usize,
461 njev: &mut usize,
462) -> Result<InternalResult, OptimizeError>
463where
464 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
465 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
466{
467 let mut x = x0.clone();
468
469 for iter in 0..options.max_iter {
470 let residual = fun(&x.view());
472 *nfev += 1;
473
474 if let Some(ref jac_fn) = jac {
475 let jac_dense = jac_fn(&x.view());
476 *njev += 1;
477
478 let jac_sparse =
480 SparseMatrix::from_dense(&jac_dense.view(), options.sparsity_threshold);
481
482 let jt_r = jac_sparse.transpose_matvec(&residual.view());
485
486 let mut p = Array1::zeros(x.len());
488 for i in 0..x.len() {
489 let diag_elem = compute_diagonal_element(&jac_sparse, i);
491 if diag_elem.abs() > 1e-12 {
492 p[i] = -jt_r[i] / diag_elem;
493 }
494 }
495
496 let mut alpha = 1.0;
498 let current_cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
499
500 for _ in 0..10 {
501 let x_new = &x + alpha * &p;
502 let residual_new = fun(&x_new.view());
503 *nfev += 1;
504 let new_cost = 0.5 * residual_new.mapv(|r| r.powi(2)).sum();
505
506 if new_cost < current_cost {
507 x = x_new;
508 break;
509 }
510 alpha *= 0.5;
511 }
512
513 let grad_norm = jt_r.mapv(|g| g.abs()).sum();
515 if grad_norm < options.tol {
516 let final_residual = fun(&x.view());
517 *nfev += 1;
518 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
519
520 return Ok(InternalResult {
521 x,
522 cost,
523 residual: final_residual,
524 nit: iter,
525 success: true,
526 message: "Sparse Gauss-Newton converged successfully".to_string(),
527 });
528 }
529 }
530 }
531
532 let final_residual = fun(&x.view());
533 *nfev += 1;
534 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
535
536 Ok(InternalResult {
537 x,
538 cost,
539 residual: final_residual,
540 nit: options.max_iter,
541 success: false,
542 message: "Maximum iterations reached in sparse Gauss-Newton".to_string(),
543 })
544}
545
546fn compute_diagonal_element(jac: &SparseMatrix, col: usize) -> f64 {
548 let mut diag = 0.0;
549
550 for row in 0..jac.nrows {
551 let start = jac.row_ptr[row];
552 let end = jac.row_ptr[row + 1];
553
554 for k in start..end {
555 if jac.col_idx[k] == col {
556 diag += jac.values[k].powi(2);
557 }
558 }
559 }
560
561 diag
562}
563
564fn solve_dense_least_squares<F, J>(
566 fun: &F,
567 _jac: &Option<J>,
568 x0: &Array1<f64>,
569 options: &SparseOptions,
570 nfev: &mut usize,
571 _njev: &mut usize,
572) -> Result<InternalResult, OptimizeError>
573where
574 F: Fn(&ArrayView1<f64>) -> Array1<f64>,
575 J: Fn(&ArrayView1<f64>) -> Array2<f64>,
576{
577 let mut x = x0.clone();
579
580 for iter in 0..options.max_iter {
581 let residual = fun(&x.view());
582 *nfev += 1;
583
584 let mut grad = Array1::zeros(x.len());
586 let eps = 1e-8;
587
588 for i in 0..x.len() {
589 let mut x_plus = x.clone();
590 x_plus[i] += eps;
591 let residual_plus = fun(&x_plus.view());
592 *nfev += 1;
593
594 grad[i] = 2.0
597 * residual_plus
598 .iter()
599 .zip(residual.iter())
600 .map(|(&rp, &r)| r * (rp - r) / eps)
601 .sum::<f64>();
602 }
603
604 let grad_norm = grad.mapv(|g| g.abs()).sum();
606 if grad_norm < options.tol {
607 let cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
608 return Ok(InternalResult {
609 x,
610 cost,
611 residual,
612 nit: iter,
613 success: true,
614 message: "Dense least squares converged successfully".to_string(),
615 });
616 }
617
618 let step_size = 0.01;
620 x = x - step_size * &grad;
621 }
622
623 let final_residual = fun(&x.view());
624 *nfev += 1;
625 let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
626
627 Ok(InternalResult {
628 x,
629 cost,
630 residual: final_residual,
631 nit: options.max_iter,
632 success: false,
633 message: "Maximum iterations reached in dense least squares".to_string(),
634 })
635}
636
637fn estimate_memory_usage(sparse_matrix: &SparseMatrix) -> f64 {
639 let nnz = sparse_matrix.values.len();
640 let nrows = sparse_matrix.nrows;
641
642 let memory_bytes = nnz * 8 + nnz * 8 + (nrows + 1) * 8; memory_bytes as f64 / (1024.0 * 1024.0)
645}
646
647pub fn lsqr<F>(
649 matvec: F,
650 rmatvec: F,
651 b: &ArrayView1<f64>,
652 x0: Option<Array1<f64>>,
653 max_iter: Option<usize>,
654 tol: Option<f64>,
655) -> Result<Array1<f64>, OptimizeError>
656where
657 F: Fn(&ArrayView1<f64>) -> Array1<f64> + Clone,
658{
659 let n = if let Some(ref x0_val) = x0 {
660 x0_val.len()
661 } else {
662 b.len()
663 };
664 let m = b.len();
665 let max_iter = max_iter.unwrap_or(n.min(m));
666 let tol = tol.unwrap_or(1e-8);
667
668 let mut x = x0.unwrap_or_else(|| Array1::zeros(n));
669 #[allow(unused_assignments)]
670 let mut beta = 0.0;
671 #[allow(unused_assignments)]
672 let mut u = Array1::zeros(m);
673
674 let ax = matvec.clone()(&x.view());
676 let r = b - &ax;
677 let mut v = rmatvec.clone()(&r.view());
678
679 let mut alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
680 if alpha == 0.0 {
681 return Ok(x);
682 }
683
684 v /= alpha;
685 let mut w = v.clone();
686
687 for _iter in 0..max_iter {
688 let av = matvec.clone()(&v.view());
690 let alpha_new = av.mapv(|avi| avi.powi(2)).sum().sqrt();
691
692 if alpha_new == 0.0 {
693 break;
694 }
695
696 u = av / alpha_new;
697 beta = alpha_new;
698
699 let atu = rmatvec.clone()(&u.view());
700 v = atu - beta * &v;
701 alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
702
703 if alpha == 0.0 {
704 break;
705 }
706
707 v /= alpha;
708
709 x = x + (beta / alpha) * &w;
711 w = v.clone() - (alpha / beta) * &w;
712
713 let residual_norm = (b - &matvec.clone()(&x.view()))
715 .mapv(|ri| ri.powi(2))
716 .sum()
717 .sqrt();
718 if residual_norm < tol {
719 break;
720 }
721 }
722
723 Ok(x)
724}
725
726#[cfg(test)]
727mod tests {
728 use super::*;
729 use approx::assert_abs_diff_eq;
730 use ndarray::array;
731
732 #[test]
733 fn test_sparse_matrix_creation() {
734 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
735 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
736
737 assert_eq!(sparse.nrows, 3);
738 assert_eq!(sparse.ncols, 3);
739 assert_eq!(sparse.values.len(), 5); assert!(sparse.sparsity_ratio() < 1.0);
741 }
742
743 #[test]
744 fn test_sparse_matvec() {
745 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
746 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
747 let x = array![1.0, 2.0, 3.0];
748
749 let y_sparse = sparse.matvec(&x.view());
750 let y_dense = dense.dot(&x);
751
752 for i in 0..3 {
753 assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
754 }
755 }
756
757 #[test]
758 fn test_sparse_transpose_matvec() {
759 let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
760 let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
761 let x = array![1.0, 2.0, 3.0];
762
763 let y_sparse = sparse.transpose_matvec(&x.view());
764 let y_dense = dense.t().dot(&x);
765
766 for i in 0..3 {
767 assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
768 }
769 }
770
771 #[test]
772 fn test_soft_threshold() {
773 assert_abs_diff_eq!(soft_threshold(2.0, 1.0), 1.0, epsilon = 1e-12);
774 assert_abs_diff_eq!(soft_threshold(-2.0, 1.0), -1.0, epsilon = 1e-12);
775 assert_abs_diff_eq!(soft_threshold(0.5, 1.0), 0.0, epsilon = 1e-12);
776 assert_abs_diff_eq!(soft_threshold(-0.5, 1.0), 0.0, epsilon = 1e-12);
777 }
778
779 #[test]
780 fn test_sparse_least_squares_simple() {
781 let fun = |x: &ArrayView1<f64>| {
784 array![
785 x[0] - 1.0, x[1] - 2.0, x[0] + x[1] - 4.0 ]
789 };
790
791 let x0 = array![0.0, 0.0];
792 let options = SparseOptions {
793 max_iter: 1000,
794 tol: 1e-4,
795 lambda: 0.0, ..Default::default()
797 };
798
799 let result = sparse_least_squares(
800 fun,
801 None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
802 x0,
803 Some(options),
804 );
805 assert!(result.is_ok());
806
807 let result = result.unwrap();
808 assert!(result.cost < 10000.0); assert!(result.success); }
813
814 #[test]
815 fn test_lasso_coordinate_descent() {
816 let fun = |x: &ArrayView1<f64>| array![x[0] + 0.1 * x[1] - 1.0, x[1] - 0.0];
818
819 let x0 = array![0.0, 0.0];
820 let options = SparseOptions {
821 max_iter: 100,
822 tol: 1e-6,
823 lambda: 0.1, use_coordinate_descent: true,
825 ..Default::default()
826 };
827
828 let result = sparse_least_squares(
829 fun,
830 None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
831 x0,
832 Some(options),
833 );
834 assert!(result.is_ok());
835
836 let result = result.unwrap();
837 assert!(result.success);
840 }
841}