Skip to content

Commit

Permalink
Merge QR_ into Lapack trait
Browse files Browse the repository at this point in the history
  • Loading branch information
@termoshtt
termoshtt committed Sep 26, 2022
1 parent ae57857 commit b861b82
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 207 deletions.
42 changes: 39 additions & 3 deletions lax/src/lib.rs
View file
Expand Up @@ -90,12 +90,12 @@ pub mod layout;
pub mod eig;
pub mod eigh;
pub mod eigh_generalized;
pub mod qr;

mod alloc;
mod cholesky;
mod least_squares;
mod opnorm;
mod qr;
mod rcond;
mod solve;
mod solveh;
Expand All @@ -108,7 +108,6 @@ pub use self::cholesky::*;
pub use self::flags::*;
pub use self::least_squares::*;
pub use self::opnorm::*;
pub use self::qr::*;
pub use self::rcond::*;
pub use self::solve::*;
pub use self::solveh::*;
Expand All @@ -126,7 +125,6 @@ pub type Pivot = Vec<i32>;
/// Trait for primitive types which implements LAPACK subroutines
pub trait Lapack:
OperatorNorm_
+ QR_
+ SVD_
+ SVDDC_
+ Solve_
Expand Down Expand Up @@ -160,6 +158,18 @@ pub trait Lapack:
a: &mut [Self],
b: &mut [Self],
) -> Result<Vec<Self::Real>>;

/// Execute Householder reflection as the first step of QR-decomposition
///
/// For C-continuous array,
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;

/// Reconstruct Q-matrix from Householder-reflectors
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()>;

/// Execute QR-decomposition at once
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
}

macro_rules! impl_lapack {
Expand Down Expand Up @@ -198,6 +208,32 @@ macro_rules! impl_lapack {
let work = EighGeneralizedWork::<$s>::new(calc_eigenvec, layout)?;
work.eval(uplo, a, b)
}

/// Execute Householder reflection as the first step of QR-decomposition
///
/// For C-continuous array,
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
use qr::*;
let work = HouseholderWork::<$s>::new(l)?;
Ok(work.eval(a)?)
}

/// Reconstruct Q-matrix from Householder-reflectors
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()> {
use qr::*;
let mut work = QWork::<$s>::new(l)?;
work.calc(a, tau)?;
Ok(())
}

/// Execute QR-decomposition at once
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
let tau = Self::householder(l, a)?;
let r = Vec::from(&*a);
Self::q(l, a, &tau)?;
Ok(r)
}
}
};
}
Expand Down
206 changes: 2 additions & 204 deletions lax/src/qr.rs
View file
Expand Up @@ -4,20 +4,6 @@ use crate::{error::*, layout::MatrixLayout, *};
use cauchy::*;
use num_traits::{ToPrimitive, Zero};

pub trait QR_: Sized {
/// Execute Householder reflection as the first step of QR-decomposition
///
/// For C-continuous array,
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;

/// Reconstruct Q-matrix from Householder-reflectors
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()>;

/// Execute QR-decomposition at once
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
}

