Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 2d6ab5f

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

File tree

5 files changed

+370
-12
lines changed

5 files changed

+370
-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: 269 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,162 @@ 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+
// Reconstruct eigenvectors into complex-array
462+
// --------------------------------------------
463+
//
464+
// From LAPACK API https://software.intel.com/en-us/node/469230
465+
//
466+
// - If the j-th eigenvalue is real,
467+
// - v(j) = VR(:,j), the j-th column of VR.
468+
//
469+
// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
470+
// - v(j) = VR(:,j) + i*VR(:,j+1)
471+
// - v(j+1) = VR(:,j) - i*VR(:,j+1).
472+
//
473+
// ```
474+
// j -> <----pair----> <----pair---->
475+
// [ ... (real), (imag), (imag), (imag), (imag), ... ] : eigs
476+
// ^ ^ ^ ^ ^
477+
// false false true false true : is_conjugate_pair
478+
// ```
479+
let n = n as usize;
480+
let v = vr.or(vl).unwrap();
481+
let mut eigvecs = unsafe { vec_uninit(n * n) };
482+
let mut is_conjugate_pair = false; // flag for check `j` is complex conjugate
483+
for j in 0..n {
484+
if alphai[j] == 0.0 {
485+
// j-th eigenvalue is real
486+
for i in 0..n {
487+
eigvecs[i + j * n] = Self::complex(v[i + j * n], 0.0);
488+
}
489+
} else {
490+
// j-th eigenvalue is complex
491+
// complex conjugated pair can be `j-1` or `j+1`
492+
if is_conjugate_pair {
493+
let j_pair = j - 1;
494+
assert!(j_pair < n);
495+
for i in 0..n {
496+
eigvecs[i + j * n] = Self::complex(v[i + j_pair * n], v[i + j * n]);
497+
}
498+
} else {
499+
let j_pair = j + 1;
500+
assert!(j_pair < n);
501+
for i in 0..n {
502+
eigvecs[i + j * n] =
503+
Self::complex(v[i + j * n], -v[i + j_pair * n]);
504+
}
505+
}
506+
is_conjugate_pair = !is_conjugate_pair;
507+
}
508+
}
509+
510+
Ok((eigs, eigvecs))
511+
}
253512
}
254513
};
255514
}
256515

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

‎ndarray-linalg/examples/eig.rs‎

Lines changed: 26 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,28 @@ 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(&[
24+
[c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)],
25+
[c64::complex(-3.84, 2.25), c64::complex(-3.84, 2.25)],
26+
]);
27+
let b = a.clone();
28+
let (e, vecs) = a.clone().eigg(&b).unwrap();
29+
println!("eigenvalues = \n{:?}", &e);
30+
println!("V = \n{:?}", vecs);
31+
}
32+
33+
fn main() {
34+
eig();
35+
eigg_real();
36+
eigg_complex();
37+
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /