nalgebra/nalgebra-lapack/src/qr.rs

271 lines
7.3 KiB
Rust
Raw Normal View History

#[cfg(feature = "serde-serialize")]
2018-10-22 13:00:10 +08:00
use serde::{Deserialize, Serialize};
use num::Zero;
use num_complex::Complex;
2020-04-06 00:49:48 +08:00
use crate::ComplexHelper;
use na::allocator::Allocator;
2021-04-12 16:34:44 +08:00
use na::dimension::{Const, Dim, DimMin, DimMinimum};
2021-04-11 17:00:38 +08:00
use na::{DefaultAllocator, Matrix, OMatrix, OVector, Scalar};
use lapack;
/// The QR decomposition of a general matrix.
#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
#[cfg_attr(
feature = "serde-serialize",
2021-04-11 17:00:38 +08:00
serde(bound(serialize = "DefaultAllocator: Allocator<T, R, C> +
Allocator<T, DimMinimum<R, C>>,
OMatrix<T, R, C>: Serialize,
OVector<T, DimMinimum<R, C>>: Serialize"))
)]
#[cfg_attr(
feature = "serde-serialize",
2021-04-11 17:00:38 +08:00
serde(bound(deserialize = "DefaultAllocator: Allocator<T, R, C> +
Allocator<T, DimMinimum<R, C>>,
OMatrix<T, R, C>: Deserialize<'de>,
OVector<T, DimMinimum<R, C>>: Deserialize<'de>"))
)]
#[derive(Clone, Debug)]
2021-04-11 17:00:38 +08:00
pub struct QR<T: Scalar, R: DimMin<C>, C: Dim>
2020-04-06 00:49:48 +08:00
where
2021-04-11 17:00:38 +08:00
DefaultAllocator: Allocator<T, R, C> + Allocator<T, DimMinimum<R, C>>,
2018-02-02 19:26:35 +08:00
{
2021-04-11 17:00:38 +08:00
qr: OMatrix<T, R, C>,
tau: OVector<T, DimMinimum<R, C>>,
}
2021-04-11 17:00:38 +08:00
impl<T: Scalar + Copy, R: DimMin<C>, C: Dim> Copy for QR<T, R, C>
2018-02-02 19:26:35 +08:00
where
2021-04-11 17:00:38 +08:00
DefaultAllocator: Allocator<T, R, C> + Allocator<T, DimMinimum<R, C>>,
OMatrix<T, R, C>: Copy,
OVector<T, DimMinimum<R, C>>: Copy,
2020-04-06 00:49:48 +08:00
{
}
2021-04-11 17:00:38 +08:00
impl<T: QRScalar + Zero, R: DimMin<C>, C: Dim> QR<T, R, C>
2020-04-06 00:49:48 +08:00
where
2021-04-11 17:00:38 +08:00
DefaultAllocator: Allocator<T, R, C>
+ Allocator<T, R, DimMinimum<R, C>>
+ Allocator<T, DimMinimum<R, C>, C>
+ Allocator<T, DimMinimum<R, C>>,
2018-02-02 19:26:35 +08:00
{
/// Computes the QR decomposition of the matrix `m`.
2021-04-11 17:00:38 +08:00
pub fn new(mut m: OMatrix<T, R, C>) -> Self {
let (nrows, ncols) = m.shape_generic();
2018-02-02 19:26:35 +08:00
let mut info = 0;
2021-08-03 23:39:45 +08:00
let mut tau = Matrix::zeros_generic(nrows.min(ncols), Const::<1>);
if nrows.value() == 0 || ncols.value() == 0 {
2021-07-28 07:18:29 +08:00
return Self { qr: m, tau };
}
2021-04-11 17:00:38 +08:00
let lwork = T::xgeqrf_work_size(
2018-02-02 19:26:35 +08:00
nrows.value() as i32,
ncols.value() as i32,
m.as_mut_slice(),
nrows.value() as i32,
tau.as_mut_slice(),
&mut info,
);
2021-08-03 23:39:45 +08:00
let mut work = vec![T::zero(); lwork as usize];
2021-04-11 17:00:38 +08:00
T::xgeqrf(
2018-02-02 19:26:35 +08:00
nrows.value() as i32,
ncols.value() as i32,
m.as_mut_slice(),
nrows.value() as i32,
tau.as_mut_slice(),
&mut work,
lwork,
&mut info,
);
2021-07-28 07:18:29 +08:00
Self { qr: m, tau }
}
/// Retrieves the upper trapezoidal submatrix `R` of this decomposition.
#[inline]
#[must_use]
2021-04-11 17:00:38 +08:00
pub fn r(&self) -> OMatrix<T, DimMinimum<R, C>, C> {
let (nrows, ncols) = self.qr.shape_generic();
self.qr.rows_generic(0, nrows.min(ncols)).upper_triangle()
}
}
2021-04-11 17:00:38 +08:00
impl<T: QRReal + Zero, R: DimMin<C>, C: Dim> QR<T, R, C>
2020-04-06 00:49:48 +08:00
where
2021-04-11 17:00:38 +08:00
DefaultAllocator: Allocator<T, R, C>
+ Allocator<T, R, DimMinimum<R, C>>
+ Allocator<T, DimMinimum<R, C>, C>
+ Allocator<T, DimMinimum<R, C>>,
2018-02-02 19:26:35 +08:00
{
/// Retrieves the matrices `(Q, R)` of this decompositions.
2018-02-02 19:26:35 +08:00
pub fn unpack(
self,
) -> (
2021-04-11 17:00:38 +08:00
OMatrix<T, R, DimMinimum<R, C>>,
OMatrix<T, DimMinimum<R, C>, C>,
2018-02-02 19:26:35 +08:00
) {
(self.q(), self.r())
}
/// Computes the orthogonal matrix `Q` of this decomposition.
#[inline]
#[must_use]
2021-04-11 17:00:38 +08:00
pub fn q(&self) -> OMatrix<T, R, DimMinimum<R, C>> {
let (nrows, ncols) = self.qr.shape_generic();
let min_nrows_ncols = nrows.min(ncols);
if min_nrows_ncols.value() == 0 {
2021-04-11 17:00:38 +08:00
return OMatrix::from_element_generic(nrows, min_nrows_ncols, T::zero());
}
2018-10-22 13:00:10 +08:00
let mut q = self
.qr
2018-02-02 19:26:35 +08:00
.generic_slice((0, 0), (nrows, min_nrows_ncols))
.into_owned();
2018-02-02 19:26:35 +08:00
let mut info = 0;
let nrows = nrows.value() as i32;
2021-04-11 17:00:38 +08:00
let lwork = T::xorgqr_work_size(
2018-02-02 19:26:35 +08:00
nrows,
min_nrows_ncols.value() as i32,
self.tau.len() as i32,
q.as_mut_slice(),
nrows,
self.tau.as_slice(),
&mut info,
);
2021-04-11 17:00:38 +08:00
let mut work = vec![T::zero(); lwork as usize];
2018-02-02 19:26:35 +08:00
2021-04-11 17:00:38 +08:00
T::xorgqr(
2018-02-02 19:26:35 +08:00
nrows,
min_nrows_ncols.value() as i32,
self.tau.len() as i32,
q.as_mut_slice(),
nrows,
self.tau.as_slice(),
&mut work,
lwork,
&mut info,
);
q
}
}
/*
*
* Lapack functions dispatch.
*
*/
2018-09-24 12:48:42 +08:00
/// Trait implemented by scalar types for which Lapack function exist to compute the
/// QR decomposition.
pub trait QRScalar: Scalar + Copy {
2018-02-02 19:26:35 +08:00
fn xgeqrf(
m: i32,
n: i32,
a: &mut [Self],
lda: i32,
tau: &mut [Self],
work: &mut [Self],
lwork: i32,
info: &mut i32,
);
fn xgeqrf_work_size(
m: i32,
n: i32,
a: &mut [Self],
lda: i32,
tau: &mut [Self],
info: &mut i32,
) -> i32;
}
2018-09-24 12:48:42 +08:00
/// Trait implemented by reals for which Lapack function exist to compute the
/// QR decomposition.
pub trait QRReal: QRScalar {
#[allow(missing_docs)]
2018-02-02 19:26:35 +08:00
fn xorgqr(
m: i32,
n: i32,
k: i32,
a: &mut [Self],
lda: i32,
tau: &[Self],
work: &mut [Self],
lwork: i32,
info: &mut i32,
);
#[allow(missing_docs)]
2018-02-02 19:26:35 +08:00
fn xorgqr_work_size(
m: i32,
n: i32,
k: i32,
a: &mut [Self],
lda: i32,
tau: &[Self],
info: &mut i32,
) -> i32;
}
macro_rules! qr_scalar_impl(
($N: ty, $xgeqrf: path) => (
impl QRScalar for $N {
#[inline]
fn xgeqrf(m: i32, n: i32, a: &mut [Self], lda: i32, tau: &mut [Self],
work: &mut [Self], lwork: i32, info: &mut i32) {
unsafe { $xgeqrf(m, n, a, lda, tau, work, lwork, info) }
}
#[inline]
fn xgeqrf_work_size(m: i32, n: i32, a: &mut [Self], lda: i32, tau: &mut [Self],
info: &mut i32) -> i32 {
let mut work = [ Zero::zero() ];
let lwork = -1 as i32;
unsafe { $xgeqrf(m, n, a, lda, tau, &mut work, lwork, info); }
ComplexHelper::real_part(work[0]) as i32
}
}
)
);
macro_rules! qr_real_impl(
($N: ty, $xorgqr: path) => (
impl QRReal for $N {
#[inline]
fn xorgqr(m: i32, n: i32, k: i32, a: &mut [Self], lda: i32, tau: &[Self],
work: &mut [Self], lwork: i32, info: &mut i32) {
unsafe { $xorgqr(m, n, k, a, lda, tau, work, lwork, info) }
}
#[inline]
fn xorgqr_work_size(m: i32, n: i32, k: i32, a: &mut [Self], lda: i32, tau: &[Self],
info: &mut i32) -> i32 {
let mut work = [ Zero::zero() ];
let lwork = -1 as i32;
unsafe { $xorgqr(m, n, k, a, lda, tau, &mut work, lwork, info); }
ComplexHelper::real_part(work[0]) as i32
}
}
)
);
qr_scalar_impl!(f32, lapack::sgeqrf);
qr_scalar_impl!(f64, lapack::dgeqrf);
qr_scalar_impl!(Complex<f32>, lapack::cgeqrf);
qr_scalar_impl!(Complex<f64>, lapack::zgeqrf);
qr_real_impl!(f32, lapack::sorgqr);
qr_real_impl!(f64, lapack::dorgqr);