Skip to content

Commit

Permalink
WIP: impl eig based on LAPACK
Browse files Browse the repository at this point in the history
  • Loading branch information
@termoshtt
termoshtt committed Jul 4, 2020
1 parent 37f9e4d commit bde41ed
Show file tree
Hide file tree
Showing 2 changed files with 145 additions and 55 deletions.
190 changes: 141 additions & 49 deletions lax/src/eig.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,12 @@
use crate::{error::*, layout::MatrixLayout};
use cauchy::*;
use num_traits::Zero;
use num_traits::{ToPrimitive, Zero};

/// Wraps `*geev` for real/complex
/// Wraps `*geev` for general matrices
pub trait Eig_: Scalar {
unsafe fn eig(
/// Calculate Right eigenvalue
fn eig(
calc_v: bool,
l: MatrixLayout,
a: &mut [Self],
Expand All @@ -16,65 +17,159 @@ pub trait Eig_: Scalar {
macro_rules! impl_eig_complex {
($scalar:ty, $ev:path) => {
impl Eig_ for $scalar {
unsafe fn eig(
fn eig(
calc_v: bool,
l: MatrixLayout,
mut a: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
let (n, _) = l.size();
let jobvr = if calc_v { b'V' } else { b'N' };
let mut w = vec![Self::Complex::zero(); n as usize];
let mut vl = Vec::new();
let mut vr = vec![Self::Complex::zero(); (n * n) as usize];
$ev(
l.lapacke_layout(),
b'N',
jobvr,
n,
&mut a,
n,
&mut w,
&mut vl,
n,
&mut vr,
n,
)
.as_lapack_result()?;
Ok((w, vr))
// Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate.
// However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A
let (jobvl, jobvr) = if calc_v {
match l {
MatrixLayout::C { .. } => (b'V', b'N'),
MatrixLayout::F { .. } => (b'N', b'V'),
}
} else {
(b'N', b'N')
};
let mut eigs = vec![Self::Complex::zero(); n as usize];
let mut rwork = vec![Self::Real::zero(); 2 * n as usize];

let mut vl = if jobvl == b'V' {
Some(vec![Self::Complex::zero(); (n * n) as usize])
} else {
None
};
let mut vr = if jobvr == b'V' {
Some(vec![Self::Complex::zero(); (n * n) as usize])
} else {
None
};

// calc work size
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
$ev(
jobvl,
jobvr,
n,
&mut a,
n,
&mut eigs,
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut work_size,
-1,
&mut rwork,
&mut info,
)
};
info.as_lapack_result()?;

// actal ev
let lwork = work_size[0].to_usize().unwrap();
let mut work = vec![Self::zero(); lwork];
unsafe {
$ev(
jobvl,
jobvr,
n,
&mut a,
n,
&mut eigs,
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut work,
lwork as i32,
&mut rwork,
&mut info,
)
};
info.as_lapack_result()?;

// Hermite conjugate
if jobvl == b'V' {
for c in vl.as_mut().unwrap().iter_mut() {
c.im = -c.im
}
}

Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
}
}
};
}

impl_eig_complex!(c64, lapack::zgeev);
impl_eig_complex!(c32, lapack::cgeev);

macro_rules! impl_eig_real {
($scalar:ty, $ev:path) => {
impl Eig_ for $scalar {
unsafe fn eig(
fn eig(
calc_v: bool,
l: MatrixLayout,
mut a: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
let (n, _) = l.size();
let jobvr = if calc_v { b'V' } else { b'N' };
let mut wr = vec![Self::Real::zero(); n as usize];
let mut wi = vec![Self::Real::zero(); n as usize];
let mut vl = Vec::new();
let mut vr = vec![Self::Real::zero(); (n * n) as usize];
let info = $ev(
l.lapacke_layout(),
b'N',
jobvr,
n,
&mut a,
n,
&mut wr,
&mut wi,
&mut vl,
n,
&mut vr,
n,
);
let w: Vec<Self::Complex> = wr
let mut wr = vec![Self::zero(); n as usize];
let mut wi = vec![Self::zero(); n as usize];
let mut vr = vec![Self::zero(); (n * n) as usize];

// calc work size
let mut info = 0;
let mut work_size = [0.0];
unsafe {
$ev(
b'N',
jobvr,
n,
&mut a,
n,
&mut wr,
&mut wi,
&mut [],
n,
&mut vr,
n,
&mut work_size,
-1,
&mut info,
)
};
info.as_lapack_result()?;

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work = vec![Self::zero(); lwork];
unsafe {
$ev(
b'N',
jobvr,
n,
&mut a,
n,
&mut wr,
&mut wi,
&mut [],
n,
&mut vr,
n,
&mut work,
lwork as i32,
&mut info,
)
};
info.as_lapack_result()?;

let eigs: Vec<Self::Complex> = wr
.iter()
.zip(wi.iter())
.map(|(&r, &i)| Self::Complex::new(r, i))
Expand Down Expand Up @@ -119,14 +214,11 @@ macro_rules! impl_eig_real {
})
.collect();

info.as_lapack_result()?;
Ok((w, v))
Ok((eigs, v))
}
}
};
}

impl_eig_real!(f64, lapacke::dgeev);
impl_eig_real!(f32, lapacke::sgeev);
impl_eig_complex!(c64, lapacke::zgeev);
impl_eig_complex!(c32, lapacke::cgeev);
impl_eig_real!(f64, lapack::dgeev);
impl_eig_real!(f32, lapack::sgeev);
10 changes: 4 additions & 6 deletions ndarray-linalg/src/eig.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,11 @@ where
fn eig(&self) -> Result<(Self::EigVal, Self::EigVec)> {
let mut a = self.to_owned();
let layout = a.square_layout()?;
let (s, t) = unsafe { A::eig(true, layout, a.as_allocated_mut()?)? };
let (n, _) = layout.size();
let (s, t) = A::eig(true, layout, a.as_allocated_mut()?)?;
let n = layout.len() as usize;
Ok((
ArrayBase::from(s),
ArrayBase::from(t)
.into_shape((n as usize, n as usize))
.unwrap(),
Array2::from_shape_vec((n, n).f(), t).unwrap(),
))
}
}
Expand All @@ -74,7 +72,7 @@ where

fn eigvals(&self) -> Result<Self::EigVal> {
let mut a = self.to_owned();
let (s, _) = unsafe { A::eig(true, a.square_layout()?, a.as_allocated_mut()?)? };
let (s, _) = A::eig(true, a.square_layout()?, a.as_allocated_mut()?)?;
Ok(ArrayBase::from(s))
}
}

0 comments on commit bde41ed

Please sign in to comment.