1use super::simple::{gees, gees_complex, geig, geigh};
18use mdarray_linalg::{get_dims, into_i32, transpose_in_place};
19
20use mdarray::{DSlice, Dense, Layout, tensor};
21
22use super::scalar::{LapackScalar, NeedsRwork};
23use mdarray_linalg::eig::{
24 Eig, EigDecomp, EigError, EigResult, SchurDecomp, SchurError, SchurResult,
25};
26use num_complex::{Complex, ComplexFloat};
27use num_traits::identities::Zero;
28
29use crate::Lapack;
30
31impl<T> Eig<T> for Lapack
32where
33 T: ComplexFloat + Default + LapackScalar + NeedsRwork<Elem = T>,
34 i8: Into<T::Real>,
35 T::Real: Into<T>,
36{
37 fn eig<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> EigResult<T>
39 where
40 T: ComplexFloat,
41 {
42 let (m, n) = get_dims!(a);
43 if m != n {
44 return Err(EigError::NotSquareMatrix);
45 }
46
47 let x = T::default();
48
49 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
50 let mut eigenvalues_imag = tensor![[T::default(); n as usize]; 1];
51 let mut eigenvalues = tensor![[Complex::new(x.re(), x.re()); n as usize]; 1];
52
53 let mut right_eigenvectors_tmp = tensor![[T::default();n as usize]; n as usize];
54 let mut right_eigenvectors = tensor![[Complex::new(x.re(), x.re());n as usize]; n as usize];
55
56 match geig::<L, Dense, Dense, Dense, Dense, T>(
57 a,
58 &mut eigenvalues_real,
59 &mut eigenvalues_imag,
60 None, Some(&mut right_eigenvectors_tmp),
62 ) {
63 Ok(_) => {
64 for i in 0..(n as usize) {
65 eigenvalues[[0, i]] = if !eigenvalues_real[[0, i]].im().is_zero() {
66 Complex::new(eigenvalues_real[[0, i]].re(), eigenvalues_real[[0, i]].im())
67 } else {
68 Complex::new(eigenvalues_real[[0, i]].re(), eigenvalues_imag[[0, i]].re())
69 }
70 }
71 let mut j = 0_usize;
72 while j < n as usize {
73 let imag = eigenvalues_imag[[0, j]];
74 if imag == T::default() {
75 for i in 0..(n as usize) {
76 let re = right_eigenvectors_tmp[[i, j]];
77 right_eigenvectors[[i, j]] = Complex::new(re.re(), re.im());
78 }
79 j += 1;
80 } else {
81 for i in 0..(n as usize) {
82 let re = right_eigenvectors_tmp[[i, j]];
83 let im = right_eigenvectors_tmp[[i, j + 1]];
84 right_eigenvectors[[i, j]] = Complex::new(re.re(), im.re()); right_eigenvectors[[i, j + 1]] =
86 ComplexFloat::conj(Complex::new(re.re(), im.re())); }
88 j += 2;
89 }
90 }
91
92 Ok(EigDecomp {
93 eigenvalues,
94 left_eigenvectors: None,
95 right_eigenvectors: Some(right_eigenvectors),
96 })
97 }
98 Err(e) => Err(e),
99 }
100 }
101
102 fn eig_full<L: Layout>(&self, _a: &mut DSlice<T, 2, L>) -> EigResult<T> {
104 todo!(); }
179
180 fn eig_values<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> EigResult<T> {
182 let (m, n) = get_dims!(a);
183 if m != n {
184 return Err(EigError::NotSquareMatrix);
185 }
186 let x = T::default();
187
188 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
189 let mut eigenvalues_imag = tensor![[T::default(); n as usize]; 1];
190 let mut eigenvalues = tensor![[Complex::new(x.re(), x.re()); n as usize]; 1];
191
192 match geig::<L, Dense, Dense, Dense, Dense, T>(
193 a,
194 &mut eigenvalues_real,
195 &mut eigenvalues_imag,
196 None,
197 None,
198 ) {
199 Ok(_) => {
200 for i in 0..(n as usize) {
201 eigenvalues[[0, i]] =
202 Complex::new(eigenvalues_real[[0, i]].re(), eigenvalues_imag[[0, i]].re());
203 }
204
205 Ok(EigDecomp {
206 eigenvalues,
207 left_eigenvectors: None,
208 right_eigenvectors: None,
209 })
210 }
211 Err(e) => Err(e),
212 }
213 }
214
215 fn eigh<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> EigResult<T> {
217 let (m, n) = get_dims!(a);
218 if m != n {
219 return Err(EigError::NotSquareMatrix);
220 }
221
222 let x = T::default();
223
224 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
225 let mut eigenvalues = tensor![[Complex::new(x.re(), x.re()); n as usize]; 1];
226 let mut right_eigenvectors_tmp = tensor![[T::default(); n as usize]; n as usize];
229
230 let x = T::default();
231 let mut right_eigenvectors =
232 tensor![[Complex::new(x.re(), x.re()); n as usize]; n as usize];
233
234 match geigh(a, &mut eigenvalues_real, &mut right_eigenvectors_tmp) {
235 Ok(_) => {
236 println!("{}", n / 2);
237 for i in 0..((n / 2 + 1) as usize) {
238 eigenvalues[[0, 2 * i]] = Complex::new(eigenvalues_real[[0, i]].re(), x.re());
239 if 2 * i + 1 < n as usize {
240 eigenvalues[[0, 2 * i + 1]] =
241 Complex::new(eigenvalues_real[[0, i]].im(), x.re());
242 }
243 }
244
245 for j in 0..(n as usize) {
246 for i in 0..(n as usize) {
247 let re = a[[j, i]];
248 right_eigenvectors[[i, j]] = Complex::new(re.re(), re.im());
249 }
250 }
251
252 Ok(EigDecomp {
253 eigenvalues,
254 left_eigenvectors: None,
255 right_eigenvectors: Some(right_eigenvectors),
256 })
257 }
258 Err(e) => Err(e),
259 }
260 }
261
262 fn eigs<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> EigResult<T> {
264 let (m, n) = get_dims!(a);
265 if m != n {
266 return Err(EigError::NotSquareMatrix);
267 }
268
269 let x = T::default();
270
271 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
272 let mut eigenvalues = tensor![[Complex::new(x.re(), x.re()); n as usize]; 1];
273 let mut right_eigenvectors_tmp = tensor![[T::default(); n as usize]; n as usize];
276
277 let x = T::default();
278 let mut right_eigenvectors =
279 tensor![[Complex::new(x.re(), x.re()); n as usize]; n as usize];
280
281 match geigh(a, &mut eigenvalues_real, &mut right_eigenvectors_tmp) {
282 Ok(_) => {
283 for i in 0..n as usize {
284 eigenvalues[[0, i]] = Complex::new(eigenvalues_real[[0, i]].re(), x.re());
285 }
286
287 for j in 0..(n as usize) {
288 for i in 0..(n as usize) {
289 let re = a[[j, i]];
290 right_eigenvectors[[i, j]] = Complex::new(re.re(), re.im());
291 }
292 }
293
294 Ok(EigDecomp {
295 eigenvalues,
296 left_eigenvectors: None,
297 right_eigenvectors: Some(right_eigenvectors),
298 })
299 }
300 Err(e) => Err(e),
301 }
302 }
303
304 fn schur<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> SchurResult<T> {
306 let (m, n) = get_dims!(a);
307 if m != n {
308 return Err(SchurError::NotSquareMatrix);
309 }
310
311 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
312 let mut eigenvalues_imag = tensor![[T::default(); n as usize]; 1];
313 let mut schur_vectors = tensor![[T::default(); n as usize]; n as usize];
314
315 match gees::<L, Dense, Dense, Dense, T>(
316 a,
317 &mut eigenvalues_real,
318 &mut eigenvalues_imag,
319 &mut schur_vectors,
320 ) {
321 Ok(_) => {
322 let mut t = tensor![[T::default(); n as usize]; n as usize];
323 for j in 0..(n as usize) {
324 for i in 0..(n as usize) {
325 t[[i, j]] = a[[j, i]];
326 }
327 }
328
329 transpose_in_place(&mut schur_vectors);
330
331 Ok(SchurDecomp {
332 t,
333 z: schur_vectors,
334 })
335 }
336 Err(e) => Err(e),
337 }
338 }
339
340 fn schur_overwrite<L: Layout>(
342 &self,
343 a: &mut DSlice<T, 2, L>,
344 t: &mut DSlice<T, 2, Dense>,
345 z: &mut DSlice<T, 2, Dense>,
346 ) -> Result<(), SchurError> {
347 let (m, n) = get_dims!(a);
348 if m != n {
349 return Err(SchurError::NotSquareMatrix);
350 }
351
352 for j in 0..(n as usize) {
353 for i in 0..(n as usize) {
354 t[[i, j]] = a[[i, j]];
355 }
356 }
357
358 let mut eigenvalues_real = tensor![[T::default(); n as usize]; 1];
359 let mut eigenvalues_imag = tensor![[T::default(); n as usize]; 1];
360
361 let result = gees::<Dense, Dense, Dense, Dense, T>(
362 t,
363 &mut eigenvalues_real,
364 &mut eigenvalues_imag,
365 z,
366 );
367 transpose_in_place(z);
368 transpose_in_place(t);
369 result
370 }
371
372 fn schur_complex<L: Layout>(&self, a: &mut DSlice<T, 2, L>) -> SchurResult<T> {
374 let (m, n) = get_dims!(a);
375 if m != n {
376 return Err(SchurError::NotSquareMatrix);
377 }
378
379 let mut eigenvalues = tensor![T::default(); n as usize];
380 let mut schur_vectors = tensor![[T::default(); n as usize]; n as usize];
381
382 match gees_complex::<L, Dense, Dense, T>(a, &mut eigenvalues, &mut schur_vectors) {
383 Ok(_) => {
384 let mut t = tensor![[T::default(); n as usize]; n as usize];
385 for j in 0..(n as usize) {
386 for i in 0..(n as usize) {
387 t[[i, j]] = a[[j, i]];
388 }
389 }
390
391 transpose_in_place(&mut schur_vectors);
392
393 Ok(SchurDecomp {
394 t,
395 z: schur_vectors,
396 })
397 }
398 Err(e) => Err(e),
399 }
400 }
401
402 fn schur_complex_overwrite<L: Layout>(
404 &self,
405 a: &mut DSlice<T, 2, L>,
406 t: &mut DSlice<T, 2, Dense>,
407 z: &mut DSlice<T, 2, Dense>,
408 ) -> Result<(), SchurError> {
409 let (m, n) = get_dims!(a);
410 if m != n {
411 return Err(SchurError::NotSquareMatrix);
412 }
413
414 for j in 0..(n as usize) {
415 for i in 0..(n as usize) {
416 t[[i, j]] = a[[i, j]];
417 }
418 }
419
420 let mut eigenvalues = tensor![T::default(); n as usize];
421 let result = gees_complex::<Dense, Dense, Dense, T>(t, &mut eigenvalues, z);
422
423 transpose_in_place(z);
424 transpose_in_place(t);
425
426 result
427 }
428}