lax/
eig.rs

1//! Eigenvalue problem for general matricies
2//!
3//! LAPACK correspondance
4//! ----------------------
5//!
6//! | f32   | f64   | c32   | c64   |
7//! |:------|:------|:------|:------|
8//! | sgeev | dgeev | cgeev | zgeev |
9//!
10
11use crate::{error::*, layout::MatrixLayout, *};
12use cauchy::*;
13use num_traits::{ToPrimitive, Zero};
14
15#[cfg_attr(doc, katexit::katexit)]
16/// Eigenvalue problem for general matrix
17///
18/// To manage memory more strictly, use [EigWork].
19///
20/// Right and Left eigenvalue problem
21/// ----------------------------------
22/// LAPACK can solve both right eigenvalue problem
23/// $$
24/// AV_R = V_R \Lambda
25/// $$
26/// where $V_R = \left( v_R^1, \cdots, v_R^n \right)$ are right eigenvectors
27/// and left eigenvalue problem
28/// $$
29/// V_L^\dagger A = V_L^\dagger \Lambda
30/// $$
31/// where $V_L = \left( v_L^1, \cdots, v_L^n \right)$ are left eigenvectors
32/// and eigenvalues
33/// $$
34/// \Lambda = \begin{pmatrix}
35///   \lambda_1 &        & 0 \\\\
36///             & \ddots &   \\\\
37///   0         &        & \lambda_n
38/// \end{pmatrix}
39/// $$
40/// which satisfies $A v_R^i = \lambda_i v_R^i$ and
41/// $\left(v_L^i\right)^\dagger A = \lambda_i \left(v_L^i\right)^\dagger$
42/// for column-major matrices, although row-major matrices are not supported.
43/// Since a row-major matrix can be interpreted
44/// as a transpose of a column-major matrix,
45/// this transforms right eigenvalue problem to left one:
46///
47/// $$
48/// A^\dagger V = V Λ ⟺ V^\dagger A = Λ V^\dagger
49/// $$
50///
51#[non_exhaustive]
52pub struct EigWork<T: Scalar> {
53    /// Problem size
54    pub n: i32,
55    /// Compute right eigenvectors or not
56    pub jobvr: JobEv,
57    /// Compute left eigenvectors or not
58    pub jobvl: JobEv,
59
60    /// Eigenvalues
61    pub eigs: Vec<MaybeUninit<T::Complex>>,
62    /// Real part of eigenvalues used in real routines
63    pub eigs_re: Option<Vec<MaybeUninit<T::Real>>>,
64    /// Imaginary part of eigenvalues used in real routines
65    pub eigs_im: Option<Vec<MaybeUninit<T::Real>>>,
66
67    /// Left eigenvectors
68    pub vc_l: Option<Vec<MaybeUninit<T::Complex>>>,
69    /// Left eigenvectors used in real routines
70    pub vr_l: Option<Vec<MaybeUninit<T::Real>>>,
71    /// Right eigenvectors
72    pub vc_r: Option<Vec<MaybeUninit<T::Complex>>>,
73    /// Right eigenvectors used in real routines
74    pub vr_r: Option<Vec<MaybeUninit<T::Real>>>,
75
76    /// Working memory
77    pub work: Vec<MaybeUninit<T>>,
78    /// Working memory with `T::Real`
79    pub rwork: Option<Vec<MaybeUninit<T::Real>>>,
80}
81
82impl<T> EigWork<T>
83where
84    T: Scalar,
85    EigWork<T>: EigWorkImpl<Elem = T>,
86{
87    /// Create new working memory for eigenvalues compution.
88    pub fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
89        EigWorkImpl::new(calc_v, l)
90    }
91
92    /// Compute eigenvalues and vectors on this working memory.
93    pub fn calc(&mut self, a: &mut [T]) -> Result<EigRef<T>> {
94        EigWorkImpl::calc(self, a)
95    }
96
97    /// Compute eigenvalues and vectors by consuming this working memory.
98    pub fn eval(self, a: &mut [T]) -> Result<EigOwned<T>> {
99        EigWorkImpl::eval(self, a)
100    }
101}
102
103/// Owned result of eigenvalue problem by [EigWork::eval]
104#[derive(Debug, Clone, PartialEq)]
105pub struct EigOwned<T: Scalar> {
106    /// Eigenvalues
107    pub eigs: Vec<T::Complex>,
108    /// Right eigenvectors
109    pub vr: Option<Vec<T::Complex>>,
110    /// Left eigenvectors
111    pub vl: Option<Vec<T::Complex>>,
112}
113
114/// Reference result of eigenvalue problem by [EigWork::calc]
115#[derive(Debug, Clone, PartialEq)]
116pub struct EigRef<'work, T: Scalar> {
117    /// Eigenvalues
118    pub eigs: &'work [T::Complex],
119    /// Right eigenvectors
120    pub vr: Option<&'work [T::Complex]>,
121    /// Left eigenvectors
122    pub vl: Option<&'work [T::Complex]>,
123}
124
125/// Helper trait for implementing [EigWork] methods
126pub trait EigWorkImpl: Sized {
127    type Elem: Scalar;
128    fn new(calc_v: bool, l: MatrixLayout) -> Result<Self>;
129    fn calc<'work>(&'work mut self, a: &mut [Self::Elem]) -> Result<EigRef<'work, Self::Elem>>;
130    fn eval(self, a: &mut [Self::Elem]) -> Result<EigOwned<Self::Elem>>;
131}
132
133macro_rules! impl_eig_work_c {
134    ($c:ty, $ev:path) => {
135        impl EigWorkImpl for EigWork<$c> {
136            type Elem = $c;
137
138            fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
139                let (n, _) = l.size();
140                let (jobvl, jobvr) = if calc_v {
141                    match l {
142                        MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
143                        MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
144                    }
145                } else {
146                    (JobEv::None, JobEv::None)
147                };
148                let mut eigs = vec_uninit(n as usize);
149                let mut rwork = vec_uninit(2 * n as usize);
150
151                let mut vc_l = jobvl.then(|| vec_uninit((n * n) as usize));
152                let mut vc_r = jobvr.then(|| vec_uninit((n * n) as usize));
153
154                // calc work size
155                let mut info = 0;
156                let mut work_size = [<$c>::zero()];
157                unsafe {
158                    $ev(
159                        jobvl.as_ptr(),
160                        jobvr.as_ptr(),
161                        &n,
162                        std::ptr::null_mut(),
163                        &n,
164                        AsPtr::as_mut_ptr(&mut eigs),
165                        AsPtr::as_mut_ptr(vc_l.as_deref_mut().unwrap_or(&mut [])),
166                        &n,
167                        AsPtr::as_mut_ptr(vc_r.as_deref_mut().unwrap_or(&mut [])),
168                        &n,
169                        AsPtr::as_mut_ptr(&mut work_size),
170                        &(-1),
171                        AsPtr::as_mut_ptr(&mut rwork),
172                        &mut info,
173                    )
174                };
175                info.as_lapack_result()?;
176
177                let lwork = work_size[0].to_usize().unwrap();
178                let work: Vec<MaybeUninit<$c>> = vec_uninit(lwork);
179                Ok(Self {
180                    n,
181                    jobvl,
182                    jobvr,
183                    eigs,
184                    eigs_re: None,
185                    eigs_im: None,
186                    rwork: Some(rwork),
187                    vc_l,
188                    vc_r,
189                    vr_l: None,
190                    vr_r: None,
191                    work,
192                })
193            }
194
195            fn calc<'work>(
196                &'work mut self,
197                a: &mut [Self::Elem],
198            ) -> Result<EigRef<'work, Self::Elem>> {
199                let lwork = self.work.len().to_i32().unwrap();
200                let mut info = 0;
201                unsafe {
202                    $ev(
203                        self.jobvl.as_ptr(),
204                        self.jobvr.as_ptr(),
205                        &self.n,
206                        AsPtr::as_mut_ptr(a),
207                        &self.n,
208                        AsPtr::as_mut_ptr(&mut self.eigs),
209                        AsPtr::as_mut_ptr(self.vc_l.as_deref_mut().unwrap_or(&mut [])),
210                        &self.n,
211                        AsPtr::as_mut_ptr(self.vc_r.as_deref_mut().unwrap_or(&mut [])),
212                        &self.n,
213                        AsPtr::as_mut_ptr(&mut self.work),
214                        &lwork,
215                        AsPtr::as_mut_ptr(self.rwork.as_mut().unwrap()),
216                        &mut info,
217                    )
218                };
219                info.as_lapack_result()?;
220                // Hermite conjugate
221                if let Some(vl) = self.vc_l.as_mut() {
222                    for value in vl {
223                        let value = unsafe { value.assume_init_mut() };
224                        value.im = -value.im;
225                    }
226                }
227                Ok(EigRef {
228                    eigs: unsafe { self.eigs.slice_assume_init_ref() },
229                    vl: self
230                        .vc_l
231                        .as_ref()
232                        .map(|v| unsafe { v.slice_assume_init_ref() }),
233                    vr: self
234                        .vc_r
235                        .as_ref()
236                        .map(|v| unsafe { v.slice_assume_init_ref() }),
237                })
238            }
239
240            fn eval(mut self, a: &mut [Self::Elem]) -> Result<EigOwned<Self::Elem>> {
241                let _eig_ref = self.calc(a)?;
242                Ok(EigOwned {
243                    eigs: unsafe { self.eigs.assume_init() },
244                    vl: self.vc_l.map(|v| unsafe { v.assume_init() }),
245                    vr: self.vc_r.map(|v| unsafe { v.assume_init() }),
246                })
247            }
248        }
249    };
250}
251
252impl_eig_work_c!(c32, lapack_sys::cgeev_);
253impl_eig_work_c!(c64, lapack_sys::zgeev_);
254
255macro_rules! impl_eig_work_r {
256    ($f:ty, $ev:path) => {
257        impl EigWorkImpl for EigWork<$f> {
258            type Elem = $f;
259
260            fn new(calc_v: bool, l: MatrixLayout) -> Result<Self> {
261                let (n, _) = l.size();
262                let (jobvl, jobvr) = if calc_v {
263                    match l {
264                        MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
265                        MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
266                    }
267                } else {
268                    (JobEv::None, JobEv::None)
269                };
270                let mut eigs_re = vec_uninit(n as usize);
271                let mut eigs_im = vec_uninit(n as usize);
272                let mut vr_l = jobvl.then(|| vec_uninit((n * n) as usize));
273                let mut vr_r = jobvr.then(|| vec_uninit((n * n) as usize));
274                let vc_l = jobvl.then(|| vec_uninit((n * n) as usize));
275                let vc_r = jobvr.then(|| vec_uninit((n * n) as usize));
276
277                // calc work size
278                let mut info = 0;
279                let mut work_size: [$f; 1] = [0.0];
280                unsafe {
281                    $ev(
282                        jobvl.as_ptr(),
283                        jobvr.as_ptr(),
284                        &n,
285                        std::ptr::null_mut(),
286                        &n,
287                        AsPtr::as_mut_ptr(&mut eigs_re),
288                        AsPtr::as_mut_ptr(&mut eigs_im),
289                        AsPtr::as_mut_ptr(vr_l.as_deref_mut().unwrap_or(&mut [])),
290                        &n,
291                        AsPtr::as_mut_ptr(vr_r.as_deref_mut().unwrap_or(&mut [])),
292                        &n,
293                        AsPtr::as_mut_ptr(&mut work_size),
294                        &(-1),
295                        &mut info,
296                    )
297                };
298                info.as_lapack_result()?;
299
300                // actual ev
301                let lwork = work_size[0].to_usize().unwrap();
302                let work = vec_uninit(lwork);
303
304                Ok(Self {
305                    n,
306                    jobvr,
307                    jobvl,
308                    eigs: vec_uninit(n as usize),
309                    eigs_re: Some(eigs_re),
310                    eigs_im: Some(eigs_im),
311                    rwork: None,
312                    vr_l,
313                    vr_r,
314                    vc_l,
315                    vc_r,
316                    work,
317                })
318            }
319
320            fn calc<'work>(
321                &'work mut self,
322                a: &mut [Self::Elem],
323            ) -> Result<EigRef<'work, Self::Elem>> {
324                let lwork = self.work.len().to_i32().unwrap();
325                let mut info = 0;
326                unsafe {
327                    $ev(
328                        self.jobvl.as_ptr(),
329                        self.jobvr.as_ptr(),
330                        &self.n,
331                        AsPtr::as_mut_ptr(a),
332                        &self.n,
333                        AsPtr::as_mut_ptr(self.eigs_re.as_mut().unwrap()),
334                        AsPtr::as_mut_ptr(self.eigs_im.as_mut().unwrap()),
335                        AsPtr::as_mut_ptr(self.vr_l.as_deref_mut().unwrap_or(&mut [])),
336                        &self.n,
337                        AsPtr::as_mut_ptr(self.vr_r.as_deref_mut().unwrap_or(&mut [])),
338                        &self.n,
339                        AsPtr::as_mut_ptr(&mut self.work),
340                        &lwork,
341                        &mut info,
342                    )
343                };
344                info.as_lapack_result()?;
345
346                let eigs_re = self
347                    .eigs_re
348                    .as_ref()
349                    .map(|e| unsafe { e.slice_assume_init_ref() })
350                    .unwrap();
351                let eigs_im = self
352                    .eigs_im
353                    .as_ref()
354                    .map(|e| unsafe { e.slice_assume_init_ref() })
355                    .unwrap();
356                reconstruct_eigs(eigs_re, eigs_im, &mut self.eigs);
357
358                if let Some(v) = self.vr_l.as_ref() {
359                    let v = unsafe { v.slice_assume_init_ref() };
360                    reconstruct_eigenvectors(true, eigs_im, v, self.vc_l.as_mut().unwrap());
361                }
362                if let Some(v) = self.vr_r.as_ref() {
363                    let v = unsafe { v.slice_assume_init_ref() };
364                    reconstruct_eigenvectors(false, eigs_im, v, self.vc_r.as_mut().unwrap());
365                }
366
367                Ok(EigRef {
368                    eigs: unsafe { self.eigs.slice_assume_init_ref() },
369                    vl: self
370                        .vc_l
371                        .as_ref()
372                        .map(|v| unsafe { v.slice_assume_init_ref() }),
373                    vr: self
374                        .vc_r
375                        .as_ref()
376                        .map(|v| unsafe { v.slice_assume_init_ref() }),
377                })
378            }
379
380            fn eval(mut self, a: &mut [Self::Elem]) -> Result<EigOwned<Self::Elem>> {
381                let _eig_ref = self.calc(a)?;
382                Ok(EigOwned {
383                    eigs: unsafe { self.eigs.assume_init() },
384                    vl: self.vc_l.map(|v| unsafe { v.assume_init() }),
385                    vr: self.vc_r.map(|v| unsafe { v.assume_init() }),
386                })
387            }
388        }
389    };
390}
391impl_eig_work_r!(f32, lapack_sys::sgeev_);
392impl_eig_work_r!(f64, lapack_sys::dgeev_);
393
394/// Reconstruct eigenvectors into complex-array
395///
396/// From LAPACK API https://software.intel.com/en-us/node/469230
397///
398/// - If the j-th eigenvalue is real,
399///   - v(j) = VR(:,j), the j-th column of VR.
400///
401/// - If the j-th and (j+1)-st eigenvalues form a complex conjugate pair,
402///   - v(j)   = VR(:,j) + i*VR(:,j+1)
403///   - v(j+1) = VR(:,j) - i*VR(:,j+1).
404///
405/// In the C-layout case, we need the conjugates of the left
406/// eigenvectors, so the signs should be reversed.
407fn reconstruct_eigenvectors<T: Scalar>(
408    take_hermite_conjugate: bool,
409    eig_im: &[T],
410    vr: &[T],
411    vc: &mut [MaybeUninit<T::Complex>],
412) {
413    let n = eig_im.len();
414    assert_eq!(vr.len(), n * n);
415    assert_eq!(vc.len(), n * n);
416
417    let mut col = 0;
418    while col < n {
419        if eig_im[col].is_zero() {
420            // The corresponding eigenvalue is real.
421            for row in 0..n {
422                let re = vr[row + col * n];
423                vc[row + col * n].write(T::complex(re, T::zero()));
424            }
425            col += 1;
426        } else {
427            // This is a complex conjugate pair.
428            assert!(col + 1 < n);
429            for row in 0..n {
430                let re = vr[row + col * n];
431                let mut im = vr[row + (col + 1) * n];
432                if take_hermite_conjugate {
433                    im = -im;
434                }
435                vc[row + col * n].write(T::complex(re, im));
436                vc[row + (col + 1) * n].write(T::complex(re, -im));
437            }
438            col += 2;
439        }
440    }
441}
442
443/// Create complex eigenvalues from real and imaginary parts.
444fn reconstruct_eigs<T: Scalar>(re: &[T], im: &[T], eigs: &mut [MaybeUninit<T::Complex>]) {
445    let n = eigs.len();
446    assert_eq!(re.len(), n);
447    assert_eq!(im.len(), n);
448    for i in 0..n {
449        eigs[i].write(T::complex(re[i], im[i]));
450    }
451}
lax/eig.rs

lax/
eig.rs
