Skip to content

Commit

Permalink
Added eigg function which is a wrapper to ggev
Browse files Browse the repository at this point in the history 
Added doc comments for eigg function

Added testcase
  • Loading branch information
@selvavm
selvavm committed Apr 18, 2021
1 parent a561e5a commit 2d6ab5f
Show file tree
Hide file tree
Showing 5 changed files with 370 additions and 12 deletions.
2 changes: 1 addition & 1 deletion lax/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ intel-mkl-system = ["intel-mkl-src/mkl-dynamic-lp64-seq"]
thiserror = "1.0.24"
cauchy = "0.4.0"
num-traits = "0.2.14"
lapack = "0.18.0"
lapack = "0.19.0"

[dependencies.intel-mkl-src]
version = "0.6.0"
Expand Down
279 changes: 269 additions & 10 deletions lax/src/eig.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,16 @@ pub trait Eig_: Scalar {
l: MatrixLayout,
a: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
fn eigg(
calc_v: bool,
l: MatrixLayout,
a: &mut [Self],
b: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
}

macro_rules! impl_eig_complex {
($scalar:ty, $ev:path) => {
($scalar:ty, $ev1:path, $ev2:path) => {
impl Eig_ for $scalar {
fn eig(
calc_v: bool,
Expand Down Expand Up @@ -51,7 +57,7 @@ macro_rules! impl_eig_complex {
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
$ev(
$ev1(
jobvl,
jobvr,
n,
Expand All @@ -74,7 +80,7 @@ macro_rules! impl_eig_complex {
let lwork = work_size[0].to_usize().unwrap();
let mut work = unsafe { vec_uninit(lwork) };
unsafe {
$ev(
$ev1(
jobvl,
jobvr,
n,
Expand Down Expand Up @@ -102,15 +108,115 @@ macro_rules! impl_eig_complex {

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

fn eigg(
calc_v: bool,
l: MatrixLayout,
mut a: &mut [Self],
mut b: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
let (n, _) = l.size();
// 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 alpha = unsafe { vec_uninit(n as usize) };
let mut beta = unsafe { vec_uninit(n as usize) };
let mut rwork = unsafe { vec_uninit(8 * n as usize) };

let mut vl = if jobvl == b'V' {
Some(unsafe { vec_uninit((n * n) as usize) })
} else {
None
};
let mut vr = if jobvr == b'V' {
Some(unsafe { vec_uninit((n * n) as usize) })
} else {
None
};

// calc work size
let mut info = 0;
let mut work_size = [Self::zero()];
unsafe {
$ev2(
jobvl,
jobvr,
n,
&mut a,
n,
&mut b,
n,
&mut alpha,
&mut beta,
&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 = unsafe { vec_uninit(lwork) };
unsafe {
$ev2(
jobvl,
jobvr,
n,
&mut a,
n,
&mut b,
n,
&mut alpha,
&mut beta,
&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
}
}
// reconstruct eigenvalues
let eigs: Vec<Self::Complex> = alpha
.iter()
.zip(beta.iter())
.map(|(&a, &b)| a / b)
.collect();

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

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

macro_rules! impl_eig_real {
($scalar:ty, $ev:path) => {
($scalar:ty, $ev1:path, $ev2:path) => {
impl Eig_ for $scalar {
fn eig(
calc_v: bool,
Expand Down Expand Up @@ -146,7 +252,7 @@ macro_rules! impl_eig_real {
let mut info = 0;
let mut work_size = [0.0];
unsafe {
$ev(
$ev1(
jobvl,
jobvr,
n,
Expand All @@ -169,7 +275,7 @@ macro_rules! impl_eig_real {
let lwork = work_size[0].to_usize().unwrap();
let mut work = unsafe { vec_uninit(lwork) };
unsafe {
$ev(
$ev1(
jobvl,
jobvr,
n,
Expand Down Expand Up @@ -250,9 +356,162 @@ macro_rules! impl_eig_real {

Ok((eigs, eigvecs))
}

fn eigg(
calc_v: bool,
l: MatrixLayout,
mut a: &mut [Self],
mut b: &mut [Self],
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
let (n, _) = l.size();
// 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 alphar = unsafe { vec_uninit(n as usize) };
let mut alphai = unsafe { vec_uninit(n as usize) };
let mut beta = unsafe { vec_uninit(n as usize) };

let mut vl = if jobvl == b'V' {
Some(unsafe { vec_uninit((n * n) as usize) })
} else {
None
};
let mut vr = if jobvr == b'V' {
Some(unsafe { vec_uninit((n * n) as usize) })
} else {
None
};

// calc work size
let mut info = 0;
let mut work_size = [0.0];
unsafe {
$ev2(
jobvl,
jobvr,
n,
&mut a,
n,
&mut b,
n,
&mut alphar,
&mut alphai,
&mut beta,
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut work_size,
-1,
&mut info,
)
};
info.as_lapack_result()?;

// actual ev
let lwork = work_size[0].to_usize().unwrap();
let mut work = unsafe { vec_uninit(lwork) };
unsafe {
$ev2(
jobvl,
jobvr,
n,
&mut a,
n,
&mut b,
n,
&mut alphar,
&mut alphai,
&mut beta,
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
n,
&mut work,
lwork as i32,
&mut info,
)
};
info.as_lapack_result()?;

// reconstruct eigenvalues
let alpha: Vec<Self::Complex> = alphar
.iter()
.zip(alphai.iter())
.map(|(&re, &im)| Self::complex(re, im))
.collect();

let eigs: Vec<Self::Complex> = alpha
.iter()
.zip(beta.iter())
.map(|(&a, &b)| a / b)
.collect();

if !calc_v {
return Ok((eigs, Vec::new()));
}

// Reconstruct eigenvectors into complex-array
// --------------------------------------------
//
// From LAPACK API https://software.intel.com/en-us/node/469230
//
// - If the j-th eigenvalue is real,
// - v(j) = VR(:,j), the j-th column of VR.
//
// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
// - v(j) = VR(:,j) + i*VR(:,j+1)
// - v(j+1) = VR(:,j) - i*VR(:,j+1).
//
// ```
// j -> <----pair----> <----pair---->
// [ ... (real), (imag), (imag), (imag), (imag), ... ] : eigs
// ^ ^ ^ ^ ^
// false false true false true : is_conjugate_pair
// ```
let n = n as usize;
let v = vr.or(vl).unwrap();
let mut eigvecs = unsafe { vec_uninit(n * n) };
let mut is_conjugate_pair = false; // flag for check `j` is complex conjugate
for j in 0..n {
if alphai[j] == 0.0 {
// j-th eigenvalue is real
for i in 0..n {
eigvecs[i + j * n] = Self::complex(v[i + j * n], 0.0);
}
} else {
// j-th eigenvalue is complex
// complex conjugated pair can be `j-1` or `j+1`
if is_conjugate_pair {
let j_pair = j - 1;
assert!(j_pair < n);
for i in 0..n {
eigvecs[i + j * n] = Self::complex(v[i + j_pair * n], v[i + j * n]);
}
} else {
let j_pair = j + 1;
assert!(j_pair < n);
for i in 0..n {
eigvecs[i + j * n] =
Self::complex(v[i + j * n], -v[i + j_pair * n]);
}
}
is_conjugate_pair = !is_conjugate_pair;
}
}

Ok((eigs, eigvecs))
}
}
};
}

impl_eig_real!(f64, lapack::dgeev);
impl_eig_real!(f32, lapack::sgeev);
impl_eig_real!(f64, lapack::dgeev, lapack::dggev);
impl_eig_real!(f32, lapack::sgeev, lapack::sggev);
27 changes: 26 additions & 1 deletion ndarray-linalg/examples/eig.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use ndarray::*;
use ndarray_linalg::*;

fn main() {
fn eig() {
let a = arr2(&[[2.0, 1.0, 2.0], [-2.0, 2.0, 1.0], [1.0, 2.0, -2.0]]);
let (e, vecs) = a.clone().eig().unwrap();
println!("eigenvalues = \n{:?}", e);
Expand All @@ -10,3 +10,28 @@ fn main() {
let av = a_c.dot(&vecs);
println!("AV = \n{:?}", av);
}

fn eigg_real() {
let a = arr2(&[[1.0 / 2.0.sqrt(), 0.0], [0.0, 1.0]]);
let b = arr2(&[[0.0, 1.0], [-1.0 / 2.0.sqrt(), 0.0]]);
let (e, vecs) = a.clone().eigg(&b).unwrap();
println!("eigenvalues = \n{:?}", e);
println!("V = \n{:?}", vecs);
}

fn eigg_complex() {
let a = arr2(&[
[c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)],
[c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)],
]);
let b = a.clone();
let (e, vecs) = a.clone().eigg(&b).unwrap();
println!("eigenvalues = \n{:?}", &e);
println!("V = \n{:?}", vecs);
}

fn main() {
eig();
eigg_real();
eigg_complex();
}
Loading

0 comments on commit 2d6ab5f

Please sign in to comment.