Skip to content

Commit b861b82

Browse files
committed
Merge QR_ into Lapack trait
1 parent ae57857 commit b861b82

File tree

2 files changed

+41
-207
lines changed

2 files changed

+41
-207
lines changed

lax/src/lib.rs

+39-3
Original file line numberDiff line numberDiff line change
@@ -90,12 +90,12 @@ pub mod layout;
9090
pub mod eig;
9191
pub mod eigh;
9292
pub mod eigh_generalized;
93+
pub mod qr;
9394

9495
mod alloc;
9596
mod cholesky;
9697
mod least_squares;
9798
mod opnorm;
98-
mod qr;
9999
mod rcond;
100100
mod solve;
101101
mod solveh;
@@ -108,7 +108,6 @@ pub use self::cholesky::*;
108108
pub use self::flags::*;
109109
pub use self::least_squares::*;
110110
pub use self::opnorm::*;
111-
pub use self::qr::*;
112111
pub use self::rcond::*;
113112
pub use self::solve::*;
114113
pub use self::solveh::*;
@@ -126,7 +125,6 @@ pub type Pivot = Vec<i32>;
126125
/// Trait for primitive types which implements LAPACK subroutines
127126
pub trait Lapack:
128127
OperatorNorm_
129-
+ QR_
130128
+ SVD_
131129
+ SVDDC_
132130
+ Solve_
@@ -160,6 +158,18 @@ pub trait Lapack:
160158
a: &mut [Self],
161159
b: &mut [Self],
162160
) -> Result<Vec<Self::Real>>;
161+
162+
/// Execute Householder reflection as the first step of QR-decomposition
163+
///
164+
/// For C-continuous array,
165+
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
166+
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
167+
168+
/// Reconstruct Q-matrix from Householder-reflectors
169+
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()>;
170+
171+
/// Execute QR-decomposition at once
172+
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
163173
}
164174

165175
macro_rules! impl_lapack {
@@ -198,6 +208,32 @@ macro_rules! impl_lapack {
198208
let work = EighGeneralizedWork::<$s>::new(calc_eigenvec, layout)?;
199209
work.eval(uplo, a, b)
200210
}
211+
212+
/// Execute Householder reflection as the first step of QR-decomposition
213+
///
214+
/// For C-continuous array,
215+
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
216+
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
217+
use qr::*;
218+
let work = HouseholderWork::<$s>::new(l)?;
219+
Ok(work.eval(a)?)
220+
}
221+
222+
/// Reconstruct Q-matrix from Householder-reflectors
223+
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()> {
224+
use qr::*;
225+
let mut work = QWork::<$s>::new(l)?;
226+
work.calc(a, tau)?;
227+
Ok(())
228+
}
229+
230+
/// Execute QR-decomposition at once
231+
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
232+
let tau = Self::householder(l, a)?;
233+
let r = Vec::from(&*a);
234+
Self::q(l, a, &tau)?;
235+
Ok(r)
236+
}
201237
}
202238
};
203239
}

lax/src/qr.rs

+2-204
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,6 @@ use crate::{error::*, layout::MatrixLayout, *};
44
use cauchy::*;
55
use num_traits::{ToPrimitive, Zero};
66

