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