Skip to content

Commit 2119604

Browse files
selvavmselvavm
committed
Added eigg function which is a wrapper to ggev
Added doc comments for eigg function Added testcase
1 parent a561e5a commit 2119604

File tree

5 files changed

+365
-12
lines changed

5 files changed

+365
-12
lines changed

‎lax/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ intel-mkl-system = ["intel-mkl-src/mkl-dynamic-lp64-seq"]
3232
thiserror = "1.0.24"
3333
cauchy = "0.4.0"
3434
num-traits = "0.2.14"
35-
lapack = "0.18.0"
35+
lapack = "0.19.0"
3636

3737
[dependencies.intel-mkl-src]
3838
version = "0.6.0"

‎lax/src/eig.rs‎

Lines changed: 270 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,16 @@ pub trait Eig_: Scalar {
1212
l: MatrixLayout,
1313
a: &mut [Self],
1414
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
15+
fn eigg(
16+
calc_v: bool,
17+
l: MatrixLayout,
18+
a: &mut [Self],
19+
b: &mut [Self],
20+
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
1521
}
1622

1723
macro_rules! impl_eig_complex {
18-
($scalar:ty, $ev:path) => {
24+
($scalar:ty, $ev1:path, $ev2:path) => {
1925
impl Eig_ for $scalar {
2026
fn eig(
2127
calc_v: bool,
@@ -51,7 +57,7 @@ macro_rules! impl_eig_complex {
5157
let mut info = 0;
5258
let mut work_size = [Self::zero()];
5359
unsafe {
54-
$ev(
60+
$ev1(
5561
jobvl,
5662
jobvr,
5763
n,
@@ -74,7 +80,7 @@ macro_rules! impl_eig_complex {
7480
let lwork = work_size[0].to_usize().unwrap();
7581
let mut work = unsafe { vec_uninit(lwork) };
7682
unsafe {
77-
$ev(
83+
$ev1(
7884
jobvl,
7985
jobvr,
8086
n,
@@ -102,15 +108,115 @@ macro_rules! impl_eig_complex {
102108

103109
Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
104110
}
111+
112+
fn eigg(
113+
calc_v: bool,
114+
l: MatrixLayout,
115+
mut a: &mut [Self],
116+
mut b: &mut [Self],
117+
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
118+
let (n, _) = l.size();
119+
// Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate.
120+
// However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A
121+
let (jobvl, jobvr) = if calc_v {
122+
match l {
123+
MatrixLayout::C { .. } => (b'V', b'N'),
124+
MatrixLayout::F { .. } => (b'N', b'V'),
125+
}
126+
} else {
127+
(b'N', b'N')
128+
};
129+
let mut alpha = unsafe { vec_uninit(n as usize) };
130+
let mut beta = unsafe { vec_uninit(n as usize) };
131+
let mut rwork = unsafe { vec_uninit(8 * n as usize) };
132+
133+
let mut vl = if jobvl == b'V' {
134+
Some(unsafe { vec_uninit((n * n) as usize) })
135+
} else {
136+
None
137+
};
138+
let mut vr = if jobvr == b'V' {
139+
Some(unsafe { vec_uninit((n * n) as usize) })
140+
} else {
141+
None
142+
};
143+
144+
// calc work size
145+
let mut info = 0;
146+
let mut work_size = [Self::zero()];
147+
unsafe {
148+
$ev2(
149+
jobvl,
150+
jobvr,
151+
n,
152+
&mut a,
153+
n,
154+
&mut b,
155+
n,
156+
&mut alpha,
157+
&mut beta,
158+
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
159+
n,
160+
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
161+
n,
162+
&mut work_size,
163+
-1,
164+
&mut rwork,
165+
&mut info,
166+
)
167+
};
168+
info.as_lapack_result()?;
169+
170+
// actal ev
171+
let lwork = work_size[0].to_usize().unwrap();
172+
let mut work = unsafe { vec_uninit(lwork) };
173+
unsafe {
174+
$ev2(
175+
jobvl,
176+
jobvr,
177+
n,
178+
&mut a,
179+
n,
180+
&mut b,
181+
n,
182+
&mut alpha,
183+
&mut beta,
184+
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
185+
n,
186+
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
187+
n,
188+
&mut work,
189+
lwork as i32,
190+
&mut rwork,
191+
&mut info,
192+
)
193+
};
194+
info.as_lapack_result()?;
195+
196+
// Hermite conjugate
197+
if jobvl == b'V' {
198+
for c in vl.as_mut().unwrap().iter_mut() {
199+
c.im = -c.im
200+
}
201+
}
202+
// reconstruct eigenvalues
203+
let eigs: Vec<Self::Complex> = alpha
204+
.iter()
205+
.zip(beta.iter())
206+
.map(|(&a, &b)| a/b)
207+
.collect();
208+
209+
Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
210+
}
105211
}
106212
};
107213
}
108214

109-
impl_eig_complex!(c64, lapack::zgeev);
110-
impl_eig_complex!(c32, lapack::cgeev);
215+
impl_eig_complex!(c64, lapack::zgeev, lapack::zggev);
216+
impl_eig_complex!(c32, lapack::cgeev, lapack::cggev);
111217

112218
macro_rules! impl_eig_real {
113-
($scalar:ty, $ev:path) => {
219+
($scalar:ty, $ev1:path, $ev2:path) => {
114220
impl Eig_ for $scalar {
115221
fn eig(
116222
calc_v: bool,
@@ -146,7 +252,7 @@ macro_rules! impl_eig_real {
146252
let mut info = 0;
147253
let mut work_size = [0.0];
148254
unsafe {
149-
$ev(
255+
$ev1(
150256
jobvl,
151257
jobvr,
152258
n,
@@ -169,7 +275,7 @@ macro_rules! impl_eig_real {
169275
let lwork = work_size[0].to_usize().unwrap();
170276
let mut work = unsafe { vec_uninit(lwork) };
171277
unsafe {
172-
$ev(
278+
$ev1(
173279
jobvl,
174280
jobvr,
175281
n,
@@ -250,9 +356,163 @@ macro_rules! impl_eig_real {
250356

251357
Ok((eigs, eigvecs))
252358
}
359+
360+
fn eigg(
361+
calc_v: bool,
362+
l: MatrixLayout,
363+
mut a: &mut [Self],
364+
mut b: &mut [Self],
365+
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
366+
let (n, _) = l.size();
367+
// Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate.
368+
// However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A
369+
let (jobvl, jobvr) = if calc_v {
370+
match l {
371+
MatrixLayout::C { .. } => (b'V', b'N'),
372+
MatrixLayout::F { .. } => (b'N', b'V'),
373+
}
374+
} else {
375+
(b'N', b'N')
376+
};
377+
let mut alphar = unsafe { vec_uninit(n as usize) };
378+
let mut alphai = unsafe { vec_uninit(n as usize) };
379+
let mut beta = unsafe { vec_uninit(n as usize) };
380+
381+
let mut vl = if jobvl == b'V' {
382+
Some(unsafe { vec_uninit((n * n) as usize) })
383+
} else {
384+
None
385+
};
386+
let mut vr = if jobvr == b'V' {
387+
Some(unsafe { vec_uninit((n * n) as usize) })
388+
} else {
389+
None
390+
};
391+
392+
// calc work size
393+
let mut info = 0;
394+
let mut work_size = [0.0];
395+
unsafe {
396+
$ev2(
397+
jobvl,
398+
jobvr,
399+
n,
400+
&mut a,
401+
n,
402+
&mut b,
403+
n,
404+
&mut alphar,
405+
&mut alphai,
406+
&mut beta,
407+
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
408+
n,
409+
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
410+
n,
411+
&mut work_size,
412+
-1,
413+
&mut info,
414+
)
415+
};
416+
info.as_lapack_result()?;
417+
418+
// actual ev
419+
let lwork = work_size[0].to_usize().unwrap();
420+
let mut work = unsafe { vec_uninit(lwork) };
421+
unsafe {
422+
$ev2(
423+
jobvl,
424+
jobvr,
425+
n,
426+
&mut a,
427+
n,
428+
&mut b,
429+
n,
430+
&mut alphar,
431+
&mut alphai,
432+
&mut beta,
433+
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
434+
n,
435+
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
436+
n,
437+
&mut work,
438+
lwork as i32,
439+
&mut info,
440+
)
441+
};
442+
info.as_lapack_result()?;
443+
444+
// reconstruct eigenvalues
445+
let alpha: Vec<Self::Complex> = alphar
446+
.iter()
447+
.zip(alphai.iter())
448+
.map(|(&re, &im)| Self::complex(re, im))
449+
.collect();
450+
451+
let eigs: Vec<Self::Complex> = alpha
452+
.iter()
453+
.zip(beta.iter())
454+
.map(|(&a, &b)| a/b)
455+
.collect();
456+
457+
if !calc_v {
458+
return Ok((eigs, Vec::new()));
459+
}
460+
461+
462+
// Reconstruct eigenvectors into complex-array
463+
// --------------------------------------------
464+
//
465+
// From LAPACK API https://software.intel.com/en-us/node/469230
466+
//
467+
// - If the j-th eigenvalue is real,
468+
// - v(j) = VR(:,j), the j-th column of VR.
469+
//
470+
// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
471+
// - v(j) = VR(:,j) + i*VR(:,j+1)
472+
// - v(j+1) = VR(:,j) - i*VR(:,j+1).
473+
//
474+
// ```
475+
// j -> <----pair----> <----pair---->
476+
// [ ... (real), (imag), (imag), (imag), (imag), ... ] : eigs
477+
// ^ ^ ^ ^ ^
478+
// false false true false true : is_conjugate_pair
479+
// ```
480+
let n = n as usize;
481+
let v = vr.or(vl).unwrap();
482+
let mut eigvecs = unsafe { vec_uninit(n * n) };
483+
let mut is_conjugate_pair = false; // flag for check `j` is complex conjugate
484+
for j in 0..n {
485+
if alphai[j] == 0.0 {
486+
// j-th eigenvalue is real
487+
for i in 0..n {
488+
eigvecs[i + j * n] = Self::complex(v[i + j * n], 0.0);
489+
}
490+
} else {
491+
// j-th eigenvalue is complex
492+
// complex conjugated pair can be `j-1` or `j+1`
493+
if is_conjugate_pair {
494+
let j_pair = j - 1;
495+
assert!(j_pair < n);
496+
for i in 0..n {
497+
eigvecs[i + j * n] = Self::complex(v[i + j_pair * n], v[i + j * n]);
498+
}
499+
} else {
500+
let j_pair = j + 1;
501+
assert!(j_pair < n);
502+
for i in 0..n {
503+
eigvecs[i + j * n] =
504+
Self::complex(v[i + j * n], -v[i + j_pair * n]);
505+
}
506+
}
507+
is_conjugate_pair = !is_conjugate_pair;
508+
}
509+
}
510+
511+
Ok((eigs, eigvecs))
512+
}
253513
}
254514
};
255515
}
256516

257-
impl_eig_real!(f64, lapack::dgeev);
258-
impl_eig_real!(f32, lapack::sgeev);
517+
impl_eig_real!(f64, lapack::dgeev, lapack::dggev);
518+
impl_eig_real!(f32, lapack::sgeev, lapack::sggev);

‎ndarray-linalg/examples/eig.rs‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use ndarray::*;
22
use ndarray_linalg::*;
33

4-
fn main() {
4+
fn eig() {
55
let a = arr2(&[[2.0, 1.0, 2.0], [-2.0, 2.0, 1.0], [1.0, 2.0, -2.0]]);
66
let (e, vecs) = a.clone().eig().unwrap();
77
println!("eigenvalues = \n{:?}", e);
@@ -10,3 +10,26 @@ fn main() {
1010
let av = a_c.dot(&vecs);
1111
println!("AV = \n{:?}", av);
1212
}
13+
14+
fn eigg_real() {
15+
let a = arr2(&[[1.0/2.0.sqrt(), 0.0], [0.0, 1.0]]);
16+
let b = arr2(&[[0.0, 1.0], [-1.0/2.0.sqrt(), 0.0]]);
17+
let (e, vecs) = a.clone().eigg(&b).unwrap();
18+
println!("eigenvalues = \n{:?}", e);
19+
println!("V = \n{:?}", vecs);
20+
}
21+
22+
fn eigg_complex() {
23+
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)]]);
24+
let b = a.clone();
25+
let (e, vecs) = a.clone().eigg(&b).unwrap();
26+
println!("eigenvalues = \n{:?}", &e);
27+
println!("V = \n{:?}", vecs);
28+
29+
}
30+
31+
fn main() {
32+
eig();
33+
eigg_real();
34+
eigg_complex();
35+
}

0 commit comments

Comments
 (0)