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