1use crate::{error::*, layout::MatrixLayout, *};
4use cauchy::*;
5use num_traits::{ToPrimitive, Zero};
6
7pub trait QR_: Sized {
8 fn householder(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
13
14 fn q(l: MatrixLayout, a: &mut [Self], tau: &[Self]) -> Result<()>;
16
17 fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>>;
19}
20
21macro_rules! impl_qr {
22 ($scalar:ty, $qrf:path, $lqf:path, $gqr:path, $glq:path) => {
23 impl QR_ for $scalar {
24 fn householder(l: MatrixLayout, mut a: &mut [Self]) -> Result<Vec<Self>> {
25 let m = l.lda();
26 let n = l.len();
27 let k = m.min(n);
28 let mut tau = unsafe { vec_uninit(k as usize) };
29
30 let mut info = 0;
32 let mut work_size = [Self::zero()];
33 unsafe {
34 match l {
35 MatrixLayout::F { .. } => {
36 $qrf(m, n, &mut a, m, &mut tau, &mut work_size, -1, &mut info);
37 }
38 MatrixLayout::C { .. } => {
39 $lqf(m, n, &mut a, m, &mut tau, &mut work_size, -1, &mut info);
40 }
41 }
42 }
43 info.as_lapack_result()?;
44
45 let lwork = work_size[0].to_usize().unwrap();
47 let mut work = unsafe { vec_uninit(lwork) };
48 unsafe {
49 match l {
50 MatrixLayout::F { .. } => {
51 $qrf(
52 m,
53 n,
54 &mut a,
55 m,
56 &mut tau,
57 &mut work,
58 lwork as i32,
59 &mut info,
60 );
61 }
62 MatrixLayout::C { .. } => {
63 $lqf(
64 m,
65 n,
66 &mut a,
67 m,
68 &mut tau,
69 &mut work,
70 lwork as i32,
71 &mut info,
72 );
73 }
74 }
75 }
76 info.as_lapack_result()?;
77
78 Ok(tau)
79 }
80
81 fn q(l: MatrixLayout, mut a: &mut [Self], tau: &[Self]) -> Result<()> {
82 let m = l.lda();
83 let n = l.len();
84 let k = m.min(n);
85 assert_eq!(tau.len(), k as usize);
86
87 let mut info = 0;
89 let mut work_size = [Self::zero()];
90 unsafe {
91 match l {
92 MatrixLayout::F { .. } => {
93 $gqr(m, k, k, &mut a, m, &tau, &mut work_size, -1, &mut info)
94 }
95 MatrixLayout::C { .. } => {
96 $glq(k, n, k, &mut a, m, &tau, &mut work_size, -1, &mut info)
97 }
98 }
99 };
100
101 let lwork = work_size[0].to_usize().unwrap();
103 let mut work = unsafe { vec_uninit(lwork) };
104 unsafe {
105 match l {
106 MatrixLayout::F { .. } => {
107 $gqr(m, k, k, &mut a, m, &tau, &mut work, lwork as i32, &mut info)
108 }
109 MatrixLayout::C { .. } => {
110 $glq(k, n, k, &mut a, m, &tau, &mut work, lwork as i32, &mut info)
111 }
112 }
113 }
114 info.as_lapack_result()?;
115 Ok(())
116 }
117
118 fn qr(l: MatrixLayout, a: &mut [Self]) -> Result<Vec<Self>> {
119 let tau = Self::householder(l, a)?;
120 let r = Vec::from(&*a);
121 Self::q(l, a, &tau)?;
122 Ok(r)
123 }
124 }
125 };
126} impl_qr!(
129 f64,
130 lapack::dgeqrf,
131 lapack::dgelqf,
132 lapack::dorgqr,
133 lapack::dorglq
134);
135impl_qr!(
136 f32,
137 lapack::sgeqrf,
138 lapack::sgelqf,
139 lapack::sorgqr,
140 lapack::sorglq
141);
142impl_qr!(
143 c64,
144 lapack::zgeqrf,
145 lapack::zgelqf,
146 lapack::zungqr,
147 lapack::zunglq
148);
149impl_qr!(
150 c32,
151 lapack::cgeqrf,
152 lapack::cgelqf,
153 lapack::cungqr,
154 lapack::cunglq
155);