From f9f16e23978d3570248b6ab5a6859a25912b4b6f Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Tue, 27 Sep 2022 16:50:49 +0900 Subject: [PATCH 1/9] Add SvdWork --- lax/src/svd.rs | 301 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 301 insertions(+) diff --git a/lax/src/svd.rs b/lax/src/svd.rs index 0d9bd353..0c44241b 100644 --- a/lax/src/svd.rs +++ b/lax/src/svd.rs @@ -30,6 +30,307 @@ pub trait SVD_: Scalar { -> Result>; } +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>>, +} + +#[derive(Debug, Clone)] +pub struct SvdRef<'work, T: Scalar> { + pub s: &'work [T::Real], + pub u: Option<&'work [T]>, + pub vt: Option<&'work [T]>, +} + +#[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 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); + let mut rwork = vec_uninit(5 * 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_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + &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: 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 { + $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_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), + &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_r!(f64, lapack_sys::dgesvd_); +impl_svd_work_r!(f32, lapack_sys::sgesvd_); + macro_rules! impl_svd { (@real, $scalar:ty, $gesvd:path) => { impl_svd!(@body, $scalar, $gesvd, ); From 3fd8c0b26bc98e5d12bd7d5b9fb59c57c60fb634 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Tue, 27 Sep 2022 16:57:43 +0900 Subject: [PATCH 2/9] Merge SVD_ to Lapack --- lax/src/lib.rs | 21 +++++-- lax/src/svd.rs | 140 +++-------------------------------------------- lax/src/svddc.rs | 8 +-- 3 files changed, 29 insertions(+), 140 deletions(-) diff --git a/lax/src/lib.rs b/lax/src/lib.rs index ffe1fb62..582a980a 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -65,7 +65,7 @@ //! Singular Value Decomposition //! ----------------------------- //! -//! - [SVD_] trait provides methods for singular value decomposition for general matrix +//! - [svd] module for singular value decomposition (SVD) for general matrix //! - [SVDDC_] trait provides methods for singular value decomposition for general matrix //! with divided-and-conquer algorithm //! - [LeastSquaresSvdDivideConquer_] trait provides methods @@ -91,6 +91,7 @@ pub mod eig; pub mod eigh; pub mod eigh_generalized; pub mod qr; +pub mod svd; mod alloc; mod cholesky; @@ -99,7 +100,6 @@ mod opnorm; mod rcond; mod solve; mod solveh; -mod svd; mod svddc; mod triangular; mod tridiagonal; @@ -111,7 +111,7 @@ pub use self::opnorm::*; pub use self::rcond::*; pub use self::solve::*; pub use self::solveh::*; -pub use self::svd::*; +pub use self::svd::SvdOwned; pub use self::svddc::*; pub use self::triangular::*; pub use self::tridiagonal::*; @@ -125,7 +125,6 @@ pub type Pivot = Vec; /// Trait for primitive types which implements LAPACK subroutines pub trait Lapack: OperatorNorm_ - + SVD_ + SVDDC_ + Solve_ + Solveh_ @@ -170,6 +169,9 @@ 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>; } macro_rules! impl_lapack { @@ -228,6 +230,17 @@ 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) + } } }; } diff --git a/lax/src/svd.rs b/lax/src/svd.rs index 0c44241b..fc695108 100644 --- a/lax/src/svd.rs +++ b/lax/src/svd.rs @@ -1,35 +1,17 @@ //! 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>, -} - -#[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>; -} - pub struct SvdWork { pub ju: JobSvd, pub jvt: JobSvd, @@ -330,109 +312,3 @@ macro_rules! impl_svd_work_r { } impl_svd_work_r!(f64, lapack_sys::dgesvd_); impl_svd_work_r!(f32, lapack_sys::sgesvd_); - -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 { - MatrixLayout::F { .. } => JobSvd::from_bool(calc_u), - MatrixLayout::C { .. } => JobSvd::from_bool(calc_vt), - }; - let jvt = match l { - MatrixLayout::F { .. } => JobSvd::from_bool(calc_vt), - MatrixLayout::C { .. } => JobSvd::from_bool(calc_u), - }; - - let m = l.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 = l.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); - - $( - let mut $rwork_ident: Vec> = vec_uninit(5 * k as usize); - )* - - // eval work size - let mut info = 0; - let mut work_size = [Self::zero()]; - unsafe { - $gesvd( - ju.as_ptr(), - jvt.as_ptr(), - &m, - &n, - AsPtr::as_mut_ptr(a), - &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_size), - &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - &mut info, - ); - } - info.as_lapack_result()?; - - // calc - let lwork = work_size[0].to_usize().unwrap(); - let mut work: Vec> = vec_uninit(lwork); - unsafe { - $gesvd( - ju.as_ptr(), - jvt.as_ptr() , - &m, - &n, - AsPtr::as_mut_ptr(a), - &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), - &(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() }); - - match l { - MatrixLayout::F { .. } => Ok(SVDOutput { s, u, vt }), - MatrixLayout::C { .. } => Ok(SVDOutput { 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_); diff --git a/lax/src/svddc.rs b/lax/src/svddc.rs index 928fd44c..22612cee 100644 --- a/lax/src/svddc.rs +++ b/lax/src/svddc.rs @@ -14,7 +14,7 @@ pub trait SVDDC_: Scalar { /// |:-------|:-------|:-------|:-------| /// | sgesdd | dgesdd | cgesdd | zgesdd | /// - fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; + fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; } macro_rules! impl_svddc { @@ -26,7 +26,7 @@ macro_rules! impl_svddc { }; (@body, $scalar:ty, $gesdd:path, $($rwork_ident:ident),*) => { impl SVDDC_ for $scalar { - fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self],) -> Result> { + fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self],) -> Result> { let m = l.lda(); let n = l.len(); let k = m.min(n); @@ -112,8 +112,8 @@ macro_rules! impl_svddc { let vt = vt.map(|v| unsafe { v.assume_init() }); match l { - MatrixLayout::F { .. } => Ok(SVDOutput { s, u, vt }), - MatrixLayout::C { .. } => Ok(SVDOutput { s, u: vt, vt: u }), + MatrixLayout::F { .. } => Ok(SvdOwned { s, u, vt }), + MatrixLayout::C { .. } => Ok(SvdOwned { s, u: vt, vt: u }), } } } From f30931d61c241774ac069923a4b341fd0714f3e9 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Tue, 27 Sep 2022 21:53:02 +0900 Subject: [PATCH 3/9] SvdDcWork and SvdDcWorkImpl --- lax/src/lib.rs | 2 +- lax/src/svddc.rs | 296 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 296 insertions(+), 2 deletions(-) diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 582a980a..9adc7fdb 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -111,7 +111,7 @@ pub use self::opnorm::*; pub use self::rcond::*; pub use self::solve::*; pub use self::solveh::*; -pub use self::svd::SvdOwned; +pub use self::svd::{SvdOwned, SvdRef}; pub use self::svddc::*; pub use self::triangular::*; pub use self::tridiagonal::*; diff --git a/lax/src/svddc.rs b/lax/src/svddc.rs index 22612cee..239e0d8c 100644 --- a/lax/src/svddc.rs +++ b/lax/src/svddc.rs @@ -14,9 +14,303 @@ pub trait SVDDC_: Scalar { /// |:-------|:-------|:-------|:-------| /// | sgesdd | dgesdd | cgesdd | zgesdd | /// - fn svddc(l: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; + fn svddc(layout: 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>>, +} + +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)), + 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 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), + }; + let mut rwork = vec_uninit(lrwork); + + 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_size), + &(-1), + AsPtr::as_mut_ptr(&mut rwork), + 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: 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 { + $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_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, + }) + } + + 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 { + $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_svd_dc_work_r!(f64, lapack_sys::dgesdd_); +impl_svd_dc_work_r!(f32, lapack_sys::sgesdd_); + macro_rules! impl_svddc { (@real, $scalar:ty, $gesdd:path) => { impl_svddc!(@body, $scalar, $gesdd, ); From 38c64c9ea4c2be2a93b166f6c04dd69434ec7ab8 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Tue, 27 Sep 2022 22:20:11 +0900 Subject: [PATCH 4/9] Merge SVDDC_ to Lapack trait --- lax/src/lib.rs | 16 ++++-- lax/src/svddc.rs | 133 ++++------------------------------------------- 2 files changed, 21 insertions(+), 128 deletions(-) diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 9adc7fdb..00ad27b3 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -66,8 +66,7 @@ //! ----------------------------- //! //! - [svd] module for singular value decomposition (SVD) for general matrix -//! - [SVDDC_] trait provides methods for singular value decomposition for general matrix -//! with divided-and-conquer algorithm +//! - [svddc] module for singular value decomposition (SVD) with divided-and-conquer algorithm for general matrix //! - [LeastSquaresSvdDivideConquer_] trait provides methods //! for solving least square problem by SVD //! @@ -92,6 +91,7 @@ pub mod eigh; pub mod eigh_generalized; pub mod qr; pub mod svd; +pub mod svddc; mod alloc; mod cholesky; @@ -100,7 +100,6 @@ mod opnorm; mod rcond; mod solve; mod solveh; -mod svddc; mod triangular; mod tridiagonal; @@ -112,7 +111,6 @@ pub use self::rcond::*; pub use self::solve::*; pub use self::solveh::*; pub use self::svd::{SvdOwned, SvdRef}; -pub use self::svddc::*; pub use self::triangular::*; pub use self::tridiagonal::*; @@ -125,7 +123,6 @@ pub type Pivot = Vec; /// Trait for primitive types which implements LAPACK subroutines pub trait Lapack: OperatorNorm_ - + SVDDC_ + Solve_ + Solveh_ + Cholesky_ @@ -172,6 +169,9 @@ pub trait Lapack: /// 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>; } macro_rules! impl_lapack { @@ -241,6 +241,12 @@ macro_rules! impl_lapack { 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) + } } }; } diff --git a/lax/src/svddc.rs b/lax/src/svddc.rs index 239e0d8c..c16db4bb 100644 --- a/lax/src/svddc.rs +++ b/lax/src/svddc.rs @@ -1,22 +1,17 @@ +//! 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(layout: MatrixLayout, jobz: JobSvd, a: &mut [Self]) -> Result>; -} - pub struct SvdDcWork { pub jobz: JobSvd, pub layout: MatrixLayout, @@ -310,111 +305,3 @@ macro_rules! impl_svd_dc_work_r { } impl_svd_dc_work_r!(f64, lapack_sys::dgesdd_); impl_svd_dc_work_r!(f32, lapack_sys::sgesdd_); - -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); - - let (u_col, vt_row) = match jobz { - JobSvd::All | JobSvd::None => (m, n), - JobSvd::Some => (k, k), - }; - 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), - }; - - $( // 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), - }; - let mut $rwork_ident: Vec> = 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()]; - unsafe { - $gesdd( - jobz.as_ptr(), - &m, - &n, - AsPtr::as_mut_ptr(a), - &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_size), - &(-1), - $(AsPtr::as_mut_ptr(&mut $rwork_ident),)* - 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); - unsafe { - $gesdd( - jobz.as_ptr(), - &m, - &n, - AsPtr::as_mut_ptr(a), - &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 iwork), - &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() }); - - match l { - MatrixLayout::F { .. } => Ok(SvdOwned { s, u, vt }), - MatrixLayout::C { .. } => Ok(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_); From ac2f7bcb18e817168367dd6ede85f8413e0a5fcf Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Wed, 28 Sep 2022 00:19:16 +0900 Subject: [PATCH 5/9] Rename LeastSquaresOutput to LeastSquaresOwned --- lax/src/least_squares.rs | 10 +++++----- ndarray-linalg/src/least_squares.rs | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/lax/src/least_squares.rs b/lax/src/least_squares.rs index 1bfd4d37..532d1e6b 100644 --- a/lax/src/least_squares.rs +++ b/lax/src/least_squares.rs @@ -5,7 +5,7 @@ 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 @@ -21,7 +21,7 @@ pub trait LeastSquaresSvdDivideConquer_: Scalar { a_layout: MatrixLayout, a: &mut [Self], b: &mut [Self], - ) -> Result>; + ) -> Result>; /// Solve least square problems $\argmin_X \| AX - B\|$ fn least_squares_nrhs( @@ -46,7 +46,7 @@ macro_rules! impl_least_squares { l: MatrixLayout, a: &mut [Self], b: &mut [Self], - ) -> Result> { + ) -> Result> { let b_layout = l.resized(b.len() as i32, 1); Self::least_squares_nrhs(l, a, b_layout, b) } @@ -56,7 +56,7 @@ macro_rules! impl_least_squares { a: &mut [Self], b_layout: MatrixLayout, b: &mut [Self], - ) -> Result> { + ) -> Result> { // Minimize |b - Ax|_2 // // where @@ -160,7 +160,7 @@ macro_rules! impl_least_squares { transpose_over(b_layout, &b_t, b); } - Ok(LeastSquaresOutput { + Ok(LeastSquaresOwned { singular_values, rank, }) diff --git a/ndarray-linalg/src/least_squares.rs b/ndarray-linalg/src/least_squares.rs index 86ca17f3..0042ef74 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( @@ -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( From da3221a74ea544b10afddca84ef75099f6670c5a Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Wed, 28 Sep 2022 00:19:39 +0900 Subject: [PATCH 6/9] LeastSquaresWork --- lax/src/least_squares.rs | 327 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 326 insertions(+), 1 deletion(-) diff --git a/lax/src/least_squares.rs b/lax/src/least_squares.rs index 532d1e6b..44445316 100644 --- a/lax/src/least_squares.rs +++ b/lax/src/least_squares.rs @@ -12,6 +12,14 @@ pub struct LeastSquaresOwned { pub rank: i32, } +/// 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, +} + #[cfg_attr(doc, katexit::katexit)] /// Solve least square problem pub trait LeastSquaresSvdDivideConquer_: Scalar { @@ -29,8 +37,325 @@ pub trait LeastSquaresSvdDivideConquer_: Scalar { a: &mut [Self], b_layout: MatrixLayout, b: &mut [Self], - ) -> Result>; + ) -> Result>; +} + +pub struct LeastSquaresWork { + pub a_layout: MatrixLayout, + pub b_layout: MatrixLayout, + pub singular_values: Vec>, + pub work: Vec>, + pub iwork: Vec>, + pub rwork: Option>>, +} + +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 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(), + &a_layout.lda(), + std::ptr::null_mut(), + &b_layout.lda(), + 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 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 { + $lsd( + &m, + &n, + &nrhs, + AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), + &a_layout.lda(), + 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 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(), + &a_layout.lda(), + std::ptr::null_mut(), + &b_layout.lda(), + AsPtr::as_mut_ptr(&mut singular_values), + &rcond, + &mut rank, + AsPtr::as_mut_ptr(&mut work_size), + &(-1), + 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 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 a_layout = 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 { + $lsd( + &m, + &n, + &nrhs, + AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), + &a_layout.lda(), + 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 self.singular_values), + &rcond, + &mut rank, + 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 { 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_r!(f64, lapack_sys::dgelsd_); +impl_least_squares_work_r!(f32, lapack_sys::sgelsd_); macro_rules! impl_least_squares { (@real, $scalar:ty, $gelsd:path) => { From b379b92737412b0aaa6207710c875bff21542c36 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Wed, 28 Sep 2022 21:44:59 +0900 Subject: [PATCH 7/9] Merge LeastSquaresSvdDivideConquer_ into Lapack trait --- lax/src/least_squares.rs | 162 ---------------------------- lax/src/lib.rs | 46 ++++++-- ndarray-linalg/src/least_squares.rs | 2 +- 3 files changed, 39 insertions(+), 171 deletions(-) diff --git a/lax/src/least_squares.rs b/lax/src/least_squares.rs index 44445316..9f8d4561 100644 --- a/lax/src/least_squares.rs +++ b/lax/src/least_squares.rs @@ -20,26 +20,6 @@ pub struct LeastSquaresRef<'work, A: Scalar> { 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>; -} - pub struct LeastSquaresWork { pub a_layout: MatrixLayout, pub b_layout: MatrixLayout, @@ -356,145 +336,3 @@ macro_rules! impl_least_squares_work_r { } impl_least_squares_work_r!(f64, lapack_sys::dgelsd_); impl_least_squares_work_r!(f32, lapack_sys::sgelsd_); - -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); - }; - - (@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) - } - - 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) - let (m, n) = a_layout.size(); - let (m_, nrhs) = b_layout.size(); - let k = m.min(n); - assert!(m_ >= m); - - // Transpose if a is C-continuous - let mut a_t = None; - let a_layout = match a_layout { - MatrixLayout::C { .. } => { - let (layout, t) = transpose(a_layout, a); - a_t = Some(t); - layout - } - MatrixLayout::F { .. } => a_layout, - }; - - // Transpose if b is C-continuous - let mut b_t = None; - let b_layout = match b_layout { - MatrixLayout::C { .. } => { - let (layout, t) = transpose(b_layout, b); - b_t = Some(t); - layout - } - MatrixLayout::F { .. } => b_layout, - }; - - let rcond: Self::Real = -1.; - let mut singular_values: Vec> = vec_uninit( k as usize); - 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( - &m, - &n, - &nrhs, - AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), - &a_layout.lda(), - 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), - &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); - )* - unsafe { - $gelsd( - &m, - &n, - &nrhs, - AsPtr::as_mut_ptr(a_t.as_mut().map(|v| v.as_mut_slice()).unwrap_or(a)), - &a_layout.lda(), - 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), - &rcond, - &mut rank, - AsPtr::as_mut_ptr(&mut work), - &(lwork as i32), - $(AsPtr::as_mut_ptr(&mut $rwork),)* - AsPtr::as_mut_ptr(&mut iwork), - &mut info, - ); - } - info.as_lapack_result()?; - - let singular_values = unsafe { singular_values.assume_init() }; - - // 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(LeastSquaresOwned { - singular_values, - rank, - }) - } - } - }; -} - -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_); diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 00ad27b3..09d15c46 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -120,16 +120,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_ - + 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( @@ -172,6 +166,22 @@ pub trait Lapack: /// 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 { @@ -247,6 +257,26 @@ macro_rules! impl_lapack { 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/ndarray-linalg/src/least_squares.rs b/ndarray-linalg/src/least_squares.rs index 0042ef74..f376f569 100644 --- a/ndarray-linalg/src/least_squares.rs +++ b/ndarray-linalg/src/least_squares.rs @@ -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, { From 6a218d389a7db1a62dac882ae06f23f4d5658457 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Thu, 29 Sep 2022 00:39:52 +0900 Subject: [PATCH 8/9] Fix refactoring bug --- lax/src/least_squares.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/lax/src/least_squares.rs b/lax/src/least_squares.rs index 9f8d4561..d0bb7def 100644 --- a/lax/src/least_squares.rs +++ b/lax/src/least_squares.rs @@ -70,9 +70,9 @@ macro_rules! impl_least_squares_work_c { &n, &nrhs, std::ptr::null_mut(), - &a_layout.lda(), + &m, std::ptr::null_mut(), - &b_layout.lda(), + &m_, AsPtr::as_mut_ptr(&mut singular_values), &rcond, &mut rank, @@ -116,7 +116,7 @@ macro_rules! impl_least_squares_work_c { // Transpose if a is C-continuous let mut a_t = None; - let a_layout = match self.a_layout { + let _ = match self.a_layout { MatrixLayout::C { .. } => { let (layout, t) = transpose(self.a_layout, a); a_t = Some(t); @@ -146,9 +146,9 @@ macro_rules! impl_least_squares_work_c { &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, @@ -218,9 +218,9 @@ macro_rules! impl_least_squares_work_r { &n, &nrhs, std::ptr::null_mut(), - &a_layout.lda(), + &m, std::ptr::null_mut(), - &b_layout.lda(), + &m_, AsPtr::as_mut_ptr(&mut singular_values), &rcond, &mut rank, @@ -261,7 +261,7 @@ macro_rules! impl_least_squares_work_r { // Transpose if a is C-continuous let mut a_t = None; - let a_layout = match self.a_layout { + let _ = match self.a_layout { MatrixLayout::C { .. } => { let (layout, t) = transpose(self.a_layout, a); a_t = Some(t); @@ -291,9 +291,9 @@ macro_rules! impl_least_squares_work_r { &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, From d75c693fd23d7c57a3cf6939006cf48f6b4497c7 Mon Sep 17 00:00:00 2001 From: Toshiki Teramura Date: Thu, 29 Sep 2022 00:48:09 +0900 Subject: [PATCH 9/9] Make least_squares submodule public, fix document --- lax/src/lib.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lax/src/lib.rs b/lax/src/lib.rs index 09d15c46..199f2dc2 100644 --- a/lax/src/lib.rs +++ b/lax/src/lib.rs @@ -67,8 +67,7 @@ //! //! - [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 -//! - [LeastSquaresSvdDivideConquer_] trait provides methods -//! for solving least square problem by SVD +//! - [least_squares] module for solving least square problem using SVD //! #![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)] @@ -89,13 +88,13 @@ 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; @@ -105,7 +104,7 @@ 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::*;