Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added eigg function which is a wrapper to ggev #280

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Added eigg function which is a wrapper to ggev #280

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion lax/Cargo.toml
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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