lax/
qr.rs

1//! QR decomposition
2
3use crate::{error::*, layout::MatrixLayout, *};
4use cauchy::*;
5use num_traits::{ToPrimitive, Zero};
6
7pub 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
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                // eval work size
31                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                // calc
46                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                // eval work size
88                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                // calc
102                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} // endmacro
127
128impl_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);