Skip to content

Commit 1d0a880

Browse files
termoshtttermoshtt
committed
impl eig based on LAPACK *geev
1 parent 0e88d0c commit 1d0a880

File tree

2 files changed

+214
-91
lines changed

2 files changed

+214
-91
lines changed

‎lax/src/eig.rs

+210-85
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,12 @@
22
33
use crate::{error::*, layout::MatrixLayout};
44
use cauchy::*;
5-
use num_traits::Zero;
5+
use num_traits::{ToPrimitive, Zero};
66

7-
/// Wraps `*geev` for real/complex
7+
/// Wraps `*geev` for general matrices
88
pub trait Eig_: Scalar {
9-
unsafe fn eig(
9+
/// Calculate Right eigenvalue
10+
fn eig(
1011
calc_v: bool,
1112
l: MatrixLayout,
1213
a: &mut [Self],
@@ -16,117 +17,241 @@ pub trait Eig_: Scalar {
1617
macro_rules! impl_eig_complex {
1718
($scalar:ty, $ev:path) => {
1819
impl Eig_ for $scalar {
19-
unsafe fn eig(
20+
fn eig(
2021
calc_v: bool,
2122
l: MatrixLayout,
2223
mut a: &mut [Self],
2324
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
2425
let (n, _) = l.size();
25-
let jobvr = if calc_v { b'V' } else { b'N' };
26-
let mut w = vec![Self::Complex::zero(); n as usize];
27-
let mut vl = Vec::new();
28-
let mut vr = vec![Self::Complex::zero(); (n * n) as usize];
29-
$ev(
30-
l.lapacke_layout(),
31-
b'N',
32-
jobvr,
33-
n,
34-
&mut a,
35-
n,
36-
&mut w,
37-
&mut vl,
38-
n,
39-
&mut vr,
40-
n,
41-
)
42-
.as_lapack_result()?;
43-
Ok((w, vr))
26+
// Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate.
27+
// However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A
28+
let (jobvl, jobvr) = if calc_v {
29+
match l {
30+
MatrixLayout::C { .. } => (b'V', b'N'),
31+
MatrixLayout::F { .. } => (b'N', b'V'),
32+
}
33+
} else {
34+
(b'N', b'N')
35+
};
36+
let mut eigs = vec![Self::Complex::zero(); n as usize];
37+
let mut rwork = vec![Self::Real::zero(); 2 * n as usize];
38+
39+
let mut vl = if jobvl == b'V' {
40+
Some(vec![Self::Complex::zero(); (n * n) as usize])
41+
} else {
42+
None
43+
};
44+
let mut vr = if jobvr == b'V' {
45+
Some(vec![Self::Complex::zero(); (n * n) as usize])
46+
} else {
47+
None
48+
};
49+
50+
// calc work size
51+
let mut info = 0;
52+
let mut work_size = [Self::zero()];
53+
unsafe {
54+
$ev(
55+
jobvl,
56+
jobvr,
57+
n,
58+
&mut a,
59+
n,
60+
&mut eigs,
61+
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
62+
n,
63+
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
64+
n,
65+
&mut work_size,
66+
-1,
67+
&mut rwork,
68+
&mut info,
69+
)
70+
};
71+
info.as_lapack_result()?;
72+
73+
// actal ev
74+
let lwork = work_size[0].to_usize().unwrap();
75+
let mut work = vec![Self::zero(); lwork];
76+
unsafe {
77+
$ev(
78+
jobvl,
79+
jobvr,
80+
n,
81+
&mut a,
82+
n,
83+
&mut eigs,
84+
&mut vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
85+
n,
86+
&mut vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
87+
n,
88+
&mut work,
89+
lwork as i32,
90+
&mut rwork,
91+
&mut info,
92+
)
93+
};
94+
info.as_lapack_result()?;
95+
96+
// Hermite conjugate
97+
if jobvl == b'V' {
98+
for c in vl.as_mut().unwrap().iter_mut() {
99+
c.im = -c.im
100+
}
101+
}
102+
103+
Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
44104
}
45105
}
46106
};
47107
}
48108

109+
impl_eig_complex!(c64, lapack::zgeev);
110+
impl_eig_complex!(c32, lapack::cgeev);
111+
49112
macro_rules! impl_eig_real {
50113
($scalar:ty, $ev:path) => {
51114
impl Eig_ for $scalar {
52-
unsafe fn eig(
115+
fn eig(
53116
calc_v: bool,
54117
l: MatrixLayout,
55118
mut a: &mut [Self],
56119
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
57120
let (n, _) = l.size();
58-
let jobvr = if calc_v { b'V' } else { b'N' };
59-
let mut wr = vec![Self::Real::zero(); n as usize];
60-
let mut wi = vec![Self::Real::zero(); n as usize];
61-
let mut vl = Vec::new();
62-
let mut vr = vec![Self::Real::zero(); (n * n) as usize];
63-
let info = $ev(
64-
l.lapacke_layout(),
65-
b'N',
66-
jobvr,
67-
n,
68-
&mut a,
69-
n,
70-
&mut wr,
71-
&mut wi,
72-
&mut vl,
73-
n,
74-
&mut vr,
75-
n,
76-
);
77-
let w: Vec<Self::Complex> = wr
121+
// Because LAPACK assumes F-continious array, C-continious array should be taken Hermitian conjugate.
122+
// However, we utilize a fact that left eigenvector of A^H corresponds to the right eigenvector of A
123+
let (jobvl, jobvr) = if calc_v {
124+
match l {
125+
MatrixLayout::C { .. } => (b'V', b'N'),
126+
MatrixLayout::F { .. } => (b'N', b'V'),
127+
}
128+
} else {
129+
(b'N', b'N')
130+
};
131+
let mut eig_re = vec![Self::zero(); n as usize];
132+
let mut eig_im = vec![Self::zero(); n as usize];
133+
134+
let mut vl = if jobvl == b'V' {
135+
Some(vec![Self::zero(); (n * n) as usize])
136+
} else {
137+
None
138+
};
139+
let mut vr = if jobvr == b'V' {
140+
Some(vec![Self::zero(); (n * n) as usize])
141+
} else {
142+
None
143+
};
144+
145+
// calc work size
146+
let mut info = 0;
147+
let mut work_size = [0.0];
148+
unsafe {
149+
$ev(
150+
jobvl,
151+
jobvr,
152+
n,
153+
&mut a,
154+
n,
155+
&mut eig_re,
156+
&mut eig_im,
157+
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
158+
n,
159+
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
160+
n,
161+
&mut work_size,
162+
-1,
163+
&mut info,
164+
)
165+
};
166+
info.as_lapack_result()?;
167+
168+
// actual ev
169+
let lwork = work_size[0].to_usize().unwrap();
170+
let mut work = vec![Self::zero(); lwork];
171+
unsafe {
172+
$ev(
173+
jobvl,
174+
jobvr,
175+
n,
176+
&mut a,
177+
n,
178+
&mut eig_re,
179+
&mut eig_im,
180+
vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
181+
n,
182+
vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut []),
183+
n,
184+
&mut work,
185+
lwork as i32,
186+
&mut info,
187+
)
188+
};
189+
info.as_lapack_result()?;
190+
191+
// reconstruct eigenvalues
192+
let eigs: Vec<Self::Complex> = eig_re
78193
.iter()
79-
.zip(wi.iter())
80-
.map(|(&r, &i)| Self::Complex::new(r, i))
194+
.zip(eig_im.iter())
195+
.map(|(&re, &im)| Self::complex(re, im))
81196
.collect();
82197

83-
// If the j-th eigenvalue is real, then
84-
// eigenvector = [ vr[j], vr[j+n], vr[j+2*n], ... ].
198+
if !calc_v {
199+
return Ok((eigs, Vec::new()));
200+
}
201+
202+
// Reconstruct eigenvectors into complex-array
203+
// --------------------------------------------
85204
//
86-
// If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
87-
// eigenvector(j) = [ vr[j] + i*vr[j+1], vr[j+n] + i*vr[j+n+1], vr[j+2*n] + i*vr[j+2*n+1], ... ] and
88-
// eigenvector(j+1) = [ vr[j] - i*vr[j+1], vr[j+n] - i*vr[j+n+1], vr[j+2*n] - i*vr[j+2*n+1], ... ].
205+
// From LAPACK API https://software.intel.com/en-us/node/469230
89206
//
90-
// Therefore, if eigenvector(j) is written as [ v_{j0}, v_{j1}, v_{j2}, ... ],
91-
// you have to make
92-
// v = vec![ v_{00}, v_{10}, v_{20}, ..., v_{jk}, v_{(j+1)k}, v_{(j+2)k}, ... ] (v.len() = n*n)
93-
// based on wi and vr.
94-
// After that, v is converted to Array2 (see ../eig.rs).
207+
// - If the j-th eigenvalue is real,
208+
// - v(j) = VR(:,j), the j-th column of VR.
209+
//
210+
// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
211+
// - v(j) = VR(:,j) + i*VR(:,j+1)
212+
// - v(j+1) = VR(:,j) - i*VR(:,j+1).
213+
//
214+
// ```
215+
// j -> <----pair----> <----pair---->
216+
// [ ... (real), (imag), (imag), (imag), (imag), ... ] : eigs
217+
// ^ ^ ^ ^ ^
218+
// false false true false true : is_conjugate_pair
219+
// ```
95220
let n = n as usize;
96-
let mut flg = false;
97-
let conj: Vec<i8> = wi
98-
.iter()
99-
.map(|&i| {
100-
if flg {
101-
flg = false;
102-
-1
103-
} else if i != 0.0 {
104-
flg = true;
105-
1
106-
} else {
107-
0
221+
let v = vr.or(vl).unwrap();
222+
let mut eigvecs = vec![Self::Complex::zero(); n * n];
223+
let mut is_conjugate_pair = false; // flag for check `j` is complex conjugate
224+
for j in 0..n {
225+
if eig_im[j] == 0.0 {
226+
// j-th eigenvalue is real
227+
for i in 0..n {
228+
eigvecs[i + j * n] = Self::complex(v[i + j * n], 0.0);
108229
}
109-
})
110-
.collect();
111-
let v: Vec<Self::Complex> = (0..n * n)
112-
.map(|i| {
113-
let j = i % n;
114-
match conj[j] {
115-
1 => Self::Complex::new(vr[i], vr[i + 1]),
116-
-1 => Self::Complex::new(vr[i - 1], -vr[i]),
117-
_ => Self::Complex::new(vr[i], 0.0),
230+
} else {
231+
// j-th eigenvalue is complex
232+
// complex conjugated pair can be `j-1` or `j+1`
233+
if is_conjugate_pair {
234+
let j_pair = j - 1;
235+
assert!(j_pair < n);
236+
for i in 0..n {
237+
eigvecs[i + j * n] = Self::complex(v[i + j_pair * n], v[i + j * n]);
238+
}
239+
} else {
240+
let j_pair = j + 1;
241+
assert!(j_pair < n);
242+
for i in 0..n {
243+
eigvecs[i + j * n] = Self::complex(v[i + j * n], v[i + j_pair * n]);
244+
}
118245
}
119-
})
120-
.collect();
246+
is_conjugate_pair = !is_conjugate_pair;
247+
}
248+
}
121249

122-
info.as_lapack_result()?;
123-
Ok((w, v))
250+
Ok((eigs, eigvecs))
124251
}
125252
}
126253
};
127254
}
128255

129-
impl_eig_real!(f64, lapacke::dgeev);
130-
impl_eig_real!(f32, lapacke::sgeev);
131-
impl_eig_complex!(c64, lapacke::zgeev);
132-
impl_eig_complex!(c32, lapacke::cgeev);
256+
impl_eig_real!(f64, lapack::dgeev);
257+
impl_eig_real!(f32, lapack::sgeev);

‎ndarray-linalg/src/eig.rs

+4-6
Original file line numberDiff line numberDiff line change
@@ -48,13 +48,11 @@ where
4848
fn eig(&self) -> Result<(Self::EigVal, Self::EigVec)> {
4949
let mut a = self.to_owned();
5050
let layout = a.square_layout()?;
51-
let (s, t) = unsafe { A::eig(true, layout, a.as_allocated_mut()?)? };
52-
let (n, _) = layout.size();
51+
let (s, t) = A::eig(true, layout, a.as_allocated_mut()?)?;
52+
let n = layout.len() as usize;
5353
Ok((
5454
ArrayBase::from(s),
55-
ArrayBase::from(t)
56-
.into_shape((n as usize, n as usize))
57-
.unwrap(),
55+
Array2::from_shape_vec((n, n).f(), t).unwrap(),
5856
))
5957
}
6058
}
@@ -74,7 +72,7 @@ where
7472

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

0 commit comments

Comments
 (0)