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
13namespace Eigen {
14
15namespace 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)
19template <typename MatrixType, typename ResultType>
20void 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
34template <typename MatrixType, typename ResultType>
35void 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()
43template <typename MatrixType, typename ResultType>
44void 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()
56template <typename MatrixType, typename ResultType>
57void 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
69template <typename MatrixType>
70void 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()
103template <typename MatrixType, typename ResultType>
104void 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
119template <typename MatrixType, typename ResultType>
120void 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.
139template <typename MatrixType, typename ResultType>
140void 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  */
181template <typename MatrixType, typename ResultType>
182void 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  */
205template <typename MatrixType, typename ResultType>
206void 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
231namespace 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  */
240template <typename MatrixType, int IsComplex = NumTraits<typename internal::traits<MatrixType>::Scalar>::IsComplex>
241struct 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
256template <typename MatrixType>
257struct 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
281template <typename MatrixType>
282struct 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  */
317template<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
353namespace internal {
354template<typename Derived>
355struct traits<MatrixSquareRootReturnValue<Derived> >
356{
357  typedef typename Derived::PlainObject ReturnType;
358};
359}
360
361template <typename Derived>
362const 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