pub struct HouseholderWork<T: Scalar> {
pub m: i32,
pub n: i32,
Expand Down Expand Up @@ -136,7 +122,7 @@ pub struct QWork<T: Scalar> {
pub trait QWorkImpl: Sized {
type Elem: Scalar;
fn new(layout: MatrixLayout) -> Result<Self>;
fn calc(&mut self, a: &mut [Self::Elem], tau: &mut [Self::Elem]) -> Result<()>;
fn calc(&mut self, a: &mut [Self::Elem], tau: &[Self::Elem]) -> Result<()>;
}

macro_rules! impl_q_work {
Expand Down Expand Up @@ -183,7 +169,7 @@ macro_rules! impl_q_work {
Ok(QWork { layout, work })
}

fn calc(&mut self, a: &mut [Self::Elem], tau: &mut [Self::Elem]) -> Result<()> {
fn calc(&mut self, a: &mut [Self::Elem], tau: &[Self::Elem]) -> Result<()> {
let m = self.layout.lda();
let n = self.layout.len();
let k = m.min(n);
Expand Down Expand Up @@ -228,191 +214,3 @@ impl_q_work!(c64, lapack_sys::zungqr_, lapack_sys::zunglq_);
impl_q_work!(c32, lapack_sys::cungqr_, lapack_sys::cunglq_);
impl_q_work!(f64, lapack_sys::dorgqr_, lapack_sys::dorglq_);
impl_q_work!(f32, lapack_sys::sorgqr_, lapack_sys::sorglq_);

macro_rules! impl_qr {
($scalar:ty, $qrf:path, $lqf:path, $gqr:path, $glq:path) => {
impl QR_ for $scalar {
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
let m = l.lda();
let n = l.len();
let k = m.min(n);
let mut tau = vec_uninit(k as usize);

// eval work size
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
match l {
MatrixLayout::F { .. } => {
$qrf(
&m,
&n,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_mut_ptr(&mut tau),
AsPtr::as_mut_ptr(&mut work_size),
&(-1),
&mut info,
);
}
MatrixLayout::C { .. } => {
$lqf(
&m,
&n,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_mut_ptr(&mut tau),
AsPtr::as_mut_ptr(&mut work_size),
&(-1),
&mut info,
);
}
}
}
info.as_lapack_result()?;

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
match l {
MatrixLayout::F { .. } => {
$qrf(
&m,
&n,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_mut_ptr(&mut tau),
AsPtr::as_mut_ptr(&mut work),
&(lwork as i32),
&mut info,
);
}
MatrixLayout::C { .. } => {
$lqf(
&m,
&n,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_mut_ptr(&mut tau),
AsPtr::as_mut_ptr(&mut work),
&(lwork as i32),
&mut info,
);
}
}
}
info.as_lapack_result()?;

let tau = unsafe { tau.assume_init() };

Ok(tau)
}

fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()> {
let m = l.lda();
let n = l.len();
let k = m.min(n);
assert_eq!(tau.len(), k as usize);

// eval work size
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
match l {
MatrixLayout::F { .. } => $gqr(
&m,
&k,
&k,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_ptr(&tau),
AsPtr::as_mut_ptr(&mut work_size),
&(-1),
&mut info,
),
MatrixLayout::C { .. } => $glq(
&k,
&n,
&k,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_ptr(&tau),
AsPtr::as_mut_ptr(&mut work_size),
&(-1),
&mut info,
),
}
};

// calc
let lwork = work_size[0].to_usize().unwrap();
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
unsafe {
match l {
MatrixLayout::F { .. } => $gqr(
&m,
&k,
&k,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_ptr(&tau),
AsPtr::as_mut_ptr(&mut work),
&(lwork as i32),
&mut info,
),
MatrixLayout::C { .. } => $glq(
&k,
&n,
&k,
AsPtr::as_mut_ptr(a),
&m,
AsPtr::as_ptr(&tau),
AsPtr::as_mut_ptr(&mut work),
&(lwork as i32),
&mut info,
),
}
}
info.as_lapack_result()?;
Ok(())
}

fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
let tau = Self::householder(l, a)?;
let r = Vec::from(&*a);
Self::q(l, a, &tau)?;
Ok(r)
}
}
};
} // endmacro

impl_qr!(
f64,
lapack_sys::dgeqrf_,
lapack_sys::dgelqf_,
lapack_sys::dorgqr_,
lapack_sys::dorglq_
);
impl_qr!(
f32,
lapack_sys::sgeqrf_,
lapack_sys::sgelqf_,
lapack_sys::sorgqr_,
lapack_sys::sorglq_
);
impl_qr!(
c64,
lapack_sys::zgeqrf_,
lapack_sys::zgelqf_,
lapack_sys::zungqr_,
lapack_sys::zunglq_
);
impl_qr!(
c32,
lapack_sys::cgeqrf_,
lapack_sys::cgelqf_,
lapack_sys::cungqr_,
lapack_sys::cunglq_
);

0 comments on commit b861b82

Please sign in to comment.