Skip to main content

mdarray_linalg_lapack/eig/
context.rs

1//! Eigenvalue Decomposition (EIG):
2//!     A * v = λ * v (right eigenvectors)
3//!     u^H * A = λ * u^H (left eigenvectors)
4//! where:
5//!     - A is n × n         (input square matrix)
6//!     - λ are eigenvalues  (can be complex)
7//!     - v are right eigenvectors
8//!     - u are left eigenvectors
9//!
10//! For Hermitian/symmetric matrices (EIGH):
11//!     A * v = λ * v
12//! where:
13//!     - A is n × n Hermitian/symmetric matrix
14//!     - λ are real eigenvalues
15//!     - v are orthonormal eigenvectors
16
17use mdarray::{Array, Dense, Dim, Layout, Shape, Slice};
18use mdarray_linalg::{
19    eig::{Eig, EigDecomp, EigError, EighDecomp, SchurDecomp, SchurError},
20    utils::transpose_in_place,
21};
22use num_complex::{Complex, ComplexFloat};
23use num_traits::identities::Zero;
24
25use super::{
26    scalar::{LapackScalar, NeedsRwork},
27    simple::{gees, gees_complex, geig, geigh},
28};
29use crate::Lapack;
30
31impl<T, D0: Dim, D1: Dim> Eig<T, D0, D1> for Lapack
32where
33    T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
34    Complex<T::Real>: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = Complex<T::Real>>,
35    i8: Into<T::Real>,
36    T::Real: Into<T>,
37{
38    type SpectralScalar = Complex<T::Real>;
39    type RealScalar = T::Real;
40
41    /// Compute eigenvalues and right eigenvectors with new allocated matrices
42    fn eig<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<EigDecomp<Self::SpectralScalar, D0, D1>, EigError>
43    where
44        T: ComplexFloat,
45    {
46        let ash = *a.shape();
47        let (m, n) = (ash.dim(0), ash.dim(1));
48
49        if m != n {
50            return Err(EigError::NotSquareMatrix);
51        }
52
53        let x = T::default();
54        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
55
56        let mut eigenvalues_real = Array::from_elem(ash1, T::default());
57        let mut eigenvalues_imag = Array::from_elem(ash1, T::default());
58        let mut eigenvalues = Array::from_elem(ash1, Complex::new(x.re(), x.re()));
59
60        let mut right_eigenvectors_tmp = Array::from_elem(ash, T::default());
61        let mut right_eigenvectors = Array::from_elem(ash, Complex::new(x.re(), x.re()));
62
63        match geig::<L, Dense, Dense, Dense, Dense, T, D0, D1>(
64            a,
65            &mut eigenvalues_real,
66            &mut eigenvalues_imag,
67            None, // no left eigenvectors
68            Some(&mut right_eigenvectors_tmp),
69        ) {
70            Ok(_) => {
71                for i in 0..n {
72                    eigenvalues[i] = if !eigenvalues_real[i].im().is_zero() {
73                        Complex::new(eigenvalues_real[i].re(), eigenvalues_real[i].im())
74                    } else {
75                        Complex::new(eigenvalues_real[i].re(), eigenvalues_imag[i].re())
76                    }
77                }
78                let mut j = 0_usize;
79                while j < n {
80                    let imag = eigenvalues_imag[[j]];
81                    if imag == T::default() {
82                        for i in 0..n {
83                            let re = right_eigenvectors_tmp[[i, j]];
84                            right_eigenvectors[[i, j]] = Complex::new(re.re(), re.im());
85                        }
86                        j += 1;
87                    } else {
88                        for i in 0..n {
89                            let re = right_eigenvectors_tmp[[i, j]];
90                            let im = right_eigenvectors_tmp[[i, j + 1]];
91                            right_eigenvectors[[i, j]] = Complex::new(re.re(), im.re()); // v = Re + i Im
92                            right_eigenvectors[[i, j + 1]] =
93                                ComplexFloat::conj(Complex::new(re.re(), im.re())); // v̄ = Re - i Im
94                        }
95                        j += 2;
96                    }
97                }
98
99                Ok(EigDecomp {
100                    eigenvalues,
101                    left_eigenvectors: None,
102                    right_eigenvectors: Some(right_eigenvectors),
103                })
104            }
105            Err(e) => Err(e),
106        }
107    }
108
109    /// Compute eigenvalues and both left/right eigenvectors with new allocated matrices
110    fn eig_full<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<EigDecomp<Self::SpectralScalar, D0, D1>, EigError> {
111        let ash = *a.shape();
112        let (m, n) = (ash.dim(0), ash.dim(1));
113
114        if m != n {
115            return Err(EigError::NotSquareMatrix);
116        }
117
118        let x = T::default();
119        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
120
121        let mut eigenvalues_real = Array::from_elem(ash1, T::default());
122        let mut eigenvalues_imag = Array::from_elem(ash1, T::default());
123        let mut eigenvalues = Array::from_elem(ash1, Complex::new(x.re(), x.re()));
124
125        let mut left_eigenvectors_tmp = Array::from_elem(ash, T::default());
126        let mut right_eigenvectors_tmp = Array::from_elem(ash, T::default());
127        let mut left_eigenvectors = Array::from_elem(ash, Complex::new(x.re(), x.re()));
128        let mut right_eigenvectors = Array::from_elem(ash, Complex::new(x.re(), x.re()));
129
130        match geig::<L, Dense, Dense, Dense, Dense, T, D0, D1>(
131            a,
132            &mut eigenvalues_real,
133            &mut eigenvalues_imag,
134            Some(&mut left_eigenvectors_tmp),
135            Some(&mut right_eigenvectors_tmp),
136        ) {
137            Ok(_) => {
138                for i in 0..n {
139                    eigenvalues[i] = if !eigenvalues_real[i].im().is_zero() {
140                        Complex::new(eigenvalues_real[i].re(), eigenvalues_real[i].im())
141                    } else {
142                        Complex::new(eigenvalues_real[i].re(), eigenvalues_imag[i].re())
143                    };
144                }
145
146                let mut j = 0_usize;
147                while j < n {
148                    let imag = eigenvalues_imag[[j]];
149                    if imag == T::default() {
150                        for i in 0..n {
151                            let re_right = right_eigenvectors_tmp[[i, j]];
152                            let re_left = left_eigenvectors_tmp[[i, j]];
153                            right_eigenvectors[[i, j]] =
154                                Complex::new(re_right.re(), re_right.im());
155                            left_eigenvectors[[i, j]] = Complex::new(re_left.re(), re_left.im());
156                        }
157                        j += 1;
158                    } else {
159                        for i in 0..n {
160                            let re_right = right_eigenvectors_tmp[[i, j]];
161                            let im_right = right_eigenvectors_tmp[[i, j + 1]];
162                            let re_left = left_eigenvectors_tmp[[i, j]];
163                            let im_left = left_eigenvectors_tmp[[i, j + 1]];
164
165                            right_eigenvectors[[i, j]] =
166                                Complex::new(re_right.re(), im_right.re());
167                            right_eigenvectors[[i, j + 1]] =
168                                ComplexFloat::conj(Complex::new(re_right.re(), im_right.re()));
169
170                            left_eigenvectors[[i, j]] = Complex::new(re_left.re(), im_left.re());
171                            left_eigenvectors[[i, j + 1]] =
172                                ComplexFloat::conj(Complex::new(re_left.re(), im_left.re()));
173                        }
174                        j += 2;
175                    }
176                }
177
178                Ok(EigDecomp {
179                    eigenvalues,
180                    left_eigenvectors: Some(left_eigenvectors),
181                    right_eigenvectors: Some(right_eigenvectors),
182                })
183            }
184            Err(e) => Err(e),
185        }
186    }
187
188    /// Compute only eigenvalues with new allocated vectors
189    fn eig_values<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<Array<Self::SpectralScalar, (D0,)>, EigError> {
190        let ash = *a.shape();
191        let (m, n) = (ash.dim(0), ash.dim(1));
192
193        if m != n {
194            return Err(EigError::NotSquareMatrix);
195        }
196
197        let x = T::default();
198        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
199
200        let mut eigenvalues_real = Array::from_elem(ash1, T::default());
201        let mut eigenvalues_imag = Array::from_elem(ash1, T::default());
202        let mut eigenvalues = Array::from_elem(ash1, Complex::new(x.re(), x.re()));
203
204        match geig::<L, Dense, Dense, Dense, Dense, T, D0, D1>(
205            a,
206            &mut eigenvalues_real,
207            &mut eigenvalues_imag,
208            None,
209            None,
210        ) {
211            Ok(_) => {
212                for i in 0..n {
213                    eigenvalues[i] =
214                        Complex::new(eigenvalues_real[i].re(), eigenvalues_imag[i].re());
215                }
216
217                Ok(eigenvalues)
218            }
219            Err(e) => Err(e),
220        }
221    }
222
223    /// Compute eigenvalues and eigenvectors of a self-adjoint matrix
224    fn eigh<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<EighDecomp<T, Self::RealScalar, D0, D1>, EigError> {
225        let ash = *a.shape();
226        let (m, n) = (ash.dim(0), ash.dim(1));
227
228        if m != n {
229            return Err(EigError::NotSquareMatrix);
230        }
231
232        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
233        let mut eigenvalues = Array::from_elem(ash1, T::Real::zero());
234        let mut eigenvectors = Array::from_elem(ash, T::default());
235
236        match geigh(a, &mut eigenvalues) {
237            Ok(_) => {
238                for j in 0..n {
239                    for i in 0..n {
240                        eigenvectors[[i, j]] = a[[j, i]];
241                    }
242                }
243
244                Ok(EighDecomp {
245                    eigenvalues,
246                    eigenvectors,
247                })
248            }
249            Err(e) => Err(e),
250        }
251    }
252
253    /// Compute Schur decomposition with new allocated matrices
254    fn schur<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<SchurDecomp<T, D0, D1>, SchurError> {
255        let ash = *a.shape();
256        let (m, n) = (ash.dim(0), ash.dim(1));
257
258        if m != n {
259            return Err(SchurError::NotSquareMatrix);
260        }
261
262        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
263
264        let mut eigenvalues_real = Array::from_elem(ash1, T::default());
265        let mut eigenvalues_imag = Array::from_elem(ash1, T::default());
266        let mut schur_vectors = Array::from_elem(ash, T::default());
267
268        match gees::<L, Dense, Dense, Dense, T, D0, D1>(
269            a,
270            &mut eigenvalues_real,
271            &mut eigenvalues_imag,
272            &mut schur_vectors,
273        ) {
274            Ok(_) => {
275                let mut t = Array::from_elem(ash, T::default());
276                for j in 0..n {
277                    for i in 0..n {
278                        t[[i, j]] = a[[j, i]];
279                    }
280                }
281
282                transpose_in_place(&mut schur_vectors);
283
284                Ok(SchurDecomp {
285                    t,
286                    z: schur_vectors,
287                })
288            }
289            Err(e) => Err(e),
290        }
291    }
292
293    /// Compute Schur decomposition overwriting existing matrices
294    fn schur_write<L: Layout>(
295        &self,
296        a: &mut Slice<T, (D0, D1), L>,
297        t: &mut Slice<T, (D0, D1), Dense>,
298        z: &mut Slice<T, (D0, D1), Dense>,
299    ) -> Result<(), SchurError> {
300        let ash = *a.shape();
301        let (m, n) = (ash.dim(0), ash.dim(1));
302
303        if m != n {
304            return Err(SchurError::NotSquareMatrix);
305        }
306
307        for j in 0..n {
308            for i in 0..n {
309                t[[i, j]] = a[[i, j]];
310            }
311        }
312
313        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
314        let mut eigenvalues_real = Array::from_elem(ash1, T::default());
315        let mut eigenvalues_imag = Array::from_elem(ash1, T::default());
316
317        let result = gees::<Dense, Dense, Dense, Dense, T, D0, D1>(
318            t,
319            &mut eigenvalues_real,
320            &mut eigenvalues_imag,
321            z,
322        );
323        transpose_in_place(z);
324        transpose_in_place(t);
325        result
326    }
327
328    /// Compute Schur (complex) decomposition with new allocated matrices
329    fn schur_complex<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<SchurDecomp<Self::SpectralScalar, D0, D1>, SchurError> {
330        let ash = *a.shape();
331        let (m, n) = (ash.dim(0), ash.dim(1));
332
333        if m != n {
334            return Err(SchurError::NotSquareMatrix);
335        }
336
337        let zero = T::Real::zero();
338        let ash1 = <(D0,) as Shape>::from_dims(&[n]);
339        let mut eigenvalues = Array::from_elem(ash1, Complex::new(zero, zero));
340        let mut a_complex = Array::from_fn(ash, |idx| {
341            let x = a[idx];
342            Complex::new(x.re(), x.im())
343        });
344        let mut schur_vectors = Array::from_elem(ash, Complex::new(zero, zero));
345
346        match gees_complex::<Dense, Dense, Dense, Self::SpectralScalar, D0, D1>(
347            &mut a_complex,
348            &mut eigenvalues,
349            &mut schur_vectors,
350        ) {
351            Ok(_) => {
352                let mut t = Array::from_elem(ash, Complex::new(zero, zero));
353                for j in 0..n {
354                    for i in 0..n {
355                        t[[i, j]] = a_complex[[j, i]];
356                    }
357                }
358
359                transpose_in_place(&mut schur_vectors);
360
361                Ok(SchurDecomp {
362                    t,
363                    z: schur_vectors,
364                })
365            }
366            Err(e) => Err(e),
367        }
368    }
369
370    /// Compute Schur (complex) decomposition overwriting existing matrices
371    fn schur_complex_write<L: Layout>(
372        &self,
373        a: &mut Slice<T, (D0, D1), L>,
374        t: &mut Slice<Self::SpectralScalar, (D0, D1), Dense>,
375        z: &mut Slice<Self::SpectralScalar, (D0, D1), Dense>,
376    ) -> Result<(), SchurError> {
377        let SchurDecomp { t: t_result, z: z_result } = self.schur_complex(a)?;
378        for (dst, src) in t.iter_mut().zip(t_result.iter()) {
379            *dst = *src;
380        }
381        for (dst, src) in z.iter_mut().zip(z_result.iter()) {
382            *dst = *src;
383        }
384        Ok(())
385    }
386}