Skip to content

Commit 184de06

Browse files
termoshtttermoshtt
committed
WIP: EigWork
1 parent ec958c6 commit 184de06

File tree

1 file changed

+206
-0
lines changed

1 file changed

+206
-0
lines changed

‎lax/src/eig.rs

Lines changed: 206 additions & 0 deletions
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,201 @@ 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_deref_mut().unwrap_or(&mut [])),
89+
&n,
90+
AsPtr::as_mut_ptr(vr.as_deref_mut().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.as_deref_mut()
133+
.unwrap_or(&mut []),
134+
),
135+
&self.n,
136+
AsPtr::as_mut_ptr(
137+
self.vr.as_deref_mut()
138+
.unwrap_or(&mut []),
139+
),
140+
&self.n,
141+
AsPtr::as_mut_ptr(&mut self.work),
142+
&lwork,
143+
AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()),
144+
&mut info,
145+
)
146+
};
147+
info.as_lapack_result()?;
148+
149+
let eigs = self
150+
.eigs
151+
.as_ref()
152+
.map(|v| unsafe { v.slice_assume_init_ref() })
153+
.unwrap();
154+
155+
// Hermite conjugate
156+
if let Some(vl) = self.vl.as_mut() {
157+
for value in vl {
158+
let value = unsafe { value.assume_init_mut() };
159+
value.im = -value.im;
160+
}
161+
}
162+
let v = match (self.vl.as_ref(), self.vr.as_ref()) {
163+
(Some(v), None) | (None, Some(v)) => Some(unsafe { v.slice_assume_init_ref() }),
164+
(None, None) => None,
165+
_ => unreachable!(),
166+
};
167+
Ok((eigs, v))
168+
}
169+
}
170+
171+
impl EigWork<f64> {
172+
pub fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
173+
let (n, _) = l.size();
174+
let (jobvl, jobvr) = if calc_v {
175+
match l {
176+
MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
177+
MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
178+
}
179+
} else {
180+
(JobEv::None, JobEv::None)
181+
};
182+
let mut eigs_re: Vec<MaybeUninit<f64>> = vec_uninit(n as usize);
183+
let mut eigs_im: Vec<MaybeUninit<f64>> = vec_uninit(n as usize);
184+
185+
let mut vl: Option<Vec<MaybeUninit<f64>>> = jobvl.then(|| vec_uninit((n * n) as usize));
186+
let mut vr: Option<Vec<MaybeUninit<f64>>> = jobvr.then(|| vec_uninit((n * n) as usize));
187+
188+
// calc work size
189+
let mut info = 0;
190+
let mut work_size: [f64; 1] = [0.0];
191+
unsafe {
192+
lapack_sys::dgeev_(
193+
jobvl.as_ptr(),
194+
jobvr.as_ptr(),
195+
&n,
196+
std::ptr::null_mut(),
197+
&n,
198+
AsPtr::as_mut_ptr(&mut eigs_re),
199+
AsPtr::as_mut_ptr(&mut eigs_im),
200+
AsPtr::as_mut_ptr(vl.as_deref_mut().unwrap_or(&mut [])),
201+
&n,
202+
AsPtr::as_mut_ptr(vr.as_deref_mut().unwrap_or(&mut [])),
203+
&n,
204+
AsPtr::as_mut_ptr(&mut work_size),
205+
&(-1),
206+
&mut info,
207+
)
208+
};
209+
info.as_lapack_result()?;
210+
211+
// actual ev
212+
let lwork = work_size[0].to_usize().unwrap();
213+
let work: Vec<MaybeUninit<f64>> = vec_uninit(lwork);
214+
215+
Ok(Self {
216+
n,
217+
jobvr,
218+
jobvl,
219+
eigs: None,
220+
eigs_re: Some(eigs_re),
221+
eigs_im: Some(eigs_im),
222+
rwork: None,
223+
vl,
224+
vr,
225+
work,
226+
})
227+
}
228+
}
229+
24230
macro_rules! impl_eig_complex {
25231
($scalar:ty, $ev:path) => {
26232
impl Eig_ for $scalar {

0 commit comments

Comments
 (0)