ndarray_linalg/
eigh.rs

1//! Eigendecomposition for Hermitian matrices.
2//!
3//! For a Hermitian matrix `A`, this solves the eigenvalue problem `A V = V D`
4//! for `D` and `V`, where `D` is the diagonal matrix of eigenvalues in
5//! ascending order and `V` is the orthonormal matrix of corresponding
6//! eigenvectors.
7//!
8//! For a pair of Hermitian matrices `A` and `B` where `B` is also positive
9//! definite, this solves the generalized eigenvalue problem `A V = B V D`,
10//! where `D` is the diagonal matrix of generalized eigenvalues in ascending
11//! order and `V` is the matrix of corresponding generalized eigenvectors. The
12//! matrix `V` is normalized such that `V^H B V = I`.
13//!
14//! # Example
15//!
16//! Find the eigendecomposition of a Hermitian (or real symmetric) matrix.
17//!
18//! ```
19//! use approx::assert_abs_diff_eq;
20//! use ndarray::{array, Array2};
21//! use ndarray_linalg::{Eigh, UPLO};
22//!
23//! let a: Array2<f64> = array![
24//!     [2., 1.],
25//!     [1., 2.],
26//! ];
27//! let (eigvals, eigvecs) = a.eigh(UPLO::Lower)?;
28//! assert_abs_diff_eq!(eigvals, array![1., 3.]);
29//! assert_abs_diff_eq!(
30//!     a.dot(&eigvecs),
31//!     eigvecs.dot(&Array2::from_diag(&eigvals)),
32//! );
33//! # Ok::<(), Box<dyn std::error::Error>>(())
34//! ```
35
36use ndarray::*;
37
38use crate::diagonal::*;
39use crate::error::*;
40use crate::layout::*;
41use crate::operator::LinearOperator;
42use crate::types::*;
43use crate::UPLO;
44
45/// Eigenvalue decomposition of Hermite matrix reference
46pub trait Eigh {
47    type EigVal;
48    type EigVec;
49    fn eigh(&self, uplo: UPLO) -> Result<(Self::EigVal, Self::EigVec)>;
50}
51
52/// Eigenvalue decomposition of mutable reference of Hermite matrix
53pub trait EighInplace {
54    type EigVal;
55    fn eigh_inplace(&mut self, uplo: UPLO) -> Result<(Self::EigVal, &mut Self)>;
56}
57
58/// Eigenvalue decomposition of Hermite matrix
59pub trait EighInto: Sized {
60    type EigVal;
61    fn eigh_into(self, uplo: UPLO) -> Result<(Self::EigVal, Self)>;
62}
63
64impl<A, S> EighInto for ArrayBase<S, Ix2>
65where
66    A: Scalar + Lapack,
67    S: DataMut<Elem = A>,
68{
69    type EigVal = Array1<A::Real>;
70
71    fn eigh_into(mut self, uplo: UPLO) -> Result<(Self::EigVal, Self)> {
72        let (val, _) = self.eigh_inplace(uplo)?;
73        Ok((val, self))
74    }
75}
76
77impl<A, S, S2> EighInto for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
78where
79    A: Scalar + Lapack,
80    S: DataMut<Elem = A>,
81    S2: DataMut<Elem = A>,
82{
83    type EigVal = Array1<A::Real>;
84
85    fn eigh_into(mut self, uplo: UPLO) -> Result<(Self::EigVal, Self)> {
86        let (val, _) = self.eigh_inplace(uplo)?;
87        Ok((val, self))
88    }
89}
90
91impl<A, S> Eigh for ArrayBase<S, Ix2>
92where
93    A: Scalar + Lapack,
94    S: Data<Elem = A>,
95{
96    type EigVal = Array1<A::Real>;
97    type EigVec = Array2<A>;
98
99    fn eigh(&self, uplo: UPLO) -> Result<(Self::EigVal, Self::EigVec)> {
100        let a = self.to_owned();
101        a.eigh_into(uplo)
102    }
103}
104
105impl<A, S, S2> Eigh for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
106where
107    A: Scalar + Lapack,
108    S: Data<Elem = A>,
109    S2: Data<Elem = A>,
110{
111    type EigVal = Array1<A::Real>;
112    type EigVec = (Array2<A>, Array2<A>);
113
114    fn eigh(&self, uplo: UPLO) -> Result<(Self::EigVal, Self::EigVec)> {
115        let (a, b) = (self.0.to_owned(), self.1.to_owned());
116        (a, b).eigh_into(uplo)
117    }
118}
119
120impl<A, S> EighInplace for ArrayBase<S, Ix2>
121where
122    A: Scalar + Lapack,
123    S: DataMut<Elem = A>,
124{
125    type EigVal = Array1<A::Real>;
126
127    fn eigh_inplace(&mut self, uplo: UPLO) -> Result<(Self::EigVal, &mut Self)> {
128        let layout = self.square_layout()?;
129        // XXX Force layout to be Fortran (see #146)
130        match layout {
131            MatrixLayout::C { .. } => self.swap_axes(0, 1),
132            MatrixLayout::F { .. } => {}
133        }
134        let s = A::eigh(true, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
135        Ok((ArrayBase::from(s), self))
136    }
137}
138
139impl<A, S, S2> EighInplace for (ArrayBase<S, Ix2>, ArrayBase<S2, Ix2>)
140where
141    A: Scalar + Lapack,
142    S: DataMut<Elem = A>,
143    S2: DataMut<Elem = A>,
144{
145    type EigVal = Array1<A::Real>;
146
147    /// Solves the generalized eigenvalue problem.
148    ///
149    /// # Panics
150    ///
151    /// Panics if the shapes of the matrices are different.
152    fn eigh_inplace(&mut self, uplo: UPLO) -> Result<(Self::EigVal, &mut Self)> {
153        assert_eq!(
154            self.0.shape(),
155            self.1.shape(),
156            "The shapes of the matrices must be identical.",
157        );
158        let layout = self.0.square_layout()?;
159        // XXX Force layout to be Fortran (see #146)
160        match layout {
161            MatrixLayout::C { .. } => self.0.swap_axes(0, 1),
162            MatrixLayout::F { .. } => {}
163        }
164
165        let layout = self.1.square_layout()?;
166        match layout {
167            MatrixLayout::C { .. } => self.1.swap_axes(0, 1),
168            MatrixLayout::F { .. } => {}
169        }
170
171        let s = A::eigh_generalized(
172            true,
173            self.0.square_layout()?,
174            uplo,
175            self.0.as_allocated_mut()?,
176            self.1.as_allocated_mut()?,
177        )?;
178
179        Ok((ArrayBase::from(s), self))
180    }
181}
182
183/// Calculate eigenvalues without eigenvectors
184pub trait EigValsh {
185    type EigVal;
186    fn eigvalsh(&self, uplo: UPLO) -> Result<Self::EigVal>;
187}
188
189/// Calculate eigenvalues without eigenvectors
190pub trait EigValshInto {
191    type EigVal;
192    fn eigvalsh_into(self, uplo: UPLO) -> Result<Self::EigVal>;
193}
194
195/// Calculate eigenvalues without eigenvectors
196pub trait EigValshInplace {
197    type EigVal;
198    fn eigvalsh_inplace(&mut self, uplo: UPLO) -> Result<Self::EigVal>;
199}
200
201impl<A, S> EigValshInto for ArrayBase<S, Ix2>
202where
203    A: Scalar + Lapack,
204    S: DataMut<Elem = A>,
205{
206    type EigVal = Array1<A::Real>;
207
208    fn eigvalsh_into(mut self, uplo: UPLO) -> Result<Self::EigVal> {
209        self.eigvalsh_inplace(uplo)
210    }
211}
212
213impl<A, S> EigValsh for ArrayBase<S, Ix2>
214where
215    A: Scalar + Lapack,
216    S: Data<Elem = A>,
217{
218    type EigVal = Array1<A::Real>;
219
220    fn eigvalsh(&self, uplo: UPLO) -> Result<Self::EigVal> {
221        let a = self.to_owned();
222        a.eigvalsh_into(uplo)
223    }
224}
225
226impl<A, S> EigValshInplace for ArrayBase<S, Ix2>
227where
228    A: Scalar + Lapack,
229    S: DataMut<Elem = A>,
230{
231    type EigVal = Array1<A::Real>;
232
233    fn eigvalsh_inplace(&mut self, uplo: UPLO) -> Result<Self::EigVal> {
234        let s = A::eigh(true, self.square_layout()?, uplo, self.as_allocated_mut()?)?;
235        Ok(ArrayBase::from(s))
236    }
237}
238
239/// Calculate symmetric square-root matrix using `eigh`
240pub trait SymmetricSqrt {
241    type Output;
242    fn ssqrt(&self, uplo: UPLO) -> Result<Self::Output>;
243}
244
245impl<A, S> SymmetricSqrt for ArrayBase<S, Ix2>
246where
247    A: Scalar + Lapack,
248    S: Data<Elem = A>,
249{
250    type Output = Array2<A>;
251
252    fn ssqrt(&self, uplo: UPLO) -> Result<Self::Output> {
253        let a = self.to_owned();
254        a.ssqrt_into(uplo)
255    }
256}
257
258/// Calculate symmetric square-root matrix using `eigh`
259pub trait SymmetricSqrtInto {
260    type Output;
261    fn ssqrt_into(self, uplo: UPLO) -> Result<Self::Output>;
262}
263
264impl<A, S> SymmetricSqrtInto for ArrayBase<S, Ix2>
265where
266    A: Scalar + Lapack,
267    S: DataMut<Elem = A> + DataOwned,
268{
269    type Output = Array2<A>;
270
271    fn ssqrt_into(self, uplo: UPLO) -> Result<Self::Output> {
272        let (e, v) = self.eigh_into(uplo)?;
273        let e_sqrt = Array::from_iter(e.iter().map(|r| Scalar::from_real(r.sqrt())));
274        let ev = e_sqrt.into_diagonal().apply2(&v.t());
275        Ok(v.apply2(&ev))
276    }
277}