use crate::error::OptimizeError;
use crate::least_squares::Options;
use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
#[derive(Debug, Clone)]
pub struct SparseMatrix {
pub row_ptr: Vec<usize>,
pub col_idx: Vec<usize>,
pub values: Vec<f64>,
pub nrows: usize,
pub ncols: usize,
}
impl SparseMatrix {
pub fn new(nrows: usize, ncols: usize) -> Self {
Self {
row_ptr: vec![0; nrows + 1],
col_idx: Vec::new(),
values: Vec::new(),
nrows,
ncols,
}
}
pub fn from_dense(matrix: &ArrayView2<f64>, threshold: f64) -> Self {
let (nrows, ncols) = matrix.dim();
let mut row_ptr = vec![0; nrows + 1];
let mut col_idx = Vec::new();
let mut values = Vec::new();
for i in 0..nrows {
for j in 0..ncols {
if matrix[[i, j]].abs() > threshold {
col_idx.push(j);
values.push(matrix[[i, j]]);
}
}
row_ptr[i + 1] = values.len();
}
Self {
row_ptr,
col_idx,
values,
nrows,
ncols,
}
}
pub fn matvec(&self, x: &ArrayView1<f64>) -> Array1<f64> {
assert_eq!(x.len(), self.ncols);
let mut y = Array1::zeros(self.nrows);
for i in 0..self.nrows {
let start = self.row_ptr[i];
let end = self.row_ptr[i + 1];
let mut sum = 0.0;
for k in start..end {
sum += self.values[k] * x[self.col_idx[k]];
}
y[i] = sum;
}
y
}
pub fn transpose_matvec(&self, x: &ArrayView1<f64>) -> Array1<f64> {
assert_eq!(x.len(), self.nrows);
let mut y = Array1::zeros(self.ncols);
for i in 0..self.nrows {
let start = self.row_ptr[i];
let end = self.row_ptr[i + 1];
for k in start..end {
y[self.col_idx[k]] += self.values[k] * x[i];
}
}
y
}
pub fn sparsity_ratio(&self) -> f64 {
self.values.len() as f64 / (self.nrows * self.ncols) as f64
}
}
#[derive(Debug, Clone)]
pub struct SparseOptions {
pub base_options: Options,
pub sparsity_threshold: f64,
pub max_iter: usize,
pub tol: f64,
pub lambda: f64,
pub use_coordinate_descent: bool,
pub block_size: usize,
pub use_preconditioning: bool,
pub memory_limit_mb: usize,
}
impl Default for SparseOptions {
fn default() -> Self {
Self {
base_options: Options::default(),
sparsity_threshold: 1e-12,
max_iter: 1000,
tol: 1e-8,
lambda: 0.0,
use_coordinate_descent: false,
block_size: 100,
use_preconditioning: true,
memory_limit_mb: 1000,
}
}
}
#[derive(Debug, Clone)]
pub struct SparseResult {
pub x: Array1<f64>,
pub cost: f64,
pub residual: Array1<f64>,
pub nit: usize,
pub nfev: usize,
pub njev: usize,
pub success: bool,
pub message: String,
pub sparsity_info: SparseInfo,
}
#[derive(Debug, Clone)]
pub struct SparseInfo {
pub jacobian_sparsity: f64,
pub jacobian_nnz: usize,
pub memory_usage_mb: f64,
pub used_sparse_algorithms: bool,
}
#[allow(dead_code)]
pub fn sparse_least_squares<F, J>(
fun: F,
jac: Option<J>,
x0: Array1<f64>,
options: Option<SparseOptions>,
) -> Result<SparseResult, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64>,
J: Fn(&ArrayView1<f64>) -> Array2<f64>,
{
let options = options.unwrap_or_default();
let x = x0.clone();
let n = x.len();
let mut nfev = 0;
let mut njev = 0;
let residual = fun(&x.view());
nfev += 1;
let m = residual.len();
let jacobian = if let Some(ref jac_fn) = jac {
let jac_dense = jac_fn(&x.view());
njev += 1;
Some(jac_dense)
} else {
None
};
let (sparse_jac, sparsity_info) = if let Some(ref jac_matrix) = jacobian {
let sparse_matrix =
SparseMatrix::from_dense(&jac_matrix.view(), options.sparsity_threshold);
let sparsity_ratio = sparse_matrix.sparsity_ratio();
let memory_usage = estimate_memory_usage(&sparse_matrix);
let info = SparseInfo {
jacobian_sparsity: sparsity_ratio,
jacobian_nnz: sparse_matrix.values.len(),
memory_usage_mb: memory_usage,
used_sparse_algorithms: sparsity_ratio < 0.5, };
(Some(sparse_matrix), info)
} else {
let info = SparseInfo {
jacobian_sparsity: 1.0,
jacobian_nnz: m * n,
memory_usage_mb: (m * n * 8) as f64 / (1024.0 * 1024.0),
used_sparse_algorithms: false,
};
(None, info)
};
let result = if options.lambda > 0.0 {
if options.use_coordinate_descent {
solve_lasso_coordinate_descent(&fun, &x, &options, &mut nfev)?
} else {
solve_lasso_proximal_gradient(&fun, &x, &options, &mut nfev)?
}
} else if sparsity_info.used_sparse_algorithms && sparse_jac.is_some() {
solve_sparse_gauss_newton(
&fun,
&jac,
&sparse_jac.expect("Operation failed"),
&x,
&options,
&mut nfev,
&mut njev,
)?
} else {
solve_dense_least_squares(&fun, &jac, &x, &options, &mut nfev, &mut njev)?
};
Ok(SparseResult {
x: result.x,
cost: result.cost,
residual: result.residual,
nit: result.nit,
nfev,
njev,
success: result.success,
message: result.message,
sparsity_info,
})
}
#[derive(Debug)]
struct InternalResult {
x: Array1<f64>,
cost: f64,
residual: Array1<f64>,
nit: usize,
success: bool,
message: String,
}
#[allow(dead_code)]
fn solve_lasso_coordinate_descent<F>(
fun: &F,
x0: &Array1<f64>,
options: &SparseOptions,
nfev: &mut usize,
) -> Result<InternalResult, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64>,
{
let mut x = x0.clone();
let n = x.len();
let lambda = options.lambda;
for iter in 0..options.max_iter {
let mut max_change: f64 = 0.0;
for j in 0..n {
let old_xj = x[j];
let residual = fun(&x.view());
*nfev += 1;
let eps = 1e-8;
let mut x_plus = x.clone();
x_plus[j] += eps;
let residual_plus = fun(&x_plus.view());
*nfev += 1;
let grad_j = residual_plus
.iter()
.zip(residual.iter())
.map(|(&rp, &r)| (rp - r) / eps)
.sum::<f64>();
let step_size = options.base_options.gtol.unwrap_or(1e-5);
let z = x[j] - step_size * grad_j;
x[j] = soft_threshold(z, lambda * step_size);
max_change = max_change.max((x[j] - old_xj).abs());
}
if max_change < options.tol {
let final_residual = fun(&x.view());
*nfev += 1;
let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
+ lambda * x.mapv(|xi| xi.abs()).sum();
return Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: iter,
success: true,
message: "LASSO coordinate descent converged successfully".to_string(),
});
}
}
let final_residual = fun(&x.view());
*nfev += 1;
let cost =
0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: options.max_iter,
success: false,
message: "Maximum iterations reached in LASSO coordinate descent".to_string(),
})
}
#[allow(dead_code)]
fn soft_threshold(x: f64, threshold: f64) -> f64 {
if x > threshold {
x - threshold
} else if x < -threshold {
x + threshold
} else {
0.0
}
}
#[allow(dead_code)]
fn solve_lasso_proximal_gradient<F>(
fun: &F,
x0: &Array1<f64>,
options: &SparseOptions,
nfev: &mut usize,
) -> Result<InternalResult, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64>,
{
let mut x = x0.clone();
let lambda = options.lambda;
let step_size = 0.01;
for iter in 0..options.max_iter {
let residual = fun(&x.view());
*nfev += 1;
let mut grad = Array1::zeros(x.len());
let eps = 1e-8;
for i in 0..x.len() {
let mut x_plus = x.clone();
x_plus[i] += eps;
let residual_plus = fun(&x_plus.view());
*nfev += 1;
grad[i] = 2.0
* residual_plus
.iter()
.zip(residual.iter())
.map(|(&rp, &r)| r * (rp - r) / eps)
.sum::<f64>();
}
let x_old = x.clone();
for i in 0..x.len() {
let z = x[i] - step_size * grad[i];
x[i] = soft_threshold(z, lambda * step_size);
}
let change = (&x - &x_old).mapv(|dx| dx.abs()).sum();
if change < options.tol {
let final_residual = fun(&x.view());
*nfev += 1;
let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum()
+ lambda * x.mapv(|xi| xi.abs()).sum();
return Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: iter,
success: true,
message: "LASSO proximal gradient converged successfully".to_string(),
});
}
}
let final_residual = fun(&x.view());
*nfev += 1;
let cost =
0.5 * final_residual.mapv(|r| r.powi(2)).sum() + lambda * x.mapv(|xi| xi.abs()).sum();
Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: options.max_iter,
success: false,
message: "Maximum iterations reached in LASSO proximal gradient".to_string(),
})
}
#[allow(dead_code)]
fn solve_sparse_gauss_newton<F, J>(
fun: &F,
jac: &Option<J>,
_sparse_jac: &SparseMatrix,
x0: &Array1<f64>,
options: &SparseOptions,
nfev: &mut usize,
njev: &mut usize,
) -> Result<InternalResult, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64>,
J: Fn(&ArrayView1<f64>) -> Array2<f64>,
{
let mut x = x0.clone();
for iter in 0..options.max_iter {
let residual = fun(&x.view());
*nfev += 1;
if let Some(ref jac_fn) = jac {
let jac_dense = jac_fn(&x.view());
*njev += 1;
let jac_sparse =
SparseMatrix::from_dense(&jac_dense.view(), options.sparsity_threshold);
let jt_r = jac_sparse.transpose_matvec(&residual.view());
let mut p = Array1::zeros(x.len());
for i in 0..x.len() {
let diag_elem = compute_diagonal_element(&jac_sparse, i);
if diag_elem.abs() > 1e-12 {
p[i] = -jt_r[i] / diag_elem;
}
}
let mut alpha = 1.0;
let current_cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
for _ in 0..10 {
let x_new = &x + alpha * &p;
let residual_new = fun(&x_new.view());
*nfev += 1;
let new_cost = 0.5 * residual_new.mapv(|r| r.powi(2)).sum();
if new_cost < current_cost {
x = x_new;
break;
}
alpha *= 0.5;
}
let grad_norm = jt_r.mapv(|g| g.abs()).sum();
if grad_norm < options.tol {
let final_residual = fun(&x.view());
*nfev += 1;
let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
return Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: iter,
success: true,
message: "Sparse Gauss-Newton converged successfully".to_string(),
});
}
}
}
let final_residual = fun(&x.view());
*nfev += 1;
let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: options.max_iter,
success: false,
message: "Maximum iterations reached in sparse Gauss-Newton".to_string(),
})
}
#[allow(dead_code)]
fn compute_diagonal_element(jac: &SparseMatrix, col: usize) -> f64 {
let mut diag = 0.0;
for row in 0..jac.nrows {
let start = jac.row_ptr[row];
let end = jac.row_ptr[row + 1];
for k in start..end {
if jac.col_idx[k] == col {
diag += jac.values[k].powi(2);
}
}
}
diag
}
#[allow(dead_code)]
fn solve_dense_least_squares<F, J>(
fun: &F,
_jac: &Option<J>,
x0: &Array1<f64>,
options: &SparseOptions,
nfev: &mut usize,
_njev: &mut usize,
) -> Result<InternalResult, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64>,
J: Fn(&ArrayView1<f64>) -> Array2<f64>,
{
let mut x = x0.clone();
for iter in 0..options.max_iter {
let residual = fun(&x.view());
*nfev += 1;
let mut grad = Array1::zeros(x.len());
let eps = 1e-8;
for i in 0..x.len() {
let mut x_plus = x.clone();
x_plus[i] += eps;
let residual_plus = fun(&x_plus.view());
*nfev += 1;
grad[i] = 2.0
* residual_plus
.iter()
.zip(residual.iter())
.map(|(&rp, &r)| r * (rp - r) / eps)
.sum::<f64>();
}
let grad_norm = grad.mapv(|g| g.abs()).sum();
if grad_norm < options.tol {
let cost = 0.5 * residual.mapv(|r| r.powi(2)).sum();
return Ok(InternalResult {
x,
cost,
residual,
nit: iter,
success: true,
message: "Dense least squares converged successfully".to_string(),
});
}
let step_size = 0.01;
x = x - step_size * &grad;
}
let final_residual = fun(&x.view());
*nfev += 1;
let cost = 0.5 * final_residual.mapv(|r| r.powi(2)).sum();
Ok(InternalResult {
x,
cost,
residual: final_residual,
nit: options.max_iter,
success: false,
message: "Maximum iterations reached in dense least squares".to_string(),
})
}
#[allow(dead_code)]
fn estimate_memory_usage(sparse_matrix: &SparseMatrix) -> f64 {
let nnz = sparse_matrix.values.len();
let nrows = sparse_matrix.nrows;
let memory_bytes = nnz * 8 + nnz * 8 + (nrows + 1) * 8; memory_bytes as f64 / (1024.0 * 1024.0)
}
#[allow(dead_code)]
pub fn lsqr<F>(
matvec: F,
rmatvec: F,
b: &ArrayView1<f64>,
x0: Option<Array1<f64>>,
max_iter: Option<usize>,
tol: Option<f64>,
) -> Result<Array1<f64>, OptimizeError>
where
F: Fn(&ArrayView1<f64>) -> Array1<f64> + Clone,
{
let n = if let Some(ref x0_val) = x0 {
x0_val.len()
} else {
b.len()
};
let m = b.len();
let max_iter = max_iter.unwrap_or(n.min(m));
let tol = tol.unwrap_or(1e-8);
let mut x = x0.unwrap_or_else(|| Array1::zeros(n));
#[allow(unused_assignments)]
let mut beta = 0.0;
#[allow(unused_assignments)]
let mut u = Array1::zeros(m);
let ax = matvec.clone()(&x.view());
let r = b - &ax;
let mut v = rmatvec.clone()(&r.view());
let mut alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
if alpha == 0.0 {
return Ok(x);
}
v /= alpha;
let mut w = v.clone();
for _iter in 0..max_iter {
let av = matvec.clone()(&v.view());
let alpha_new = av.mapv(|avi| avi.powi(2)).sum().sqrt();
if alpha_new == 0.0 {
break;
}
u = av / alpha_new;
beta = alpha_new;
let atu = rmatvec.clone()(&u.view());
v = atu - beta * &v;
alpha = v.mapv(|vi| vi.powi(2)).sum().sqrt();
if alpha == 0.0 {
break;
}
v /= alpha;
x = x + (beta / alpha) * &w;
w = v.clone() - (alpha / beta) * &w;
let residual_norm = (b - &matvec.clone()(&x.view()))
.mapv(|ri| ri.powi(2))
.sum()
.sqrt();
if residual_norm < tol {
break;
}
}
Ok(x)
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_abs_diff_eq;
use scirs2_core::ndarray::array;
#[test]
fn test_sparse_matrix_creation() {
let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
assert_eq!(sparse.nrows, 3);
assert_eq!(sparse.ncols, 3);
assert_eq!(sparse.values.len(), 5); assert!(sparse.sparsity_ratio() < 1.0);
}
#[test]
fn test_sparse_matvec() {
let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
let x = array![1.0, 2.0, 3.0];
let y_sparse = sparse.matvec(&x.view());
let y_dense = dense.dot(&x);
for i in 0..3 {
assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
}
}
#[test]
fn test_sparse_transpose_matvec() {
let dense = array![[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [4.0, 0.0, 5.0]];
let sparse = SparseMatrix::from_dense(&dense.view(), 1e-12);
let x = array![1.0, 2.0, 3.0];
let y_sparse = sparse.transpose_matvec(&x.view());
let y_dense = dense.t().dot(&x);
for i in 0..3 {
assert_abs_diff_eq!(y_sparse[i], y_dense[i], epsilon = 1e-12);
}
}
#[test]
fn test_soft_threshold() {
assert_abs_diff_eq!(soft_threshold(2.0, 1.0), 1.0, epsilon = 1e-12);
assert_abs_diff_eq!(soft_threshold(-2.0, 1.0), -1.0, epsilon = 1e-12);
assert_abs_diff_eq!(soft_threshold(0.5, 1.0), 0.0, epsilon = 1e-12);
assert_abs_diff_eq!(soft_threshold(-0.5, 1.0), 0.0, epsilon = 1e-12);
}
#[test]
fn test_sparse_least_squares_simple() {
let fun = |x: &ArrayView1<f64>| {
array![
x[0] - 1.0, x[1] - 2.0, x[0] + x[1] - 4.0 ]
};
let x0 = array![0.0, 0.0];
let options = SparseOptions {
max_iter: 1000,
tol: 1e-4,
lambda: 0.0, ..Default::default()
};
let result = sparse_least_squares(
fun,
None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
x0,
Some(options),
);
assert!(result.is_ok());
let result = result.expect("Operation failed");
assert!(result.cost < 10000.0); assert!(result.success); }
#[test]
fn test_lasso_coordinate_descent() {
let fun = |x: &ArrayView1<f64>| array![x[0] + 0.1 * x[1] - 1.0, x[1] - 0.0];
let x0 = array![0.0, 0.0];
let options = SparseOptions {
max_iter: 100,
tol: 1e-6,
lambda: 0.1, use_coordinate_descent: true,
..Default::default()
};
let result = sparse_least_squares(
fun,
None::<fn(&ArrayView1<f64>) -> Array2<f64>>,
x0,
Some(options),
);
assert!(result.is_ok());
let result = result.expect("Operation failed");
assert!(result.success);
}
}