1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2011, 2013 Jitse Niesen <jitse@maths.leeds.ac.uk>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9
10 #ifndef EIGEN_MATRIX_SQUARE_ROOT
11 #define EIGEN_MATRIX_SQUARE_ROOT
12
13 namespace Eigen {
14
15 namespace internal {
16
17 // pre: T.block(i,i,2,2) has complex conjugate eigenvalues
18 // post: sqrtT.block(i,i,2,2) is square root of T.block(i,i,2,2)
19 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_2x2_diagonal_block(const MatrixType & T,typename MatrixType::Index i,ResultType & sqrtT)20 void matrix_sqrt_quasi_triangular_2x2_diagonal_block(const MatrixType& T, typename MatrixType::Index i, ResultType& sqrtT)
21 {
22 // TODO: This case (2-by-2 blocks with complex conjugate eigenvalues) is probably hidden somewhere
23 // in EigenSolver. If we expose it, we could call it directly from here.
24 typedef typename traits<MatrixType>::Scalar Scalar;
25 Matrix<Scalar,2,2> block = T.template block<2,2>(i,i);
26 EigenSolver<Matrix<Scalar,2,2> > es(block);
27 sqrtT.template block<2,2>(i,i)
28 = (es.eigenvectors() * es.eigenvalues().cwiseSqrt().asDiagonal() * es.eigenvectors().inverse()).real();
29 }
30
31 // pre: block structure of T is such that (i,j) is a 1x1 block,
32 // all blocks of sqrtT to left of and below (i,j) are correct
33 // post: sqrtT(i,j) has the correct value
34 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_1x1_off_diagonal_block(const MatrixType & T,typename MatrixType::Index i,typename MatrixType::Index j,ResultType & sqrtT)35 void matrix_sqrt_quasi_triangular_1x1_off_diagonal_block(const MatrixType& T, typename MatrixType::Index i, typename MatrixType::Index j, ResultType& sqrtT)
36 {
37 typedef typename traits<MatrixType>::Scalar Scalar;
38 Scalar tmp = (sqrtT.row(i).segment(i+1,j-i-1) * sqrtT.col(j).segment(i+1,j-i-1)).value();
39 sqrtT.coeffRef(i,j) = (T.coeff(i,j) - tmp) / (sqrtT.coeff(i,i) + sqrtT.coeff(j,j));
40 }
41
42 // similar to compute1x1offDiagonalBlock()
43 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_1x2_off_diagonal_block(const MatrixType & T,typename MatrixType::Index i,typename MatrixType::Index j,ResultType & sqrtT)44 void matrix_sqrt_quasi_triangular_1x2_off_diagonal_block(const MatrixType& T, typename MatrixType::Index i, typename MatrixType::Index j, ResultType& sqrtT)
45 {
46 typedef typename traits<MatrixType>::Scalar Scalar;
47 Matrix<Scalar,1,2> rhs = T.template block<1,2>(i,j);
48 if (j-i > 1)
49 rhs -= sqrtT.block(i, i+1, 1, j-i-1) * sqrtT.block(i+1, j, j-i-1, 2);
50 Matrix<Scalar,2,2> A = sqrtT.coeff(i,i) * Matrix<Scalar,2,2>::Identity();
51 A += sqrtT.template block<2,2>(j,j).transpose();
52 sqrtT.template block<1,2>(i,j).transpose() = A.fullPivLu().solve(rhs.transpose());
53 }
54
55 // similar to compute1x1offDiagonalBlock()
56 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_2x1_off_diagonal_block(const MatrixType & T,typename MatrixType::Index i,typename MatrixType::Index j,ResultType & sqrtT)57 void matrix_sqrt_quasi_triangular_2x1_off_diagonal_block(const MatrixType& T, typename MatrixType::Index i, typename MatrixType::Index j, ResultType& sqrtT)
58 {
59 typedef typename traits<MatrixType>::Scalar Scalar;
60 Matrix<Scalar,2,1> rhs = T.template block<2,1>(i,j);
61 if (j-i > 2)
62 rhs -= sqrtT.block(i, i+2, 2, j-i-2) * sqrtT.block(i+2, j, j-i-2, 1);
63 Matrix<Scalar,2,2> A = sqrtT.coeff(j,j) * Matrix<Scalar,2,2>::Identity();
64 A += sqrtT.template block<2,2>(i,i);
65 sqrtT.template block<2,1>(i,j) = A.fullPivLu().solve(rhs);
66 }
67
68 // solves the equation A X + X B = C where all matrices are 2-by-2
69 template <typename MatrixType>
matrix_sqrt_quasi_triangular_solve_auxiliary_equation(MatrixType & X,const MatrixType & A,const MatrixType & B,const MatrixType & C)70 void matrix_sqrt_quasi_triangular_solve_auxiliary_equation(MatrixType& X, const MatrixType& A, const MatrixType& B, const MatrixType& C)
71 {
72 typedef typename traits<MatrixType>::Scalar Scalar;
73 Matrix<Scalar,4,4> coeffMatrix = Matrix<Scalar,4,4>::Zero();
74 coeffMatrix.coeffRef(0,0) = A.coeff(0,0) + B.coeff(0,0);
75 coeffMatrix.coeffRef(1,1) = A.coeff(0,0) + B.coeff(1,1);
76 coeffMatrix.coeffRef(2,2) = A.coeff(1,1) + B.coeff(0,0);
77 coeffMatrix.coeffRef(3,3) = A.coeff(1,1) + B.coeff(1,1);
78 coeffMatrix.coeffRef(0,1) = B.coeff(1,0);
79 coeffMatrix.coeffRef(0,2) = A.coeff(0,1);
80 coeffMatrix.coeffRef(1,0) = B.coeff(0,1);
81 coeffMatrix.coeffRef(1,3) = A.coeff(0,1);
82 coeffMatrix.coeffRef(2,0) = A.coeff(1,0);
83 coeffMatrix.coeffRef(2,3) = B.coeff(1,0);
84 coeffMatrix.coeffRef(3,1) = A.coeff(1,0);
85 coeffMatrix.coeffRef(3,2) = B.coeff(0,1);
86
87 Matrix<Scalar,4,1> rhs;
88 rhs.coeffRef(0) = C.coeff(0,0);
89 rhs.coeffRef(1) = C.coeff(0,1);
90 rhs.coeffRef(2) = C.coeff(1,0);
91 rhs.coeffRef(3) = C.coeff(1,1);
92
93 Matrix<Scalar,4,1> result;
94 result = coeffMatrix.fullPivLu().solve(rhs);
95
96 X.coeffRef(0,0) = result.coeff(0);
97 X.coeffRef(0,1) = result.coeff(1);
98 X.coeffRef(1,0) = result.coeff(2);
99 X.coeffRef(1,1) = result.coeff(3);
100 }
101
102 // similar to compute1x1offDiagonalBlock()
103 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_2x2_off_diagonal_block(const MatrixType & T,typename MatrixType::Index i,typename MatrixType::Index j,ResultType & sqrtT)104 void matrix_sqrt_quasi_triangular_2x2_off_diagonal_block(const MatrixType& T, typename MatrixType::Index i, typename MatrixType::Index j, ResultType& sqrtT)
105 {
106 typedef typename traits<MatrixType>::Scalar Scalar;
107 Matrix<Scalar,2,2> A = sqrtT.template block<2,2>(i,i);
108 Matrix<Scalar,2,2> B = sqrtT.template block<2,2>(j,j);
109 Matrix<Scalar,2,2> C = T.template block<2,2>(i,j);
110 if (j-i > 2)
111 C -= sqrtT.block(i, i+2, 2, j-i-2) * sqrtT.block(i+2, j, j-i-2, 2);
112 Matrix<Scalar,2,2> X;
113 matrix_sqrt_quasi_triangular_solve_auxiliary_equation(X, A, B, C);
114 sqrtT.template block<2,2>(i,j) = X;
115 }
116
117 // pre: T is quasi-upper-triangular and sqrtT is a zero matrix of the same size
118 // post: the diagonal blocks of sqrtT are the square roots of the diagonal blocks of T
119 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_diagonal(const MatrixType & T,ResultType & sqrtT)120 void matrix_sqrt_quasi_triangular_diagonal(const MatrixType& T, ResultType& sqrtT)
121 {
122 using std::sqrt;
123 typedef typename MatrixType::Index Index;
124 const Index size = T.rows();
125 for (Index i = 0; i < size; i++) {
126 if (i == size - 1 || T.coeff(i+1, i) == 0) {
127 eigen_assert(T(i,i) >= 0);
128 sqrtT.coeffRef(i,i) = sqrt(T.coeff(i,i));
129 }
130 else {
131 matrix_sqrt_quasi_triangular_2x2_diagonal_block(T, i, sqrtT);
132 ++i;
133 }
134 }
135 }
136
137 // pre: T is quasi-upper-triangular and diagonal blocks of sqrtT are square root of diagonal blocks of T.
138 // post: sqrtT is the square root of T.
139 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular_off_diagonal(const MatrixType & T,ResultType & sqrtT)140 void matrix_sqrt_quasi_triangular_off_diagonal(const MatrixType& T, ResultType& sqrtT)
141 {
142 typedef typename MatrixType::Index Index;
143 const Index size = T.rows();
144 for (Index j = 1; j < size; j++) {
145 if (T.coeff(j, j-1) != 0) // if T(j-1:j, j-1:j) is a 2-by-2 block
146 continue;
147 for (Index i = j-1; i >= 0; i--) {
148 if (i > 0 && T.coeff(i, i-1) != 0) // if T(i-1:i, i-1:i) is a 2-by-2 block
149 continue;
150 bool iBlockIs2x2 = (i < size - 1) && (T.coeff(i+1, i) != 0);
151 bool jBlockIs2x2 = (j < size - 1) && (T.coeff(j+1, j) != 0);
152 if (iBlockIs2x2 && jBlockIs2x2)
153 matrix_sqrt_quasi_triangular_2x2_off_diagonal_block(T, i, j, sqrtT);
154 else if (iBlockIs2x2 && !jBlockIs2x2)
155 matrix_sqrt_quasi_triangular_2x1_off_diagonal_block(T, i, j, sqrtT);
156 else if (!iBlockIs2x2 && jBlockIs2x2)
157 matrix_sqrt_quasi_triangular_1x2_off_diagonal_block(T, i, j, sqrtT);
158 else if (!iBlockIs2x2 && !jBlockIs2x2)
159 matrix_sqrt_quasi_triangular_1x1_off_diagonal_block(T, i, j, sqrtT);
160 }
161 }
162 }
163
164 } // end of namespace internal
165
166 /** \ingroup MatrixFunctions_Module
167 * \brief Compute matrix square root of quasi-triangular matrix.
168 *
169 * \tparam MatrixType type of \p arg, the argument of matrix square root,
170 * expected to be an instantiation of the Matrix class template.
171 * \tparam ResultType type of \p result, where result is to be stored.
172 * \param[in] arg argument of matrix square root.
173 * \param[out] result matrix square root of upper Hessenberg part of \p arg.
174 *
175 * This function computes the square root of the upper quasi-triangular matrix stored in the upper
176 * Hessenberg part of \p arg. Only the upper Hessenberg part of \p result is updated, the rest is
177 * not touched. See MatrixBase::sqrt() for details on how this computation is implemented.
178 *
179 * \sa MatrixSquareRoot, MatrixSquareRootQuasiTriangular
180 */
181 template <typename MatrixType, typename ResultType>
matrix_sqrt_quasi_triangular(const MatrixType & arg,ResultType & result)182 void matrix_sqrt_quasi_triangular(const MatrixType &arg, ResultType &result)
183 {
184 eigen_assert(arg.rows() == arg.cols());
185 result.resize(arg.rows(), arg.cols());
186 internal::matrix_sqrt_quasi_triangular_diagonal(arg, result);
187 internal::matrix_sqrt_quasi_triangular_off_diagonal(arg, result);
188 }
189
190
191 /** \ingroup MatrixFunctions_Module
192 * \brief Compute matrix square root of triangular matrix.
193 *
194 * \tparam MatrixType type of \p arg, the argument of matrix square root,
195 * expected to be an instantiation of the Matrix class template.
196 * \tparam ResultType type of \p result, where result is to be stored.
197 * \param[in] arg argument of matrix square root.
198 * \param[out] result matrix square root of upper triangular part of \p arg.
199 *
200 * Only the upper triangular part (including the diagonal) of \p result is updated, the rest is not
201 * touched. See MatrixBase::sqrt() for details on how this computation is implemented.
202 *
203 * \sa MatrixSquareRoot, MatrixSquareRootQuasiTriangular
204 */
205 template <typename MatrixType, typename ResultType>
matrix_sqrt_triangular(const MatrixType & arg,ResultType & result)206 void matrix_sqrt_triangular(const MatrixType &arg, ResultType &result)
207 {
208 using std::sqrt;
209 typedef typename MatrixType::Index Index;
210 typedef typename MatrixType::Scalar Scalar;
211
212 eigen_assert(arg.rows() == arg.cols());
213
214 // Compute square root of arg and store it in upper triangular part of result
215 // This uses that the square root of triangular matrices can be computed directly.
216 result.resize(arg.rows(), arg.cols());
217 for (Index i = 0; i < arg.rows(); i++) {
218 result.coeffRef(i,i) = sqrt(arg.coeff(i,i));
219 }
220 for (Index j = 1; j < arg.cols(); j++) {
221 for (Index i = j-1; i >= 0; i--) {
222 // if i = j-1, then segment has length 0 so tmp = 0
223 Scalar tmp = (result.row(i).segment(i+1,j-i-1) * result.col(j).segment(i+1,j-i-1)).value();
224 // denominator may be zero if original matrix is singular
225 result.coeffRef(i,j) = (arg.coeff(i,j) - tmp) / (result.coeff(i,i) + result.coeff(j,j));
226 }
227 }
228 }
229
230
231 namespace internal {
232
233 /** \ingroup MatrixFunctions_Module
234 * \brief Helper struct for computing matrix square roots of general matrices.
235 * \tparam MatrixType type of the argument of the matrix square root,
236 * expected to be an instantiation of the Matrix class template.
237 *
238 * \sa MatrixSquareRootTriangular, MatrixSquareRootQuasiTriangular, MatrixBase::sqrt()
239 */
240 template <typename MatrixType, int IsComplex = NumTraits<typename internal::traits<MatrixType>::Scalar>::IsComplex>
241 struct matrix_sqrt_compute
242 {
243 /** \brief Compute the matrix square root
244 *
245 * \param[in] arg matrix whose square root is to be computed.
246 * \param[out] result square root of \p arg.
247 *
248 * See MatrixBase::sqrt() for details on how this computation is implemented.
249 */
250 template <typename ResultType> static void run(const MatrixType &arg, ResultType &result);
251 };
252
253
254 // ********** Partial specialization for real matrices **********
255
256 template <typename MatrixType>
257 struct matrix_sqrt_compute<MatrixType, 0>
258 {
259 template <typename ResultType>
260 static void run(const MatrixType &arg, ResultType &result)
261 {
262 eigen_assert(arg.rows() == arg.cols());
263
264 // Compute Schur decomposition of arg
265 const RealSchur<MatrixType> schurOfA(arg);
266 const MatrixType& T = schurOfA.matrixT();
267 const MatrixType& U = schurOfA.matrixU();
268
269 // Compute square root of T
270 MatrixType sqrtT = MatrixType::Zero(arg.rows(), arg.cols());
271 matrix_sqrt_quasi_triangular(T, sqrtT);
272
273 // Compute square root of arg
274 result = U * sqrtT * U.adjoint();
275 }
276 };
277
278
279 // ********** Partial specialization for complex matrices **********
280
281 template <typename MatrixType>
282 struct matrix_sqrt_compute<MatrixType, 1>
283 {
284 template <typename ResultType>
285 static void run(const MatrixType &arg, ResultType &result)
286 {
287 eigen_assert(arg.rows() == arg.cols());
288
289 // Compute Schur decomposition of arg
290 const ComplexSchur<MatrixType> schurOfA(arg);
291 const MatrixType& T = schurOfA.matrixT();
292 const MatrixType& U = schurOfA.matrixU();
293
294 // Compute square root of T
295 MatrixType sqrtT;
296 matrix_sqrt_triangular(T, sqrtT);
297
298 // Compute square root of arg
299 result = U * (sqrtT.template triangularView<Upper>() * U.adjoint());
300 }
301 };
302
303 } // end namespace internal
304
305 /** \ingroup MatrixFunctions_Module
306 *
307 * \brief Proxy for the matrix square root of some matrix (expression).
308 *
309 * \tparam Derived Type of the argument to the matrix square root.
310 *
311 * This class holds the argument to the matrix square root until it
312 * is assigned or evaluated for some other reason (so the argument
313 * should not be changed in the meantime). It is the return type of
314 * MatrixBase::sqrt() and most of the time this is the only way it is
315 * used.
316 */
317 template<typename Derived> class MatrixSquareRootReturnValue
318 : public ReturnByValue<MatrixSquareRootReturnValue<Derived> >
319 {
320 protected:
321 typedef typename Derived::Index Index;
322 typedef typename internal::ref_selector<Derived>::type DerivedNested;
323
324 public:
325 /** \brief Constructor.
326 *
327 * \param[in] src %Matrix (expression) forming the argument of the
328 * matrix square root.
329 */
330 explicit MatrixSquareRootReturnValue(const Derived& src) : m_src(src) { }
331
332 /** \brief Compute the matrix square root.
333 *
334 * \param[out] result the matrix square root of \p src in the
335 * constructor.
336 */
337 template <typename ResultType>
338 inline void evalTo(ResultType& result) const
339 {
340 typedef typename internal::nested_eval<Derived, 10>::type DerivedEvalType;
341 typedef typename internal::remove_all<DerivedEvalType>::type DerivedEvalTypeClean;
342 DerivedEvalType tmp(m_src);
343 internal::matrix_sqrt_compute<DerivedEvalTypeClean>::run(tmp, result);
344 }
345
346 Index rows() const { return m_src.rows(); }
347 Index cols() const { return m_src.cols(); }
348
349 protected:
350 const DerivedNested m_src;
351 };
352
353 namespace internal {
354 template<typename Derived>
355 struct traits<MatrixSquareRootReturnValue<Derived> >
356 {
357 typedef typename Derived::PlainObject ReturnType;
358 };
359 }
360
361 template <typename Derived>
362 const MatrixSquareRootReturnValue<Derived> MatrixBase<Derived>::sqrt() const
363 {
364 eigen_assert(rows() == cols());
365 return MatrixSquareRootReturnValue<Derived>(derived());
366 }
367
368 } // end namespace Eigen
369
370 #endif // EIGEN_MATRIX_FUNCTION
371