xref: /libCEED/examples/fluids/qfunctions/utils_eigensolver_jacobi.h (revision 49a40c8a2d720db341b0b117b89656b473cbebfb)
1 // Copyright (c) 2017-2023, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 /// @file
9 /// Eigen system solver for symmetric NxN matrices. Modified from the CC0 code provided at https://github.com/jewettaij/jacobi_pd
10 
11 #ifndef utils_eigensolver_jacobi_h
12 #define utils_eigensolver_jacobi_h
13 
14 #include <ceed.h>
15 #include <math.h>
16 
17 #include "utils.h"
18 
19 // @typedef choose the criteria for sorting eigenvalues and eigenvectors
20 typedef enum eSortCriteria {
21   SORT_NONE,
22   SORT_DECREASING_EVALS,
23   SORT_INCREASING_EVALS,
24   SORT_DECREASING_ABS_EVALS,
25   SORT_INCREASING_ABS_EVALS
26 } SortCriteria;
27 
28 ///@brief Find the off-diagonal index in row i whose absolute value is largest
29 ///
30 /// @param[in] *A matrix
31 /// @param[in] i row index
32 /// @returns   Index of absolute largest off-diagonal element in row i
33 CEED_QFUNCTION_HELPER CeedInt MaxEntryRow(const CeedScalar *A, CeedInt N, CeedInt i) {
34   CeedInt j_max = i + 1;
35   for (CeedInt j = i + 2; j < N; j++)
36     if (fabs(A[i * N + j]) > fabs(A[i * N + j_max])) j_max = j;
37   return j_max;
38 }
39 
40 /// @brief Find the indices (i_max, j_max) marking the location of the
41 ///        entry in the matrix with the largest absolute value.  This
42 ///        uses the max_idx_row[] array to find the answer in O(n) time.
43 ///
44 /// @param[in]    *A    matrix
45 /// @param[inout] i_max row index
46 /// @param[inout] j_max column index
47 CEED_QFUNCTION_HELPER void MaxEntry(const CeedScalar *A, CeedInt N, CeedInt *max_idx_row, CeedInt *i_max, CeedInt *j_max) {
48   *i_max               = 0;
49   *j_max               = max_idx_row[*i_max];
50   CeedScalar max_entry = fabs(A[*i_max * N + *j_max]);
51   for (CeedInt i = 1; i < N - 1; i++) {
52     CeedInt j = max_idx_row[i];
53     if (fabs(A[i * N + j]) > max_entry) {
54       max_entry = fabs(A[i * N + j]);
55       *i_max    = i;
56       *j_max    = j;
57     }
58   }
59 }
60 
61 /// @brief Calculate the components of a rotation matrix which performs a
62 ///        rotation in the i,j plane by an angle (θ) that (when multiplied on
63 ///        both sides) will zero the ij'th element of A, so that afterwards
64 ///        A[i][j] = 0.  The results will be stored in c, s, and t
65 ///        (which store cos(θ), sin(θ), and tan(θ), respectively).
66 ///
67 /// @param[in] *A matrix
68 /// @param[in] i row index
69 /// @param[in] j column index
70 CEED_QFUNCTION_HELPER void CalcRot(const CeedScalar *A, CeedInt N, CeedInt i, CeedInt j, CeedScalar *rotmat_cst) {
71   rotmat_cst[2]      = 1.0;  // = tan(θ)
72   CeedScalar A_jj_ii = (A[j * N + j] - A[i * N + i]);
73   if (A_jj_ii != 0.0) {
74     // kappa = (A[j][j] - A[i][i]) / (2*A[i][j])
75     CeedScalar kappa = A_jj_ii;
76     rotmat_cst[2]    = 0.0;
77     CeedScalar A_ij  = A[i * N + j];
78     if (A_ij != 0.0) {
79       kappa /= (2.0 * A_ij);
80       // t satisfies: t^2 + 2*t*kappa - 1 = 0
81       // (choose the root which has the smaller absolute value)
82       rotmat_cst[2] = 1.0 / (sqrt(1 + kappa * kappa) + fabs(kappa));
83       if (kappa < 0.0) rotmat_cst[2] = -rotmat_cst[2];
84     }
85   }
86   rotmat_cst[0] = 1.0 / sqrt(1 + rotmat_cst[2] * rotmat_cst[2]);
87   rotmat_cst[1] = rotmat_cst[0] * rotmat_cst[2];
88 }
89 
90 /// @brief  Perform a similarity transformation by multiplying matrix A on both
91 ///         sides by a rotation matrix (and its transpose) to eliminate A[i][j].
92 /// @details This rotation matrix performs a rotation in the i,j plane by
93 ///         angle θ.  This function assumes that c=cos(θ). s=sin(θ), t=tan(θ)
94 ///         have been calculated in advance (using the CalcRot() function).
95 ///         It also assumes that i<j.  The max_idx_row[] array is also updated.
96 ///         To save time, since the matrix is symmetric, the elements
97 ///         below the diagonal (ie. A[u][v] where u>v) are not computed.
98 /// @verbatim
99 ///   A' = R^T * A * R
100 /// where R the rotation in the i,j plane and ^T denotes the transpose.
101 ///                 i         j
102 ///       _                             _
103 ///      |  1                            |
104 ///      |    .                          |
105 ///      |      .                        |
106 ///      |        1                      |
107 ///      |          c   ...   s          |
108 ///      |          .  .      .          |
109 /// R  = |          .    1    .          |
110 ///      |          .      .  .          |
111 ///      |          -s  ...   c          |
112 ///      |                      1        |
113 ///      |                        .      |
114 ///      |                          .    |
115 ///      |_                           1 _|
116 /// @endverbatim
117 ///
118 /// Let A' denote the matrix A after multiplication by R^T and R.
119 /// The components of A' are:
120 ///
121 /// @verbatim
122 ///   A'_uv =  Σ_w  Σ_z   R_wu * A_wz * R_zv
123 /// @endverbatim
124 ///
125 /// Note that a the rotation at location i,j will modify all of the matrix
126 /// elements containing at least one index which is either i or j
127 /// such as: A[w][i], A[i][w], A[w][j], A[j][w].
128 /// Check and see whether these modified matrix elements exceed the
129 /// corresponding values in max_idx_row[] array for that row.
130 /// If so, then update max_idx_row for that row.
131 /// This is somewhat complicated by the fact that we must only consider
132 /// matrix elements in the upper-right triangle strictly above the diagonal.
133 /// (ie. matrix elements whose second index is > the first index).
134 /// The modified elements we must consider are marked with an "X" below:
135 ///
136 /// @verbatim
137 ///                 i         j
138 ///       _                             _
139 ///      |  .       X         X          |
140 ///      |    .     X         X          |
141 ///      |      .   X         X          |
142 ///      |        . X         X          |
143 ///      |          X X X X X 0 X X X X  |  i
144 ///      |            .       X          |
145 ///      |              .     X          |
146 /// A  = |                .   X          |
147 ///      |                  . X          |
148 ///      |                    X X X X X  |  j
149 ///      |                      .        |
150 ///      |                        .      |
151 ///      |                          .    |
152 ///      |_                           . _|
153 /// @endverbatim
154 ///
155 /// @param[in] *A matrix
156 /// @param[in] i row index
157 /// @param[in] j column index
158 CEED_QFUNCTION_HELPER void ApplyRot(CeedScalar *A, CeedInt N, CeedInt i, CeedInt j, CeedInt *max_idx_row, CeedScalar *rotmat_cst) {
159   // Compute the diagonal elements of A which have changed:
160   A[i * N + i] -= rotmat_cst[2] * A[i * N + j];
161   A[j * N + j] += rotmat_cst[2] * A[i * N + j];
162   // Note: This is algebraically equivalent to:
163   // A[i][i] = c*c*A[i][i] + s*s*A[j][j] - 2*s*c*A[i][j]
164   // A[j][j] = s*s*A[i][i] + c*c*A[j][j] + 2*s*c*A[i][j]
165 
166   // Update the off-diagonal elements of A which will change (above the diagonal)
167 
168   A[i * N + j] = 0.0;
169 
170   // compute A[w][i] and A[i][w] for all w!=i,considering above-diagonal elements
171   for (CeedInt w = 0; w < i; w++) {                                              // 0 <= w <  i  <  j < N
172     A[i * N + w] = A[w * N + i];                                                 // backup the previous value. store below diagonal (i>w)
173     A[w * N + i] = rotmat_cst[0] * A[w * N + i] - rotmat_cst[1] * A[w * N + j];  // A[w][i], A[w][j] from previous iteration
174     if (i == max_idx_row[w]) max_idx_row[w] = MaxEntryRow(A, N, w);
175     else if (fabs(A[w * N + i]) > fabs(A[w * N + max_idx_row[w]])) max_idx_row[w] = i;
176   }
177   for (CeedInt w = i + 1; w < j; w++) {                                          // 0 <= i <  w  <  j < N
178     A[w * N + i] = A[i * N + w];                                                 // backup the previous value. store below diagonal (w>i)
179     A[i * N + w] = rotmat_cst[0] * A[i * N + w] - rotmat_cst[1] * A[w * N + j];  // A[i][w], A[w][j] from previous iteration
180   }
181   for (CeedInt w = j + 1; w < N; w++) {                                          // 0 <= i < j+1 <= w < N
182     A[w * N + i] = A[i * N + w];                                                 // backup the previous value. store below diagonal (w>i)
183     A[i * N + w] = rotmat_cst[0] * A[i * N + w] - rotmat_cst[1] * A[j * N + w];  // A[i][w], A[j][w] from previous iteration
184   }
185 
186   // now that we're done modifying row i, we can update max_idx_row[i]
187   max_idx_row[i] = MaxEntryRow(A, N, i);
188 
189   // compute A[w][j] and A[j][w] for all w!=j,considering above-diagonal elements
190   for (CeedInt w = 0; w < i; w++) {                                              // 0 <=  w  <  i <  j < N
191     A[w * N + j] = rotmat_cst[1] * A[i * N + w] + rotmat_cst[0] * A[w * N + j];  // A[i][w], A[w][j] from previous iteration
192     if (j == max_idx_row[w]) max_idx_row[w] = MaxEntryRow(A, N, w);
193     else if (fabs(A[w * N + j]) > fabs(A[w * N + max_idx_row[w]])) max_idx_row[w] = j;
194   }
195   for (CeedInt w = i + 1; w < j; w++) {                                          // 0 <= i+1 <= w <  j < N
196     A[w * N + j] = rotmat_cst[1] * A[w * N + i] + rotmat_cst[0] * A[w * N + j];  // A[w][i], A[w][j] from previous iteration
197     if (j == max_idx_row[w]) max_idx_row[w] = MaxEntryRow(A, N, w);
198     else if (fabs(A[w * N + j]) > fabs(A[w * N + max_idx_row[w]])) max_idx_row[w] = j;
199   }
200   for (CeedInt w = j + 1; w < N; w++) {                                          // 0 <=  i  <  j <  w < N
201     A[j * N + w] = rotmat_cst[1] * A[w * N + i] + rotmat_cst[0] * A[j * N + w];  // A[w][i], A[j][w] from previous iteration
202   }
203   // now that we're done modifying row j, we can update max_idx_row[j]
204   max_idx_row[j] = MaxEntryRow(A, N, j);
205 }
206 
207 ///@brief Multiply matrix A on the LEFT side by a transposed rotation matrix R^T
208 ///       This matrix performs a rotation in the i,j plane by angle θ  (where
209 ///       the arguments "s" and "c" refer to cos(θ) and sin(θ), respectively).
210 /// @verbatim
211 ///   A'_uv = Σ_w  R_wu * A_wv
212 /// @endverbatim
213 ///
214 /// @param[in] *A matrix
215 /// @param[in] i row index
216 /// @param[in] j column index
217 CEED_QFUNCTION_HELPER void ApplyRotLeft(CeedScalar *A, CeedInt N, CeedInt i, CeedInt j, CeedScalar *rotmat_cst) {
218   // Recall that c = cos(θ) and s = sin(θ)
219   for (CeedInt v = 0; v < N; v++) {
220     CeedScalar Aiv = A[i * N + v];
221     A[i * N + v]   = rotmat_cst[0] * A[i * N + v] - rotmat_cst[1] * A[j * N + v];
222     A[j * N + v]   = rotmat_cst[1] * Aiv + rotmat_cst[0] * A[j * N + v];
223   }
224 }
225 
226 /// @brief Sort the rows in evec according to the numbers in v (also sorted)
227 ///
228 /// @param[inout] *eval vector containing the keys used for sorting
229 /// @param[inout] *evec matrix whose rows will be sorted according to v
230 /// @param[in]    n  size of the vector and matrix
231 /// @param[in]    s  sort decreasing order?
232 CEED_QFUNCTION_HELPER void SortRows(CeedScalar *eval, CeedScalar *evec, CeedInt N, SortCriteria sort_criteria) {
233   if (sort_criteria == SORT_NONE) return;
234 
235   for (CeedInt i = 0; i < N - 1; i++) {
236     CeedInt i_max = i;
237     for (CeedInt j = i + 1; j < N; j++) {
238       // find the "maximum" element in the array starting at position i+1
239       switch (sort_criteria) {
240         case SORT_DECREASING_EVALS:
241           if (eval[j] > eval[i_max]) i_max = j;
242           break;
243         case SORT_INCREASING_EVALS:
244           if (eval[j] < eval[i_max]) i_max = j;
245           break;
246         case SORT_DECREASING_ABS_EVALS:
247           if (fabs(eval[j]) > fabs(eval[i_max])) i_max = j;
248           break;
249         case SORT_INCREASING_ABS_EVALS:
250           if (fabs(eval[j]) < fabs(eval[i_max])) i_max = j;
251           break;
252         default:
253           break;
254       }
255     }
256     SwapScalar(&eval[i], &eval[i_max]);
257     for (CeedInt k = 0; k < N; k++) SwapScalar(&evec[i * N + k], &evec[i_max * N + k]);
258   }
259 }
260 
261 /// @brief Calculate all the eigenvalues and eigevectors of a symmetric matrix
262 ///        using the Jacobi eigenvalue algorithm:
263 ///        https://en.wikipedia.org/wiki/Jacobi_eigenvalue_algorithm
264 /// @returns The number of Jacobi iterations attempted, which should be > 0.
265 ///          If the return value is not strictly > 0 then convergence failed.
266 /// @note  To reduce the computation time further, set calc_evecs=false.
267 ///        Additionally, note that the output evecs should be normalized. It
268 ///        simply takes the Identity matrix and performs (isometric) rotations
269 ///        on it, so divergence from normalized is due to finite-precision
270 ///        arithmetic of the rotations.
271 //
272 // @param[in]  A              the matrix you wish to diagonalize (size NxN)
273 // @param[in]  N              size of the matrix
274 // @param[out] eval           store the eigenvalues here (size N)
275 // @param[out] evec           store the eigenvectors here (in rows, size NxN)
276 // @param[out] max_idx_row    work vector of size N
277 // @param[in]  sort_criteria  sort results?
278 // @param[in]  calc_evecs     calculate the eigenvectors?
279 // @param[in]  max_num_sweeps maximum number of iterations = max_num_sweeps * number of off-diagonals (N*(N-1)/2)
280 CEED_QFUNCTION_HELPER CeedInt Diagonalize(CeedScalar *A, CeedInt N, CeedScalar *eval, CeedScalar *evec, CeedInt *max_idx_row,
281                                           SortCriteria sort_criteria, bool calc_evec, const CeedInt max_num_sweeps) {
282   CeedScalar rotmat_cst[3] = {0.};  // cos(θ), sin(θ), and tan(θ),
283 
284   if (calc_evec)
285     for (CeedInt i = 0; i < N; i++)
286       for (CeedInt j = 0; j < N; j++) evec[i * N + j] = (i == j) ? 1.0 : 0.0;  // Set evec equal to the identity matrix
287 
288   for (CeedInt i = 0; i < N - 1; i++) max_idx_row[i] = MaxEntryRow(A, N, i);
289 
290   // -- Iteration --
291   CeedInt n_iters;
292   CeedInt max_num_iters = max_num_sweeps * N * (N - 1) / 2;
293   for (n_iters = 1; n_iters <= max_num_iters; n_iters++) {
294     CeedInt i, j;
295     MaxEntry(A, N, max_idx_row, &i, &j);
296 
297     // If A[i][j] is small compared to A[i][i] and A[j][j], set it to 0.
298     if ((A[i * N + i] + A[i * N + j] == A[i * N + i]) && (A[j * N + j] + A[i * N + j] == A[j * N + j])) {
299       A[i * N + j]   = 0.0;
300       max_idx_row[i] = MaxEntryRow(A, N, i);
301     }
302 
303     if (A[i * N + j] == 0.0) break;
304 
305     CalcRot(A, N, i, j, rotmat_cst);                // Calculate the parameters of the rotation matrix.
306     ApplyRot(A, N, i, j, max_idx_row, rotmat_cst);  // Apply this rotation to the A matrix.
307     if (calc_evec) ApplyRotLeft(evec, N, i, j, rotmat_cst);
308   }
309 
310   for (CeedInt i = 0; i < N; i++) eval[i] = A[i * N + i];
311 
312   // Optional: Sort results by eigenvalue.
313   SortRows(eval, evec, N, sort_criteria);
314 
315   if ((n_iters > max_num_iters) && (N > 1))  // If we exceeded max_num_iters,
316     return 0;                                // indicate an error occured.
317 
318   return n_iters;
319 }
320 
321 // @brief Interface to Diagonalize for 3x3 systems
322 CEED_QFUNCTION_HELPER CeedInt Diagonalize3(CeedScalar A[3][3], CeedScalar eval[3], CeedScalar evec[3][3], CeedInt max_idx_row[3],
323                                            SortCriteria sort_criteria, bool calc_evec, const CeedInt max_num_sweeps) {
324   return Diagonalize((CeedScalar *)A, 3, (CeedScalar *)eval, (CeedScalar *)evec, (CeedInt *)max_idx_row, sort_criteria, calc_evec, max_num_sweeps);
325 }
326 
327 #endif  // utils_eigensolver_jacobi_h
328