diff --git a/lax/src/least_squares.rs b/lax/src/least_squares.rs index 1bfd4d37..d0bb7def 100644 --- a/lax/src/least_squares.rs +++ b/lax/src/least_squares.rs @@ -5,154 +5,307 @@ use cauchy::*; use num_traits::{ToPrimitive, Zero}; /// Result of LeastSquares -pub struct LeastSquaresOutput { +pub struct LeastSquaresOwned { /// singular values pub singular_values: Vec, /// The rank of the input matrix A pub rank: i32, } -#[cfg_attr(doc, katexit::katexit)] -/// Solve least square problem -pub trait LeastSquaresSvdDivideConquer_: Scalar { - /// Compute a vector $x$ which minimizes Euclidian norm $\| Ax - b\|$ - /// for a given matrix $A$ and a vector $b$. - fn least_squares( - a_layout: MatrixLayout, - a: &mut [Self], - b: &mut [Self], - ) -> Result>; - - /// Solve least square problems $\argmin_X \| AX - B\|$ - fn least_squares_nrhs( - a_layout: MatrixLayout, - a: &mut [Self], - b_layout: MatrixLayout, - b: &mut [Self], - ) -> Result>; +/// Result of LeastSquares +pub struct LeastSquaresRef<'work, A: Scalar> { + /// singular values + pub singular_values: &'work [A::Real], + /// The rank of the input matrix A + pub rank: i32, } -macro_rules! impl_least_squares { - (@real, $scalar:ty, $gelsd:path) => { - impl_least_squares!(@body, $scalar, $gelsd, ); - }; - (@complex, $scalar:ty, $gelsd:path) => { - impl_least_squares!(@body, $scalar, $gelsd, rwork); - }; +pub struct LeastSquaresWork { + pub a_layout: MatrixLayout, + pub b_layout: MatrixLayout, + pub singular_values: Vec>, + pub work: Vec>, + pub iwork: Vec>, + pub rwork: Option>>, +} - (@body, $scalar:ty, $gelsd:path, $($rwork:ident),*) => { - impl LeastSquaresSvdDivideConquer_ for $scalar { - fn least_squares( - l: MatrixLayout, - a: &mut [Self], - b: &mut [Self], - ) -> Result> { - let b_layout = l.resized(b.len() as i32, 1); - Self::least_squares_nrhs(l, a, b_layout, b) - } +pub trait LeastSquaresWorkImpl: Sized { + type Elem: Scalar; + fn new(a_layout: MatrixLayout, b_layout: MatrixLayout) -> Result; + fn calc( + &mut self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result>; + fn eval( + self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result>; +} + +macro_rules! impl_least_squares_work_c { + ($c:ty, $lsd:path) => { + impl LeastSquaresWorkImpl for LeastSquaresWork<$c> { + type Elem = $c; - fn least_squares_nrhs( - a_layout: MatrixLayout, - a: &mut [Self], - b_layout: MatrixLayout, - b: &mut [Self], - ) -> Result> { - // Minimize |b - Ax|_2 - // - // where - // A : (m, n) - // b : (max(m, n), nrhs) // `b` has to store `x` on exit - // x : (n, nrhs) + fn new(a_layout: MatrixLayout, b_layout: MatrixLayout) -> Result { let (m, n) = a_layout.size(); let (m_, nrhs) = b_layout.size(); let k = m.min(n); assert!(m_ >= m); + let rcond = -1.; + let mut singular_values = vec_uninit(k as usize); + let mut rank: i32 = 0; + + // eval work size + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + let mut iwork_size = [0]; + let mut rwork = [::Real::zero()]; + unsafe { + $lsd( + &m, + &n, + &nrhs, + std::ptr::null_mut(), + &m, + std::ptr::null_mut(), + &m_, + AsPtr::as_mut_ptr(&mut singular_values), + &rcond, + &mut rank, + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + iwork_size.as_mut_ptr(), + &mut info, + ) + }; + info.as_lapack_result()?; + + let lwork = work_size[0].to_usize().unwrap(); + let liwork = iwork_size[0].to_usize().unwrap(); + let lrwork = rwork[0].to_usize().unwrap(); + + let work = vec_uninit(lwork); + let iwork = vec_uninit(liwork); + let rwork = vec_uninit(lrwork); + + Ok(LeastSquaresWork { + a_layout, + b_layout, + work, + iwork, + rwork: Some(rwork), + singular_values, + }) + } + + fn calc( + &mut self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result> { + let (m, n) = self.a_layout.size(); + let (m_, nrhs) = self.b_layout.size(); + assert!(m_ >= m); + + let lwork = self.work.len().to_i32().unwrap(); + // Transpose if a is C-continuous let mut a_t = None; - let a_layout = match a_layout { + let _ = match self.a_layout { MatrixLayout::C { .. } => { - let (layout, t) = transpose(a_layout, a); + let (layout, t) = transpose(self.a_layout, a); a_t = Some(t); layout } - MatrixLayout::F { .. } => a_layout, + MatrixLayout::F { .. } => self.a_layout, }; // Transpose if b is C-continuous let mut b_t = None; - let b_layout = match b_layout { + let b_layout = match self.b_layout { MatrixLayout::C { .. } => { - let (layout, t) = transpose(b_layout, b); + let (layout, t) = transpose(self.b_layout, b); b_t = Some(t); layout } - MatrixLayout::F { .. } => b_layout, + MatrixLayout::F { .. } => self.b_layout, }; - let rcond: Self::Real = -1.; - let mut singular_values: Vec> = vec_uninit( k as usize); + let rcond: ::Real = -1.; let mut rank: i32 = 0; - // eval work size let mut info = 0; - let mut work_size = [Self::zero()]; - let mut iwork_size = [0]; - $( - let mut $rwork = [Self::Real::zero()]; - )* unsafe { - $gelsd( + $lsd( &m, &n, &nrhs, AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), - &a_layout.lda(), + &m, AsPtr::as_mut_ptr(b_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(b)), - &b_layout.lda(), + &m_, + AsPtr::as_mut_ptr(&mut self.singular_values), + &rcond, + &mut rank, + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + AsPtr::as_mut_ptr(&mut self.iwork), + &mut info, + ); + } + info.as_lapack_result()?; + + let singular_values = unsafe { self.singular_values.slice_assume_init_ref() }; + + // Skip a_t -> a transpose because A has been destroyed + // Re-transpose b + if let Some(b_t) = b_t { + transpose_over(b_layout, &b_t, b); + } + + Ok(LeastSquaresRef { + singular_values, + rank, + }) + } + + fn eval( + mut self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result> { + let LeastSquaresRef { rank, .. } = self.calc(a, b)?; + let singular_values = unsafe { self.singular_values.assume_init() }; + Ok(LeastSquaresOwned { + singular_values, + rank, + }) + } + } + }; +} +impl_least_squares_work_c!(c64, lapack_sys::zgelsd_); +impl_least_squares_work_c!(c32, lapack_sys::cgelsd_); + +macro_rules! impl_least_squares_work_r { + ($c:ty, $lsd:path) => { + impl LeastSquaresWorkImpl for LeastSquaresWork<$c> { + type Elem = $c; + + fn new(a_layout: MatrixLayout, b_layout: MatrixLayout) -> Result { + let (m, n) = a_layout.size(); + let (m_, nrhs) = b_layout.size(); + let k = m.min(n); + assert!(m_ >= m); + + let rcond = -1.; + let mut singular_values = vec_uninit(k as usize); + let mut rank: i32 = 0; + + // eval work size + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + let mut iwork_size = [0]; + unsafe { + $lsd( + &m, + &n, + &nrhs, + std::ptr::null_mut(), + &m, + std::ptr::null_mut(), + &m_, AsPtr::as_mut_ptr(&mut singular_values), &rcond, &mut rank, AsPtr::as_mut_ptr(&mut work_size), &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork),)* iwork_size.as_mut_ptr(), &mut info, ) }; info.as_lapack_result()?; - // calc let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); let liwork = iwork_size[0].to_usize().unwrap(); - let mut iwork: Vec> = vec_uninit(liwork); - $( - let lrwork = $rwork[0].to_usize().unwrap(); - let mut $rwork: Vec> = vec_uninit(lrwork); - )* + + let work = vec_uninit(lwork); + let iwork = vec_uninit(liwork); + + Ok(LeastSquaresWork { + a_layout, + b_layout, + work, + iwork, + rwork: None, + singular_values, + }) + } + + fn calc( + &mut self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result> { + let (m, n) = self.a_layout.size(); + let (m_, nrhs) = self.b_layout.size(); + assert!(m_ >= m); + + let lwork = self.work.len().to_i32().unwrap(); + + // Transpose if a is C-continuous + let mut a_t = None; + let _ = match self.a_layout { + MatrixLayout::C { .. } => { + let (layout, t) = transpose(self.a_layout, a); + a_t = Some(t); + layout + } + MatrixLayout::F { .. } => self.a_layout, + }; + + // Transpose if b is C-continuous + let mut b_t = None; + let b_layout = match self.b_layout { + MatrixLayout::C { .. } => { + let (layout, t) = transpose(self.b_layout, b); + b_t = Some(t); + layout + } + MatrixLayout::F { .. } => self.b_layout, + }; + + let rcond: ::Real = -1.; + let mut rank: i32 = 0; + + let mut info = 0; unsafe { - $gelsd( + $lsd( &m, &n, &nrhs, AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), - &a_layout.lda(), + &m, AsPtr::as_mut_ptr(b_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(b)), - &b_layout.lda(), - AsPtr::as_mut_ptr(&mut singular_values), + &m_, + AsPtr::as_mut_ptr(&mut self.singular_values), &rcond, &mut rank, - AsPtr::as_mut_ptr(&mut work), - &(lwork as i32), - $(AsPtr::as_mut_ptr(&mut $rwork),)* - AsPtr::as_mut_ptr(&mut iwork), + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(&mut self.iwork), &mut info, ); } info.as_lapack_result()?; - let singular_values = unsafe { singular_values.assume_init() }; + let singular_values = unsafe { self.singular_values.slice_assume_init_ref() }; // Skip a_t -> a transpose because A has been destroyed // Re-transpose b @@ -160,7 +313,20 @@ macro_rules! impl_least_squares { transpose_over(b_layout, &b_t, b); } - Ok(LeastSquaresOutput { + Ok(LeastSquaresRef { + singular_values, + rank, + }) + } + + fn eval( + mut self, + a: &mut [Self::Elem], + b: &mut [Self::Elem], + ) -> Result> { + let LeastSquaresRef { rank, .. } = self.calc(a, b)?; + let singular_values = unsafe { self.singular_values.assume_init() }; + Ok(LeastSquaresOwned { singular_values, rank, }) @@ -168,8 +334,5 @@ macro_rules! impl_least_squares { } }; } - -impl_least_squares!(@real, f64, lapack_sys::dgelsd_); -impl_least_squares!(@real, f32, lapack_sys::sgelsd_); -impl_least_squares!(@complex, c64, lapack_sys::zgelsd_); -impl_least_squares!(@complex, c32, lapack_sys::cgelsd_); +impl_least_squares_work_r!(f64, lapack_sys::dgelsd_); +impl_least_squares_work_r!(f32, lapack_sys::sgelsd_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index ffe1fb62..199f2dc2 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -65,11 +65,9 @@ //! Singular Value Decomposition //! ----------------------------- //! -//! - [SVD_] trait provides methods for singular value decomposition for general matrix -//! - [SVDDC_] trait provides methods for singular value decomposition for general matrix -//! with divided-and-conquer algorithm -//! - [LeastSquaresSvdDivideConquer_] trait provides methods -//! for solving least square problem by SVD +//! - [svd] module for singular value decomposition (SVD) for general matrix +//! - [svddc] module for singular value decomposition (SVD) with divided-and-conquer algorithm for general matrix +//! - [least_squares] module for solving least square problem using SVD //! #![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)] @@ -90,29 +88,28 @@ pub mod layout; pub mod eig; pub mod eigh; pub mod eigh_generalized; +pub mod least_squares; pub mod qr; +pub mod svd; +pub mod svddc; mod alloc; mod cholesky; -mod least_squares; mod opnorm; mod rcond; mod solve; mod solveh; -mod svd; -mod svddc; mod triangular; mod tridiagonal; pub use self::cholesky::*; pub use self::flags::*; -pub use self::least_squares::*; +pub use self::least_squares::LeastSquaresOwned; pub use self::opnorm::*; pub use self::rcond::*; pub use self::solve::*; pub use self::solveh::*; -pub use self::svd::*; -pub use self::svddc::*; +pub use self::svd::{SvdOwned, SvdRef}; pub use self::triangular::*; pub use self::tridiagonal::*; @@ -122,18 +119,10 @@ use std::mem::MaybeUninit; pub type Pivot = Vec; +#[cfg_attr(doc, katexit::katexit)] /// Trait for primitive types which implements LAPACK subroutines pub trait Lapack: - OperatorNorm_ - + SVD_ - + SVDDC_ - + Solve_ - + Solveh_ - + Cholesky_ - + Triangular_ - + Tridiagonal_ - + Rcond_ - + LeastSquaresSvdDivideConquer_ + OperatorNorm_ + Solve_ + Solveh_ + Cholesky_ + Triangular_ + Tridiagonal_ + Rcond_ { /// Compute right eigenvalue and eigenvectors for a general matrix fn eig( @@ -170,6 +159,28 @@ pub trait Lapack: /// Execute QR-decomposition at once fn qr(l: MatrixLayout, a: &mut [Self]) -> Result>; + + /// Compute singular-value decomposition (SVD) + fn svd(l: MatrixLayout, calc_u: bool, calc_vt: bool, a: &mut [Self]) -> Result>; + + /// Compute singular value decomposition (SVD) with divide-and-conquer algorithm + fn svddc(layout: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; + + /// Compute a vector $x$ which minimizes Euclidian norm $\| Ax - b\|$ + /// for a given matrix $A$ and a vector $b$. + fn least_squares( + a_layout: MatrixLayout, + a: &mut [Self], + b: &mut [Self], + ) -> Result>; + + /// Solve least square problems $\argmin_X \| AX - B\|$ + fn least_squares_nrhs( + a_layout: MatrixLayout, + a: &mut [Self], + b_layout: MatrixLayout, + b: &mut [Self], + ) -> Result>; } macro_rules! impl_lapack { @@ -228,6 +239,43 @@ macro_rules! impl_lapack { Self::q(l, a, &tau)?; Ok(r) } + + fn svd( + l: MatrixLayout, + calc_u: bool, + calc_vt: bool, + a: &mut [Self], + ) -> Result> { + use svd::*; + let work = SvdWork::<$s>::new(l, calc_u, calc_vt)?; + work.eval(a) + } + + fn svddc(layout: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result> { + use svddc::*; + let work = SvdDcWork::<$s>::new(layout, jobz)?; + work.eval(a) + } + + fn least_squares( + l: MatrixLayout, + a: &mut [Self], + b: &mut [Self], + ) -> Result> { + let b_layout = l.resized(b.len() as i32, 1); + Self::least_squares_nrhs(l, a, b_layout, b) + } + + fn least_squares_nrhs( + a_layout: MatrixLayout, + a: &mut [Self], + b_layout: MatrixLayout, + b: &mut [Self], + ) -> Result> { + use least_squares::*; + let work = LeastSquaresWork::<$s>::new(a_layout, b_layout)?; + work.eval(a, b) + } } }; } diff --git a/lax/src/svd.rs b/lax/src/svd.rs index 0d9bd353..fc695108 100644 --- a/lax/src/svd.rs +++ b/lax/src/svd.rs @@ -1,85 +1,92 @@ //! Singular-value decomposition +//! +//! LAPACK correspondance +//! ---------------------- +//! +//! | f32 | f64 | c32 | c64 | +//! |:-------|:-------|:-------|:-------| +//! | sgesvd | dgesvd | cgesvd | zgesvd | +//! use super::{error::*, layout::*, *}; use cauchy::*; use num_traits::{ToPrimitive, Zero}; -/// Result of SVD -pub struct SVDOutput { - /// diagonal values - pub s: Vec, - /// Unitary matrix for destination space - pub u: Option>, - /// Unitary matrix for departure space - pub vt: Option>, +pub struct SvdWork { + pub ju: JobSvd, + pub jvt: JobSvd, + pub layout: MatrixLayout, + pub s: Vec>, + pub u: Option>>, + pub vt: Option>>, + pub work: Vec>, + pub rwork: Option>>, } -#[cfg_attr(doc, katexit::katexit)] -/// Singular value decomposition -pub trait SVD_: Scalar { - /// Compute singular value decomposition $A = U \Sigma V^T$ - /// - /// LAPACK correspondance - /// ---------------------- - /// - /// | f32 | f64 | c32 | c64 | - /// |:-------|:-------|:-------|:-------| - /// | sgesvd | dgesvd | cgesvd | zgesvd | - /// - fn svd(l: MatrixLayout, calc_u: bool, calc_vt: bool, a: &mut [Self]) - -> Result>; +#[derive(Debug, Clone)] +pub struct SvdRef<'work, T: Scalar> { + pub s: &'work [T::Real], + pub u: Option<&'work [T]>, + pub vt: Option<&'work [T]>, } -macro_rules! impl_svd { - (@real, $scalar:ty, $gesvd:path) => { - impl_svd!(@body, $scalar, $gesvd, ); - }; - (@complex, $scalar:ty, $gesvd:path) => { - impl_svd!(@body, $scalar, $gesvd, rwork); - }; - (@body, $scalar:ty, $gesvd:path, $($rwork_ident:ident),*) => { - impl SVD_ for $scalar { - fn svd(l: MatrixLayout, calc_u: bool, calc_vt: bool, a: &mut [Self],) -> Result> { - let ju = match l { +#[derive(Debug, Clone)] +pub struct SvdOwned { + pub s: Vec, + pub u: Option>, + pub vt: Option>, +} + +pub trait SvdWorkImpl: Sized { + type Elem: Scalar; + fn new(layout: MatrixLayout, calc_u: bool, calc_vt: bool) -> Result; + fn calc(&mut self, a: &mut [Self::Elem]) -> Result>; + fn eval(self, a: &mut [Self::Elem]) -> Result>; +} + +macro_rules! impl_svd_work_c { + ($s:ty, $svd:path) => { + impl SvdWorkImpl for SvdWork<$s> { + type Elem = $s; + + fn new(layout: MatrixLayout, calc_u: bool, calc_vt: bool) -> Result { + let ju = match layout { MatrixLayout::F { .. } => JobSvd::from_bool(calc_u), MatrixLayout::C { .. } => JobSvd::from_bool(calc_vt), }; - let jvt = match l { + let jvt = match layout { MatrixLayout::F { .. } => JobSvd::from_bool(calc_vt), MatrixLayout::C { .. } => JobSvd::from_bool(calc_u), }; - let m = l.lda(); + let m = layout.lda(); let mut u = match ju { - JobSvd::All => Some(vec_uninit( (m * m) as usize)), + JobSvd::All => Some(vec_uninit((m * m) as usize)), JobSvd::None => None, - _ => unimplemented!("SVD with partial vector output is not supported yet") + _ => unimplemented!("SVD with partial vector output is not supported yet"), }; - let n = l.len(); + let n = layout.len(); let mut vt = match jvt { - JobSvd::All => Some(vec_uninit( (n * n) as usize)), + JobSvd::All => Some(vec_uninit((n * n) as usize)), JobSvd::None => None, - _ => unimplemented!("SVD with partial vector output is not supported yet") + _ => unimplemented!("SVD with partial vector output is not supported yet"), }; let k = std::cmp::min(m, n); - let mut s = vec_uninit( k as usize); - - $( - let mut $rwork_ident: Vec> = vec_uninit(5 * k as usize); - )* + let mut s = vec_uninit(k as usize); + let mut rwork = vec_uninit(5 * k as usize); // eval work size let mut info = 0; - let mut work_size = [Self::zero()]; + let mut work_size = [Self::Elem::zero()]; unsafe { - $gesvd( + $svd( ju.as_ptr(), jvt.as_ptr(), &m, &n, - AsPtr::as_mut_ptr(a), + std::ptr::null_mut(), &m, AsPtr::as_mut_ptr(&mut s), AsPtr::as_mut_ptr(u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), @@ -88,50 +95,220 @@ macro_rules! impl_svd { &n, AsPtr::as_mut_ptr(&mut work_size), &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + AsPtr::as_mut_ptr(&mut rwork), &mut info, ); } info.as_lapack_result()?; - - // calc let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); + let work = vec_uninit(lwork); + Ok(SvdWork { + layout, + ju, + jvt, + s, + u, + vt, + work, + rwork: Some(rwork), + }) + } + + fn calc(&mut self, a: &mut [Self::Elem]) -> Result> { + let m = self.layout.lda(); + let n = self.layout.len(); + let lwork = self.work.len().to_i32().unwrap(); + + let mut info = 0; unsafe { - $gesvd( - ju.as_ptr(), - jvt.as_ptr() , + $svd( + self.ju.as_ptr(), + self.jvt.as_ptr(), &m, &n, AsPtr::as_mut_ptr(a), &m, + AsPtr::as_mut_ptr(&mut self.s), + AsPtr::as_mut_ptr( + self.u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut []), + ), + &m, + AsPtr::as_mut_ptr( + self.vt + .as_mut() + .map(|x| x.as_mut_slice()) + .unwrap_or(&mut []), + ), + &n, + AsPtr::as_mut_ptr(&mut self.work), + &(lwork as i32), + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + &mut info, + ); + } + info.as_lapack_result()?; + + let s = unsafe { self.s.slice_assume_init_ref() }; + let u = self + .u + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + let vt = self + .vt + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + + match self.layout { + MatrixLayout::F { .. } => Ok(SvdRef { s, u, vt }), + MatrixLayout::C { .. } => Ok(SvdRef { s, u: vt, vt: u }), + } + } + + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _ref = self.calc(a)?; + let s = unsafe { self.s.assume_init() }; + let u = self.u.map(|v| unsafe { v.assume_init() }); + let vt = self.vt.map(|v| unsafe { v.assume_init() }); + match self.layout { + MatrixLayout::F { .. } => Ok(SvdOwned { s, u, vt }), + MatrixLayout::C { .. } => Ok(SvdOwned { s, u: vt, vt: u }), + } + } + } + }; +} +impl_svd_work_c!(c64, lapack_sys::zgesvd_); +impl_svd_work_c!(c32, lapack_sys::cgesvd_); + +macro_rules! impl_svd_work_r { + ($s:ty, $svd:path) => { + impl SvdWorkImpl for SvdWork<$s> { + type Elem = $s; + + fn new(layout: MatrixLayout, calc_u: bool, calc_vt: bool) -> Result { + let ju = match layout { + MatrixLayout::F { .. } => JobSvd::from_bool(calc_u), + MatrixLayout::C { .. } => JobSvd::from_bool(calc_vt), + }; + let jvt = match layout { + MatrixLayout::F { .. } => JobSvd::from_bool(calc_vt), + MatrixLayout::C { .. } => JobSvd::from_bool(calc_u), + }; + + let m = layout.lda(); + let mut u = match ju { + JobSvd::All => Some(vec_uninit((m * m) as usize)), + JobSvd::None => None, + _ => unimplemented!("SVD with partial vector output is not supported yet"), + }; + + let n = layout.len(); + let mut vt = match jvt { + JobSvd::All => Some(vec_uninit((n * n) as usize)), + JobSvd::None => None, + _ => unimplemented!("SVD with partial vector output is not supported yet"), + }; + + let k = std::cmp::min(m, n); + let mut s = vec_uninit(k as usize); + + // eval work size + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $svd( + ju.as_ptr(), + jvt.as_ptr(), + &m, + &n, + std::ptr::null_mut(), + &m, AsPtr::as_mut_ptr(&mut s), AsPtr::as_mut_ptr(u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), &m, AsPtr::as_mut_ptr(vt.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), &n, - AsPtr::as_mut_ptr(&mut work), + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + &mut info, + ); + } + info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(SvdWork { + layout, + ju, + jvt, + s, + u, + vt, + work, + rwork: None, + }) + } + + fn calc(&mut self, a: &mut [Self::Elem]) -> Result> { + let m = self.layout.lda(); + let n = self.layout.len(); + let lwork = self.work.len().to_i32().unwrap(); + + let mut info = 0; + unsafe { + $svd( + self.ju.as_ptr(), + self.jvt.as_ptr(), + &m, + &n, + AsPtr::as_mut_ptr(a), + &m, + AsPtr::as_mut_ptr(&mut self.s), + AsPtr::as_mut_ptr( + self.u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut []), + ), + &m, + AsPtr::as_mut_ptr( + self.vt + .as_mut() + .map(|x| x.as_mut_slice()) + .unwrap_or(&mut []), + ), + &n, + AsPtr::as_mut_ptr(&mut self.work), &(lwork as i32), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* &mut info, ); } info.as_lapack_result()?; - let s = unsafe { s.assume_init() }; - let u = u.map(|v| unsafe { v.assume_init() }); - let vt = vt.map(|v| unsafe { v.assume_init() }); + let s = unsafe { self.s.slice_assume_init_ref() }; + let u = self + .u + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + let vt = self + .vt + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); - match l { - MatrixLayout::F { .. } => Ok(SVDOutput { s, u, vt }), - MatrixLayout::C { .. } => Ok(SVDOutput { s, u: vt, vt: u }), + match self.layout { + MatrixLayout::F { .. } => Ok(SvdRef { s, u, vt }), + MatrixLayout::C { .. } => Ok(SvdRef { s, u: vt, vt: u }), + } + } + + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _ref = self.calc(a)?; + let s = unsafe { self.s.assume_init() }; + let u = self.u.map(|v| unsafe { v.assume_init() }); + let vt = self.vt.map(|v| unsafe { v.assume_init() }); + match self.layout { + MatrixLayout::F { .. } => Ok(SvdOwned { s, u, vt }), + MatrixLayout::C { .. } => Ok(SvdOwned { s, u: vt, vt: u }), } } } }; -} // impl_svd! - -impl_svd!(@real, f64, lapack_sys::dgesvd_); -impl_svd!(@real, f32, lapack_sys::sgesvd_); -impl_svd!(@complex, c64, lapack_sys::zgesvd_); -impl_svd!(@complex, c32, lapack_sys::cgesvd_); +} +impl_svd_work_r!(f64, lapack_sys::dgesvd_); +impl_svd_work_r!(f32, lapack_sys::sgesvd_); diff --git a/lax/src/svddc.rs b/lax/src/svddc.rs index 928fd44c..c16db4bb 100644 --- a/lax/src/svddc.rs +++ b/lax/src/svddc.rs @@ -1,41 +1,50 @@ +//! Compute singular value decomposition with divide-and-conquer algorithm +//! +//! LAPACK correspondance +//! ---------------------- +//! +//! | f32 | f64 | c32 | c64 | +//! |:-------|:-------|:-------|:-------| +//! | sgesdd | dgesdd | cgesdd | zgesdd | +//! + use crate::{error::*, layout::MatrixLayout, *}; use cauchy::*; use num_traits::{ToPrimitive, Zero}; -#[cfg_attr(doc, katexit::katexit)] -/// Singular value decomposition with divide-and-conquer method -pub trait SVDDC_: Scalar { - /// Compute singular value decomposition $A = U \Sigma V^T$ - /// - /// LAPACK correspondance - /// ---------------------- - /// - /// | f32 | f64 | c32 | c64 | - /// |:-------|:-------|:-------|:-------| - /// | sgesdd | dgesdd | cgesdd | zgesdd | - /// - fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; +pub struct SvdDcWork { + pub jobz: JobSvd, + pub layout: MatrixLayout, + pub s: Vec>, + pub u: Option>>, + pub vt: Option>>, + pub work: Vec>, + pub iwork: Vec>, + pub rwork: Option>>, } -macro_rules! impl_svddc { - (@real, $scalar:ty, $gesdd:path) => { - impl_svddc!(@body, $scalar, $gesdd, ); - }; - (@complex, $scalar:ty, $gesdd:path) => { - impl_svddc!(@body, $scalar, $gesdd, rwork); - }; - (@body, $scalar:ty, $gesdd:path, $($rwork_ident:ident),*) => { - impl SVDDC_ for $scalar { - fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self],) -> Result> { - let m = l.lda(); - let n = l.len(); - let k = m.min(n); - let mut s = vec_uninit(k as usize); +pub trait SvdDcWorkImpl: Sized { + type Elem: Scalar; + fn new(layout: MatrixLayout, jobz: JobSvd) -> Result; + fn calc(&mut self, a: &mut [Self::Elem]) -> Result>; + fn eval(self, a: &mut [Self::Elem]) -> Result>; +} +macro_rules! impl_svd_dc_work_c { + ($s:ty, $sdd:path) => { + impl SvdDcWorkImpl for SvdDcWork<$s> { + type Elem = $s; + + fn new(layout: MatrixLayout, jobz: JobSvd) -> Result { + let m = layout.lda(); + let n = layout.len(); + let k = m.min(n); let (u_col, vt_row) = match jobz { JobSvd::All | JobSvd::None => (m, n), JobSvd::Some => (k, k), }; + + let mut s = vec_uninit(k as usize); let (mut u, mut vt) = match jobz { JobSvd::All => ( Some(vec_uninit((m * m) as usize)), @@ -47,27 +56,24 @@ macro_rules! impl_svddc { ), JobSvd::None => (None, None), }; + let mut iwork = vec_uninit(8 * k as usize); - $( // for complex only let mx = n.max(m) as usize; let mn = n.min(m) as usize; let lrwork = match jobz { JobSvd::None => 7 * mn, - _ => std::cmp::max(5*mn*mn + 5*mn, 2*mx*mn + 2*mn*mn + mn), + _ => std::cmp::max(5 * mn * mn + 5 * mn, 2 * mx * mn + 2 * mn * mn + mn), }; - let mut $rwork_ident: Vec> = vec_uninit(lrwork); - )* + let mut rwork = vec_uninit(lrwork); - // eval work size let mut info = 0; - let mut iwork: Vec> = vec_uninit(8 * k as usize); - let mut work_size = [Self::zero()]; + let mut work_size = [Self::Elem::zero()]; unsafe { - $gesdd( + $sdd( jobz.as_ptr(), &m, &n, - AsPtr::as_mut_ptr(a), + std::ptr::null_mut(), &m, AsPtr::as_mut_ptr(&mut s), AsPtr::as_mut_ptr(u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), @@ -76,51 +82,226 @@ macro_rules! impl_svddc { &vt_row, AsPtr::as_mut_ptr(&mut work_size), &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + AsPtr::as_mut_ptr(&mut rwork), AsPtr::as_mut_ptr(&mut iwork), &mut info, ); } info.as_lapack_result()?; - - // do svd let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); + let work = vec_uninit(lwork); + Ok(SvdDcWork { + layout, + jobz, + iwork, + work, + rwork: Some(rwork), + u, + vt, + s, + }) + } + + fn calc(&mut self, a: &mut [Self::Elem]) -> Result> { + let m = self.layout.lda(); + let n = self.layout.len(); + let k = m.min(n); + let (_, vt_row) = match self.jobz { + JobSvd::All | JobSvd::None => (m, n), + JobSvd::Some => (k, k), + }; + let lwork = self.work.len().to_i32().unwrap(); + + let mut info = 0; unsafe { - $gesdd( - jobz.as_ptr(), + $sdd( + self.jobz.as_ptr(), &m, &n, AsPtr::as_mut_ptr(a), &m, + AsPtr::as_mut_ptr(&mut self.s), + AsPtr::as_mut_ptr( + self.u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut []), + ), + &m, + AsPtr::as_mut_ptr( + self.vt + .as_mut() + .map(|x| x.as_mut_slice()) + .unwrap_or(&mut []), + ), + &vt_row, + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()), + AsPtr::as_mut_ptr(&mut self.iwork), + &mut info, + ); + } + info.as_lapack_result()?; + + let s = unsafe { self.s.slice_assume_init_ref() }; + let u = self + .u + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + let vt = self + .vt + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + + Ok(match self.layout { + MatrixLayout::F { .. } => SvdRef { s, u, vt }, + MatrixLayout::C { .. } => SvdRef { s, u: vt, vt: u }, + }) + } + + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _ref = self.calc(a)?; + let s = unsafe { self.s.assume_init() }; + let u = self.u.map(|v| unsafe { v.assume_init() }); + let vt = self.vt.map(|v| unsafe { v.assume_init() }); + Ok(match self.layout { + MatrixLayout::F { .. } => SvdOwned { s, u, vt }, + MatrixLayout::C { .. } => SvdOwned { s, u: vt, vt: u }, + }) + } + } + }; +} +impl_svd_dc_work_c!(c64, lapack_sys::zgesdd_); +impl_svd_dc_work_c!(c32, lapack_sys::cgesdd_); + +macro_rules! impl_svd_dc_work_r { + ($s:ty, $sdd:path) => { + impl SvdDcWorkImpl for SvdDcWork<$s> { + type Elem = $s; + + fn new(layout: MatrixLayout, jobz: JobSvd) -> Result { + let m = layout.lda(); + let n = layout.len(); + let k = m.min(n); + let (u_col, vt_row) = match jobz { + JobSvd::All | JobSvd::None => (m, n), + JobSvd::Some => (k, k), + }; + + let mut s = vec_uninit(k as usize); + let (mut u, mut vt) = match jobz { + JobSvd::All => ( + Some(vec_uninit((m * m) as usize)), + Some(vec_uninit((n * n) as usize)), + ), + JobSvd::Some => ( + Some(vec_uninit((m * u_col) as usize)), + Some(vec_uninit((n * vt_row) as usize)), + ), + JobSvd::None => (None, None), + }; + let mut iwork = vec_uninit(8 * k as usize); + + let mut info = 0; + let mut work_size = [Self::Elem::zero()]; + unsafe { + $sdd( + jobz.as_ptr(), + &m, + &n, + std::ptr::null_mut(), + &m, AsPtr::as_mut_ptr(&mut s), AsPtr::as_mut_ptr(u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), &m, AsPtr::as_mut_ptr(vt.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut [])), &vt_row, - AsPtr::as_mut_ptr(&mut work), - &(lwork as i32), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* + AsPtr::as_mut_ptr(&mut work_size), + &(-1), AsPtr::as_mut_ptr(&mut iwork), &mut info, ); } info.as_lapack_result()?; + let lwork = work_size[0].to_usize().unwrap(); + let work = vec_uninit(lwork); + Ok(SvdDcWork { + layout, + jobz, + iwork, + work, + rwork: None, + u, + vt, + s, + }) + } - let s = unsafe { s.assume_init() }; - let u = u.map(|v| unsafe { v.assume_init() }); - let vt = vt.map(|v| unsafe { v.assume_init() }); + fn calc(&mut self, a: &mut [Self::Elem]) -> Result> { + let m = self.layout.lda(); + let n = self.layout.len(); + let k = m.min(n); + let (_, vt_row) = match self.jobz { + JobSvd::All | JobSvd::None => (m, n), + JobSvd::Some => (k, k), + }; + let lwork = self.work.len().to_i32().unwrap(); - match l { - MatrixLayout::F { .. } => Ok(SVDOutput { s, u, vt }), - MatrixLayout::C { .. } => Ok(SVDOutput { s, u: vt, vt: u }), + let mut info = 0; + unsafe { + $sdd( + self.jobz.as_ptr(), + &m, + &n, + AsPtr::as_mut_ptr(a), + &m, + AsPtr::as_mut_ptr(&mut self.s), + AsPtr::as_mut_ptr( + self.u.as_mut().map(|x| x.as_mut_slice()).unwrap_or(&mut []), + ), + &m, + AsPtr::as_mut_ptr( + self.vt + .as_mut() + .map(|x| x.as_mut_slice()) + .unwrap_or(&mut []), + ), + &vt_row, + AsPtr::as_mut_ptr(&mut self.work), + &lwork, + AsPtr::as_mut_ptr(&mut self.iwork), + &mut info, + ); } + info.as_lapack_result()?; + + let s = unsafe { self.s.slice_assume_init_ref() }; + let u = self + .u + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + let vt = self + .vt + .as_ref() + .map(|v| unsafe { v.slice_assume_init_ref() }); + + Ok(match self.layout { + MatrixLayout::F { .. } => SvdRef { s, u, vt }, + MatrixLayout::C { .. } => SvdRef { s, u: vt, vt: u }, + }) + } + + fn eval(mut self, a: &mut [Self::Elem]) -> Result> { + let _ref = self.calc(a)?; + let s = unsafe { self.s.assume_init() }; + let u = self.u.map(|v| unsafe { v.assume_init() }); + let vt = self.vt.map(|v| unsafe { v.assume_init() }); + Ok(match self.layout { + MatrixLayout::F { .. } => SvdOwned { s, u, vt }, + MatrixLayout::C { .. } => SvdOwned { s, u: vt, vt: u }, + }) } } }; } - -impl_svddc!(@real, f32, lapack_sys::sgesdd_); -impl_svddc!(@real, f64, lapack_sys::dgesdd_); -impl_svddc!(@complex, c32, lapack_sys::cgesdd_); -impl_svddc!(@complex, c64, lapack_sys::zgesdd_); +impl_svd_dc_work_r!(f64, lapack_sys::dgesdd_); +impl_svd_dc_work_r!(f32, lapack_sys::sgesdd_); diff --git a/ndarray-linalg/src/least_squares.rs b/ndarray-linalg/src/least_squares.rs index 86ca17f3..f376f569 100644 --- a/ndarray-linalg/src/least_squares.rs +++ b/ndarray-linalg/src/least_squares.rs @@ -297,7 +297,7 @@ where D1: DataMut, D2: DataMut, { - let LeastSquaresOutput:: { + let LeastSquaresOwned:: { singular_values, rank, } = E::least_squares( @@ -340,7 +340,7 @@ fn compute_residual_scalar>( /// valid representation for `ArrayBase` (over `E`). impl LeastSquaresSvdInPlace for ArrayBase where - E: Scalar + Lapack + LeastSquaresSvdDivideConquer_, + E: Scalar + Lapack, D1: DataMut, D2: DataMut, { @@ -386,7 +386,7 @@ where { let a_layout = a.layout()?; let rhs_layout = rhs.layout()?; - let LeastSquaresOutput:: { + let LeastSquaresOwned:: { singular_values, rank, } = E::least_squares_nrhs(