Skip to content

Commit 5f00eff

Browse files
termoshtttermoshtt
committed
WIP: EigWork
1 parent ec958c6 commit 5f00eff

File tree

1 file changed

+210
-0
lines changed

1 file changed

+210
-0
lines changed

‎lax/src/eig.rs

+210
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,17 @@ use num_traits::{ToPrimitive, Zero};
44

55
#[cfg_attr(doc, katexit::katexit)]
66
/// Eigenvalue problem for general matrix
7+
///
8+
/// LAPACK assumes a column-major input. A row-major input can
9+
/// be interpreted as the transpose of a column-major input. So,
10+
/// for row-major inputs, we we want to solve the following,
11+
/// given the column-major input `A`:
12+
///
13+
/// A^T V = V Λ ⟺ V^T A = Λ V^T ⟺ conj(V)^H A = Λ conj(V)^H
14+
///
15+
/// So, in this case, the right eigenvectors are the conjugates
16+
/// of the left eigenvectors computed with `A`, and the
17+
/// eigenvalues are the eigenvalues computed with `A`.
718
pub trait Eig_: Scalar {
819
/// Compute right eigenvalue and eigenvectors $Ax = \lambda x$
920
///
@@ -21,6 +32,205 @@ pub trait Eig_: Scalar {
2132
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
2233
}
2334

35+
/// Working memory for [Eig_]
36+
#[derive(Debug, Clone)]
37+
pub struct EigWork<T: Scalar> {
38+
pub n: i32,
39+
pub jobvr: JobEv,
40+
pub jobvl: JobEv,
41+
42+
/// Eigenvalues used in complex routines
43+
pub eigs: Option<Vec<MaybeUninit<T>>>,
44+
/// Real part of eigenvalues used in real routines
45+
pub eigs_re: Option<Vec<MaybeUninit<T>>>,
46+
/// Imaginary part of eigenvalues used in real routines
47+
pub eigs_im: Option<Vec<MaybeUninit<T>>>,
48+
49+
/// Left eigenvectors
50+
pub vl: Option<Vec<MaybeUninit<T>>>,
51+
/// Right eigenvectors
52+
pub vr: Option<Vec<MaybeUninit<T>>>,
53+
54+
/// Working memory
55+
pub work: Vec<MaybeUninit<T>>,
56+
/// Working memory with `T::Real`
57+
pub rwork: Option<Vec<MaybeUninit<T::Real>>>,
58+
}
59+
60+
impl EigWork<c64> {
61+
pub fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
62+
let (n, _) = l.size();
63+
let (jobvl, jobvr) = if calc_v {
64+
match l {
65+
MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
66+
MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
67+
}
68+
} else {
69+
(JobEv::None, JobEv::None)
70+
};
71+
let mut eigs: Vec<MaybeUninit<c64>> = vec_uninit(n as usize);
72+
let mut rwork: Vec<MaybeUninit<f64>> = vec_uninit(2 * n as usize);
73+
74+
let mut vl: Option<Vec<MaybeUninit<c64>>> = jobvl.then(|| vec_uninit((n * n) as usize));
75+
let mut vr: Option<Vec<MaybeUninit<c64>>> = jobvr.then(|| vec_uninit((n * n) as usize));
76+
77+
// calc work size
78+
let mut info = 0;
79+
let mut work_size = [c64::zero()];
80+
unsafe {
81+
lapack_sys::zgeev_(
82+
jobvl.as_ptr(),
83+
jobvr.as_ptr(),
84+
&n,
85+
std::ptr::null_mut(),
86+
&n,
87+
AsPtr::as_mut_ptr(&mut eigs),
88+
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
89+
&n,
90+
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
91+
&n,
92+
AsPtr::as_mut_ptr(&mut work_size),
93+
&(-1),
94+
AsPtr::as_mut_ptr(&mut rwork),
95+
&mut info,
96+
)
97+
};
98+
info.as_lapack_result()?;
99+
100+
let lwork = work_size[0].to_usize().unwrap();
101+
let work: Vec<MaybeUninit<c64>> = vec_uninit(lwork);
102+
Ok(Self {
103+
n,
104+
jobvl,
105+
jobvr,
106+
eigs: Some(eigs),
107+
eigs_re: None,
108+
eigs_im: None,
109+
rwork: Some(rwork),
110+
vl,
111+
vr,
112+
work,
113+
})
114+
}
115+
116+
/// Compute eigenvalues and vectors on this working memory.
117+
pub fn calc<'work>(
118+
&'work mut self,
119+
a: &mut [c64],
120+
) -> Result<(&'work [c64], Option<&'work [c64]>)> {
121+
let lwork = self.work.len().to_i32().unwrap();
122+
let mut info = 0;
123+
unsafe {
124+
lapack_sys::zgeev_(
125+
self.jobvl.as_ptr(),
126+
self.jobvr.as_ptr(),
127+
&self.n,
128+
AsPtr::as_mut_ptr(a),
129+
&self.n,
130+
AsPtr::as_mut_ptr(self.eigs.as_mut().unwrap()),
131+
AsPtr::as_mut_ptr(
132+
self.vl
133+
.as_mut()
134+
.map(|v| v.as_mut_slice())
135+
.unwrap_or(&mut []),
136+
),
137+
&self.n,
138+
AsPtr::as_mut_ptr(
139+
self.vr
140+
.as_mut()
141+
.map(|v| v.as_mut_slice())
142+
.unwrap_or(&mut []),
143+
),
144+
&self.n,
145+
AsPtr::as_mut_ptr(&mut self.work),
146+
&lwork,
147+
AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()),
148+
&mut info,
149+
)
150+
};
151+
info.as_lapack_result()?;
152+
153+
let eigs = self
154+
.eigs
155+
.as_ref()
156+
.map(|v| unsafe { v.slice_assume_init_ref() })
157+
.unwrap();
158+
159+
// Hermite conjugate
160+
if let Some(vl) = self.vl.as_mut() {
161+
for value in vl {
162+
let value = unsafe { value.assume_init_mut() };
163+
value.im = -value.im;
164+
}
165+
}
166+
let v = match (self.vl.as_ref(), self.vr.as_ref()) {
167+
(Some(v), None) | (None, Some(v)) => Some(unsafe { v.slice_assume_init_ref() }),
168+
(None, None) => None,
169+
_ => unreachable!(),
170+
};
171+
Ok((eigs, v))
172+
}
173+
}
174+
175+
impl EigWork<f64> {
176+
pub fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
177+
let (n, _) = l.size();
178+
let (jobvl, jobvr) = if calc_v {
179+
match l {
180+
MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
181+
MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
182+
}
183+
} else {
184+
(JobEv::None, JobEv::None)
185+
};
186+
let mut eigs_re: Vec<MaybeUninit<f64>> = vec_uninit(n as usize);
187+
let mut eigs_im: Vec<MaybeUninit<f64>> = vec_uninit(n as usize);
188+
189+
let mut vl: Option<Vec<MaybeUninit<f64>>> = jobvl.then(|| vec_uninit((n * n) as usize));
190+
let mut vr: Option<Vec<MaybeUninit<f64>>> = jobvr.then(|| vec_uninit((n * n) as usize));
191+
192+
// calc work size
193+
let mut info = 0;
194+
let mut work_size: [f64; 1] = [0.0];
195+
unsafe {
196+
lapack_sys::dgeev_(
197+
jobvl.as_ptr(),
198+
jobvr.as_ptr(),
199+
&n,
200+
std::ptr::null_mut(),
201+
&n,
202+
AsPtr::as_mut_ptr(&mut eigs_re),
203+
AsPtr::as_mut_ptr(&mut eigs_im),
204+
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
205+
&n,
206+
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
207+
&n,
208+
AsPtr::as_mut_ptr(&mut work_size),
209+
&(-1),
210+
&mut info,
211+
)
212+
};
213+
info.as_lapack_result()?;
214+
215+
// actual ev
216+
let lwork = work_size[0].to_usize().unwrap();
217+
let work: Vec<MaybeUninit<f64>> = vec_uninit(lwork);
218+
219+
Ok(Self {
220+
n,
221+
jobvr,
222+
jobvl,
223+
eigs: None,
224+
eigs_re: Some(eigs_re),
225+
eigs_im: Some(eigs_im),
226+
rwork: None,
227+
vl,
228+
vr,
229+
work,
230+
})
231+
}
232+
}
233+
24234
macro_rules! impl_eig_complex {
25235
($scalar:ty, $ev:path) => {
26236
impl Eig_ for $scalar {

0 commit comments

Comments
 (0)