Skip to content

Commit 0d7ca20

Browse files
termoshtttermoshtt
committed
Use EigWork in Eig_ impl
1 parent 1c8b639 commit 0d7ca20

File tree

1 file changed

+20
-222
lines changed

1 file changed

+20
-222
lines changed

‎lax/src/eig.rs‎

Lines changed: 20 additions & 222 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,26 @@ pub trait Eig_: Scalar {
3232
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)>;
3333
}
3434

35+
macro_rules! impl_eig {
36+
($s:ty) => {
37+
impl Eig_ for $s {
38+
fn eig(
39+
calc_v: bool,
40+
l: MatrixLayout,
41+
a: &mut [Self],
42+
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
43+
let work = EigWork::<$s>::new(calc_v, l)?;
44+
let Eig { eigs, vr, vl } = work.eval(a)?;
45+
Ok((eigs, vr.or(vl).unwrap_or_default()))
46+
}
47+
}
48+
};
49+
}
50+
impl_eig!(c64);
51+
impl_eig!(c32);
52+
impl_eig!(f64);
53+
impl_eig!(f32);
54+
3555
/// Working memory for [Eig_]
3656
#[derive(Debug, Clone)]
3757
pub struct EigWork<T: Scalar> {
@@ -413,228 +433,6 @@ macro_rules! impl_eig_work_r {
413433
impl_eig_work_r!(f32, lapack_sys::sgeev_);
414434
impl_eig_work_r!(f64, lapack_sys::dgeev_);
415435

416-
macro_rules! impl_eig_complex {
417-
($scalar:ty, $ev:path) => {
418-
impl Eig_ for $scalar {
419-
fn eig(
420-
calc_v: bool,
421-
l: MatrixLayout,
422-
a: &mut [Self],
423-
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
424-
let (n, _) = l.size();
425-
// LAPACK assumes a column-major input. A row-major input can
426-
// be interpreted as the transpose of a column-major input. So,
427-
// for row-major inputs, we we want to solve the following,
428-
// given the column-major input `A`:
429-
//
430-
// A^T V = V Λ ⟺ V^T A = Λ V^T ⟺ conj(V)^H A = Λ conj(V)^H
431-
//
432-
// So, in this case, the right eigenvectors are the conjugates
433-
// of the left eigenvectors computed with `A`, and the
434-
// eigenvalues are the eigenvalues computed with `A`.
435-
let (jobvl, jobvr) = if calc_v {
436-
match l {
437-
MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
438-
MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
439-
}
440-
} else {
441-
(JobEv::None, JobEv::None)
442-
};
443-
let mut eigs: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
444-
let mut rwork: Vec<MaybeUninit<Self::Real>> = vec_uninit(2 * n as usize);
445-
446-
let mut vl: Option<Vec<MaybeUninit<Self>>> =
447-
jobvl.then(|| vec_uninit((n * n) as usize));
448-
let mut vr: Option<Vec<MaybeUninit<Self>>> =
449-
jobvr.then(|| vec_uninit((n * n) as usize));
450-
451-
// calc work size
452-
let mut info = 0;
453-
let mut work_size = [Self::zero()];
454-
unsafe {
455-
$ev(
456-
jobvl.as_ptr(),
457-
jobvr.as_ptr(),
458-
&n,
459-
AsPtr::as_mut_ptr(a),
460-
&n,
461-
AsPtr::as_mut_ptr(&mut eigs),
462-
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
463-
&n,
464-
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
465-
&n,
466-
AsPtr::as_mut_ptr(&mut work_size),
467-
&(-1),
468-
AsPtr::as_mut_ptr(&mut rwork),
469-
&mut info,
470-
)
471-
};
472-
info.as_lapack_result()?;
473-
474-
// actal ev
475-
let lwork = work_size[0].to_usize().unwrap();
476-
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
477-
let lwork = lwork as i32;
478-
unsafe {
479-
$ev(
480-
jobvl.as_ptr(),
481-
jobvr.as_ptr(),
482-
&n,
483-
AsPtr::as_mut_ptr(a),
484-
&n,
485-
AsPtr::as_mut_ptr(&mut eigs),
486-
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
487-
&n,
488-
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
489-
&n,
490-
AsPtr::as_mut_ptr(&mut work),
491-
&lwork,
492-
AsPtr::as_mut_ptr(&mut rwork),
493-
&mut info,
494-
)
495-
};
496-
info.as_lapack_result()?;
497-
498-
let eigs = unsafe { eigs.assume_init() };
499-
let vr = unsafe { vr.map(|v| v.assume_init()) };
500-
let mut vl = unsafe { vl.map(|v| v.assume_init()) };
501-
502-
// Hermite conjugate
503-
if jobvl.is_calc() {
504-
for c in vl.as_mut().unwrap().iter_mut() {
505-
c.im = -c.im;
506-
}
507-
}
508-
509-
Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
510-
}
511-
}
512-
};
513-
}
514-
515-
impl_eig_complex!(c64, lapack_sys::zgeev_);
516-
impl_eig_complex!(c32, lapack_sys::cgeev_);
517-
518-
macro_rules! impl_eig_real {
519-
($scalar:ty, $ev:path) => {
520-
impl Eig_ for $scalar {
521-
fn eig(
522-
calc_v: bool,
523-
l: MatrixLayout,
524-
a: &mut [Self],
525-
) -> Result<(Vec<Self::Complex>, Vec<Self::Complex>)> {
526-
let (n, _) = l.size();
527-
// LAPACK assumes a column-major input. A row-major input can
528-
// be interpreted as the transpose of a column-major input. So,
529-
// for row-major inputs, we we want to solve the following,
530-
// given the column-major input `A`:
531-
//
532-
// A^T V = V Λ ⟺ V^T A = Λ V^T ⟺ conj(V)^H A = Λ conj(V)^H
533-
//
534-
// So, in this case, the right eigenvectors are the conjugates
535-
// of the left eigenvectors computed with `A`, and the
536-
// eigenvalues are the eigenvalues computed with `A`.
537-
//
538-
// We could conjugate the eigenvalues instead of the
539-
// eigenvectors, but we have to reconstruct the eigenvectors
540-
// into new matrices anyway, and by not modifying the
541-
// eigenvalues, we preserve the nice ordering specified by
542-
// `sgeev`/`dgeev`.
543-
let (jobvl, jobvr) = if calc_v {
544-
match l {
545-
MatrixLayout::C { .. } => (JobEv::All, JobEv::None),
546-
MatrixLayout::F { .. } => (JobEv::None, JobEv::All),
547-
}
548-
} else {
549-
(JobEv::None, JobEv::None)
550-
};
551-
let mut eig_re: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
552-
let mut eig_im: Vec<MaybeUninit<Self>> = vec_uninit(n as usize);
553-
554-
let mut vl: Option<Vec<MaybeUninit<Self>>> =
555-
jobvl.then(|| vec_uninit((n * n) as usize));
556-
let mut vr: Option<Vec<MaybeUninit<Self>>> =
557-
jobvr.then(|| vec_uninit((n * n) as usize));
558-
559-
// calc work size
560-
let mut info = 0;
561-
let mut work_size: [Self; 1] = [0.0];
562-
unsafe {
563-
$ev(
564-
jobvl.as_ptr(),
565-
jobvr.as_ptr(),
566-
&n,
567-
AsPtr::as_mut_ptr(a),
568-
&n,
569-
AsPtr::as_mut_ptr(&mut eig_re),
570-
AsPtr::as_mut_ptr(&mut eig_im),
571-
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
572-
&n,
573-
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
574-
&n,
575-
AsPtr::as_mut_ptr(&mut work_size),
576-
&(-1),
577-
&mut info,
578-
)
579-
};
580-
info.as_lapack_result()?;
581-
582-
// actual ev
583-
let lwork = work_size[0].to_usize().unwrap();
584-
let mut work: Vec<MaybeUninit<Self>> = vec_uninit(lwork);
585-
let lwork = lwork as i32;
586-
unsafe {
587-
$ev(
588-
jobvl.as_ptr(),
589-
jobvr.as_ptr(),
590-
&n,
591-
AsPtr::as_mut_ptr(a),
592-
&n,
593-
AsPtr::as_mut_ptr(&mut eig_re),
594-
AsPtr::as_mut_ptr(&mut eig_im),
595-
AsPtr::as_mut_ptr(vl.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
596-
&n,
597-
AsPtr::as_mut_ptr(vr.as_mut().map(|v| v.as_mut_slice()).unwrap_or(&mut [])),
598-
&n,
599-
AsPtr::as_mut_ptr(&mut work),
600-
&lwork,
601-
&mut info,
602-
)
603-
};
604-
info.as_lapack_result()?;
605-
606-
let eig_re = unsafe { eig_re.assume_init() };
607-
let eig_im = unsafe { eig_im.assume_init() };
608-
let vl = unsafe { vl.map(|v| v.assume_init()) };
609-
let vr = unsafe { vr.map(|v| v.assume_init()) };
610-
611-
// reconstruct eigenvalues
612-
let eigs: Vec<Self::Complex> = eig_re
613-
.iter()
614-
.zip(eig_im.iter())
615-
.map(|(&re, &im)| Self::complex(re, im))
616-
.collect();
617-
618-
if calc_v {
619-
let mut eigvecs = vec_uninit((n * n) as usize);
620-
reconstruct_eigenvectors(
621-
jobvl.is_calc(),
622-
&eig_im,
623-
&vr.or(vl).unwrap(),
624-
&mut eigvecs,
625-
);
626-
Ok((eigs, unsafe { eigvecs.assume_init() }))
627-
} else {
628-
Ok((eigs, Vec::new()))
629-
}
630-
}
631-
}
632-
};
633-
}
634-
635-
impl_eig_real!(f64, lapack_sys::dgeev_);
636-
impl_eig_real!(f32, lapack_sys::sgeev_);
637-
638436
/// Reconstruct eigenvectors into complex-array
639437
///
640438
/// From LAPACK API https://software.intel.com/en-us/node/469230

0 commit comments

Comments
 (0)