1use 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 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, 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()); right_eigenvectors[[i, j + 1]] =
93 ComplexFloat::conj(Complex::new(re.re(), im.re())); }
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 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 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 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 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 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 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 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}