7-
pub trait QR_: Sized {
8-
/// Execute Householder reflection as the first step of QR-decomposition
9-
///
10-
/// For C-continuous array,
11-
/// this will call LQ-decomposition of the transposed matrix $ A^T = LQ^T $
12-
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
13-
14-
/// Reconstruct Q-matrix from Householder-reflectors
15-
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()>;
16-
17-
/// Execute QR-decomposition at once
18-
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
19-
}
20-
217
pub struct HouseholderWork<T: Scalar> {
228
pub m: i32,
239
pub n: i32,
@@ -136,7 +122,7 @@ pub struct QWork<T: Scalar> {
136122
pub trait QWorkImpl: Sized {
137123
type Elem: Scalar;
138124
fn new(layout: MatrixLayout) -> Result<Self>;
139-
fn calc(&mut self, a: &mut [Self::Elem], tau: &mut [Self::Elem]) -> Result<()>;
125+
fn calc(&mut self, a: &mut [Self::Elem], tau: &[Self::Elem]) -> Result<()>;
140126
}
141127

142128
macro_rules! impl_q_work {
@@ -183,7 +169,7 @@ macro_rules! impl_q_work {
183169
Ok(QWork { layout, work })
184170
}
185171

186-
fn calc(&mut self, a: &mut [Self::Elem], tau: &mut [Self::Elem]) -> Result<()> {
172+
fn calc(&mut self, a: &mut [Self::Elem], tau: &[Self::Elem]) -> Result<()> {
187173
let m = self.layout.lda();
188174
let n = self.layout.len();
189175
let k = m.min(n);
@@ -228,191 +214,3 @@ impl_q_work!(c64, lapack_sys::zungqr_, lapack_sys::zunglq_);
228214
impl_q_work!(c32, lapack_sys::cungqr_, lapack_sys::cunglq_);
229215
impl_q_work!(f64, lapack_sys::dorgqr_, lapack_sys::dorglq_);
230216
impl_q_work!(f32, lapack_sys::sorgqr_, lapack_sys::sorglq_);
231-
232-
macro_rules! impl_qr {
233-
($scalar:ty, $qrf:path, $lqf:path, $gqr:path, $glq:path) => {
234-
impl QR_ for $scalar {
235-
fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
236-
let m = l.lda();
237-
let n = l.len();
238-
let k = m.min(n);
239-
let mut tau = vec_uninit(k as usize);
240-
241-
// eval work size
242-
let mut info = 0;
243-
let mut work_size = [Self::zero()];
244-
unsafe {
245-
match l {
246-
MatrixLayout::F { .. } => {
247-
$qrf(
248-
&m,
249-
&n,
250-
AsPtr::as_mut_ptr(a),
251-
&m,
252-
AsPtr::as_mut_ptr(&mut tau),
253-
AsPtr::as_mut_ptr(&mut work_size),
254-
&(-1),
255-
&mut info,
256-
);
257-
}
258-
MatrixLayout::C { .. } => {
259-
$lqf(
260-
&m,
261-
&n,
262-
AsPtr::as_mut_ptr(a),
263-
&m,
264-
AsPtr::as_mut_ptr(&mut tau),
265-
AsPtr::as_mut_ptr(&mut work_size),
266-
&(-1),
267-
&mut info,
268-
);
269-
}
270-
}
271-
}
272-
info.as_lapack_result()?;
273-
274-
// calc
275-
let lwork = work_size[0].to_usize().unwrap();
276-
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
277-
unsafe {
278-
match l {
279-
MatrixLayout::F { .. } => {
280-
$qrf(
281-
&m,
282-
&n,
283-
AsPtr::as_mut_ptr(a),
284-
&m,
285-
AsPtr::as_mut_ptr(&mut tau),
286-
AsPtr::as_mut_ptr(&mut work),
287-
&(lwork as i32),
288-
&mut info,
289-
);
290-
}
291-
MatrixLayout::C { .. } => {
292-
$lqf(
293-
&m,
294-
&n,
295-
AsPtr::as_mut_ptr(a),
296-
&m,
297-
AsPtr::as_mut_ptr(&mut tau),
298-
AsPtr::as_mut_ptr(&mut work),
299-
&(lwork as i32),
300-
&mut info,
301-
);
302-
}
303-
}
304-
}
305-
info.as_lapack_result()?;
306-
307-
let tau = unsafe { tau.assume_init() };
308-
309-
Ok(tau)
310-
}
311-
312-
fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()> {
313-
let m = l.lda();
314-
let n = l.len();
315-
let k = m.min(n);
316-
assert_eq!(tau.len(), k as usize);
317-
318-
// eval work size
319-
let mut info = 0;
320-
let mut work_size = [Self::zero()];
321-
unsafe {
322-
match l {
323-
MatrixLayout::F { .. } => $gqr(
324-
&m,
325-
&k,
326-
&k,
327-
AsPtr::as_mut_ptr(a),
328-
&m,
329-
AsPtr::as_ptr(&tau),
330-
AsPtr::as_mut_ptr(&mut work_size),
331-
&(-1),
332-
&mut info,
333-
),
334-
MatrixLayout::C { .. } => $glq(
335-
&k,
336-
&n,
337-
&k,
338-
AsPtr::as_mut_ptr(a),
339-
&m,
340-
AsPtr::as_ptr(&tau),
341-
AsPtr::as_mut_ptr(&mut work_size),
342-
&(-1),
343-
&mut info,
344-
),
345-
}
346-
};
347-
348-
// calc
349-
let lwork = work_size[0].to_usize().unwrap();
350-
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
351-
unsafe {
352-
match l {
353-
MatrixLayout::F { .. } => $gqr(
354-
&m,
355-
&k,
356-
&k,
357-
AsPtr::as_mut_ptr(a),
358-
&m,
359-
AsPtr::as_ptr(&tau),
360-
AsPtr::as_mut_ptr(&mut work),
361-
&(lwork as i32),
362-
&mut info,
363-
),
364-
MatrixLayout::C { .. } => $glq(
365-
&k,
366-
&n,
367-
&k,
368-
AsPtr::as_mut_ptr(a),
369-
&m,
370-
AsPtr::as_ptr(&tau),
371-
AsPtr::as_mut_ptr(&mut work),
372-
&(lwork as i32),
373-
&mut info,
374-
),
375-
}
376-
}
377-
info.as_lapack_result()?;
378-
Ok(())
379-
}
380-
381-
fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
382-
let tau = Self::householder(l, a)?;
383-
let r = Vec::from(&*a);
384-
Self::q(l, a, &tau)?;
385-
Ok(r)
386-
}
387-
}
388-
};
389-
} // endmacro
390-
391-
impl_qr!(
392-
f64,
393-
lapack_sys::dgeqrf_,
394-
lapack_sys::dgelqf_,
395-
lapack_sys::dorgqr_,
396-
lapack_sys::dorglq_
397-
);
398-
impl_qr!(
399-
f32,
400-
lapack_sys::sgeqrf_,
401-
lapack_sys::sgelqf_,
402-
lapack_sys::sorgqr_,
403-
lapack_sys::sorglq_
404-
);
405-
impl_qr!(
406-
c64,
407-
lapack_sys::zgeqrf_,
408-
lapack_sys::zgelqf_,
409-
lapack_sys::zungqr_,
410-
lapack_sys::zunglq_
411-
);
412-
impl_qr!(
413-
c32,
414-
lapack_sys::cgeqrf_,
415-
lapack_sys::cgelqf_,
416-
lapack_sys::cungqr_,
417-
lapack_sys::cunglq_
418-
);

0 commit comments

Comments
 (0)