xref: /petsc/src/ts/tutorials/autodiff/adolc-utils/drivers.cxx (revision a69119a591a03a9d906b29c0a4e9802e4d7c9795)
1 #include "contexts.cxx"
2 #include "sparse.cxx"
3 #include "init.cxx"
4 #include <adolc/drivers/drivers.h>
5 #include <adolc/interfaces.h>
6 
7 /*
8    REQUIRES configuration of PETSc with option --download-adolc.
9 
10    For documentation on ADOL-C, see
11      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
12 */
13 
14 /* --------------------------------------------------------------------------------
15    Drivers for RHSJacobian and IJacobian
16    ----------------------------------------------------------------------------- */
17 
18 /*
19   Compute Jacobian for explicit TS in compressed format and recover from this, using
20   precomputed seed and recovery matrices. If sparse mode is not used, full Jacobian is
21   assembled (not recommended for non-toy problems!).
22 
23   Input parameters:
24   tag   - tape identifier
25   u_vec - vector at which to evaluate Jacobian
26   ctx   - ADOL-C context, as defined above
27 
28   Output parameter:
29   A     - Mat object corresponding to Jacobian
30 */
31 PetscErrorCode PetscAdolcComputeRHSJacobian(PetscInt tag, Mat A, const PetscScalar *u_vec, void *ctx) {
32   AdolcCtx     *adctx = (AdolcCtx *)ctx;
33   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
34   PetscScalar **J;
35 
36   PetscFunctionBegin;
37   PetscCall(AdolcMalloc2(m, p, &J));
38   if (adctx->Seed) fov_forward(tag, m, n, p, u_vec, adctx->Seed, NULL, J);
39   else jacobian(tag, m, n, u_vec, J);
40   if (adctx->sparse) {
41     PetscCall(RecoverJacobian(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
42   } else {
43     for (i = 0; i < m; i++) {
44       for (j = 0; j < n; j++) {
45         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValues(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
46       }
47     }
48   }
49   PetscCall(AdolcFree2(J));
50   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
51   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
52   PetscFunctionReturn(0);
53 }
54 
55 /*
56   Compute Jacobian for explicit TS in compressed format and recover from this, using
57   precomputed seed and recovery matrices. If sparse mode is not used, full Jacobian is
58   assembled (not recommended for non-toy problems!).
59 
60   Input parameters:
61   tag   - tape identifier
62   u_vec - vector at which to evaluate Jacobian
63   ctx   - ADOL-C context, as defined above
64 
65   Output parameter:
66   A     - Mat object corresponding to Jacobian
67 */
68 PetscErrorCode PetscAdolcComputeRHSJacobianLocal(PetscInt tag, Mat A, const PetscScalar *u_vec, void *ctx) {
69   AdolcCtx     *adctx = (AdolcCtx *)ctx;
70   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
71   PetscScalar **J;
72 
73   PetscFunctionBegin;
74   PetscCall(AdolcMalloc2(m, p, &J));
75   if (adctx->Seed) fov_forward(tag, m, n, p, u_vec, adctx->Seed, NULL, J);
76   else jacobian(tag, m, n, u_vec, J);
77   if (adctx->sparse) {
78     PetscCall(RecoverJacobianLocal(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
79   } else {
80     for (i = 0; i < m; i++) {
81       for (j = 0; j < n; j++) {
82         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
83       }
84     }
85   }
86   PetscCall(AdolcFree2(J));
87   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
88   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
89   PetscFunctionReturn(0);
90 }
91 
92 /*
93   Compute Jacobian for implicit TS in compressed format and recover from this, using
94   precomputed seed and recovery matrices. If sparse mode is not used, full Jacobian is
95   assembled (not recommended for non-toy problems!).
96 
97   Input parameters:
98   tag1   - tape identifier for df/dx part
99   tag2   - tape identifier for df/d(xdot) part
100   u_vec - vector at which to evaluate Jacobian
101   ctx   - ADOL-C context, as defined above
102 
103   Output parameter:
104   A     - Mat object corresponding to Jacobian
105 */
106 PetscErrorCode PetscAdolcComputeIJacobian(PetscInt tag1, PetscInt tag2, Mat A, const PetscScalar *u_vec, PetscReal a, void *ctx) {
107   AdolcCtx     *adctx = (AdolcCtx *)ctx;
108   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
109   PetscScalar **J;
110 
111   PetscFunctionBegin;
112   PetscCall(AdolcMalloc2(m, p, &J));
113 
114   /* dF/dx part */
115   if (adctx->Seed) fov_forward(tag1, m, n, p, u_vec, adctx->Seed, NULL, J);
116   else jacobian(tag1, m, n, u_vec, J);
117   PetscCall(MatZeroEntries(A));
118   if (adctx->sparse) {
119     PetscCall(RecoverJacobian(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
120   } else {
121     for (i = 0; i < m; i++) {
122       for (j = 0; j < n; j++) {
123         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValues(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
124       }
125     }
126   }
127   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
128   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
129 
130   /* a * dF/d(xdot) part */
131   if (adctx->Seed) fov_forward(tag2, m, n, p, u_vec, adctx->Seed, NULL, J);
132   else jacobian(tag2, m, n, u_vec, J);
133   if (adctx->sparse) {
134     PetscCall(RecoverJacobian(A, ADD_VALUES, m, p, adctx->Rec, J, &a));
135   } else {
136     for (i = 0; i < m; i++) {
137       for (j = 0; j < n; j++) {
138         if (fabs(J[i][j]) > 1.e-16) {
139           J[i][j] *= a;
140           PetscCall(MatSetValues(A, 1, &i, 1, &j, &J[i][j], ADD_VALUES));
141         }
142       }
143     }
144   }
145   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
146   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
147   PetscCall(AdolcFree2(J));
148   PetscFunctionReturn(0);
149 }
150 
151 /*
152   Compute Jacobian for implicit TS in the special case where it is
153   known that the mass matrix is simply the identity. i.e. We have
154   a problem of the form
155                         du/dt = F(u).
156 
157   Input parameters:
158   tag   - tape identifier for df/dx part
159   u_vec - vector at which to evaluate Jacobian
160   ctx   - ADOL-C context, as defined above
161 
162   Output parameter:
163   A     - Mat object corresponding to Jacobian
164 */
165 PetscErrorCode PetscAdolcComputeIJacobianIDMass(PetscInt tag, Mat A, PetscScalar *u_vec, PetscReal a, void *ctx) {
166   AdolcCtx     *adctx = (AdolcCtx *)ctx;
167   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
168   PetscScalar **J;
169 
170   PetscFunctionBegin;
171   PetscCall(AdolcMalloc2(m, p, &J));
172 
173   /* dF/dx part */
174   if (adctx->Seed) fov_forward(tag, m, n, p, u_vec, adctx->Seed, NULL, J);
175   else jacobian(tag, m, n, u_vec, J);
176   PetscCall(MatZeroEntries(A));
177   if (adctx->sparse) {
178     PetscCall(RecoverJacobian(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
179   } else {
180     for (i = 0; i < m; i++) {
181       for (j = 0; j < n; j++) {
182         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValues(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
183       }
184     }
185   }
186   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
187   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
188   PetscCall(AdolcFree2(J));
189 
190   /* a * dF/d(xdot) part */
191   PetscCall(MatShift(A, a));
192   PetscFunctionReturn(0);
193 }
194 
195 /*
196   Compute local portion of Jacobian for implicit TS in compressed format and recover from this, using
197   precomputed seed and recovery matrices. If sparse mode is not used, full Jacobian is
198   assembled (not recommended for non-toy problems!).
199 
200   Input parameters:
201   tag1   - tape identifier for df/dx part
202   tag2   - tape identifier for df/d(xdot) part
203   u_vec - vector at which to evaluate Jacobian
204   ctx   - ADOL-C context, as defined above
205 
206   Output parameter:
207   A     - Mat object corresponding to Jacobian
208 */
209 PetscErrorCode PetscAdolcComputeIJacobianLocal(PetscInt tag1, PetscInt tag2, Mat A, PetscScalar *u_vec, PetscReal a, void *ctx) {
210   AdolcCtx     *adctx = (AdolcCtx *)ctx;
211   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
212   PetscScalar **J;
213 
214   PetscFunctionBegin;
215   PetscCall(AdolcMalloc2(m, p, &J));
216 
217   /* dF/dx part */
218   if (adctx->Seed) fov_forward(tag1, m, n, p, u_vec, adctx->Seed, NULL, J);
219   else jacobian(tag1, m, n, u_vec, J);
220   if (adctx->sparse) {
221     PetscCall(RecoverJacobianLocal(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
222   } else {
223     for (i = 0; i < m; i++) {
224       for (j = 0; j < n; j++) {
225         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
226       }
227     }
228   }
229   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
230   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
231 
232   /* a * dF/d(xdot) part */
233   if (adctx->Seed) fov_forward(tag2, m, n, p, u_vec, adctx->Seed, NULL, J);
234   else jacobian(tag2, m, n, u_vec, J);
235   if (adctx->sparse) {
236     PetscCall(RecoverJacobianLocal(A, ADD_VALUES, m, p, adctx->Rec, J, &a));
237   } else {
238     for (i = 0; i < m; i++) {
239       for (j = 0; j < n; j++) {
240         if (fabs(J[i][j]) > 1.e-16) {
241           J[i][j] *= a;
242           PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &J[i][j], ADD_VALUES));
243         }
244       }
245     }
246   }
247   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
248   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
249   PetscCall(AdolcFree2(J));
250   PetscFunctionReturn(0);
251 }
252 
253 /*
254   Compute local portion of Jacobian for implicit TS in the special case where it is
255   known that the mass matrix is simply the identity. i.e. We have
256   a problem of the form
257                         du/dt = F(u).
258 
259   Input parameters:
260   tag   - tape identifier for df/dx part
261   u_vec - vector at which to evaluate Jacobian
262   ctx   - ADOL-C context, as defined above
263 
264   Output parameter:
265   A     - Mat object corresponding to Jacobian
266 */
267 PetscErrorCode PetscAdolcComputeIJacobianLocalIDMass(PetscInt tag, Mat A, const PetscScalar *u_vec, PetscReal a, void *ctx) {
268   AdolcCtx     *adctx = (AdolcCtx *)ctx;
269   PetscInt      i, j, m = adctx->m, n = adctx->n, p = adctx->p;
270   PetscScalar **J;
271 
272   PetscFunctionBegin;
273   PetscCall(AdolcMalloc2(m, p, &J));
274 
275   /* dF/dx part */
276   if (adctx->Seed) fov_forward(tag, m, n, p, u_vec, adctx->Seed, NULL, J);
277   else jacobian(tag, m, n, u_vec, J);
278   if (adctx->sparse) {
279     PetscCall(RecoverJacobianLocal(A, INSERT_VALUES, m, p, adctx->Rec, J, NULL));
280   } else {
281     for (i = 0; i < m; i++) {
282       for (j = 0; j < n; j++) {
283         if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
284       }
285     }
286   }
287   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
288   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
289   PetscCall(AdolcFree2(J));
290 
291   /* a * dF/d(xdot) part */
292   PetscCall(MatShift(A, a));
293   PetscFunctionReturn(0);
294 }
295 
296 /* --------------------------------------------------------------------------------
297    Drivers for Jacobian w.r.t. a parameter
298    ----------------------------------------------------------------------------- */
299 
300 /*
301   Compute Jacobian w.r.t a parameter for explicit TS.
302 
303   Input parameters:
304   tag    - tape identifier
305   u_vec  - vector at which to evaluate Jacobian
306   params - the parameters w.r.t. which we differentiate
307   ctx    - ADOL-C context, as defined above
308 
309   Output parameter:
310   A      - Mat object corresponding to Jacobian
311 */
312 PetscErrorCode PetscAdolcComputeRHSJacobianP(PetscInt tag, Mat A, const PetscScalar *u_vec, PetscScalar *params, void *ctx) {
313   AdolcCtx     *adctx = (AdolcCtx *)ctx;
314   PetscInt      i, j = 0, m = adctx->m, n = adctx->n, p = adctx->num_params;
315   PetscScalar **J, *concat, **S;
316 
317   PetscFunctionBegin;
318 
319   /* Allocate memory and concatenate independent variable values with parameter */
320   PetscCall(AdolcMalloc2(m, p, &J));
321   PetscCall(PetscMalloc1(n + p, &concat));
322   PetscCall(AdolcMalloc2(n + p, p, &S));
323   PetscCall(Subidentity(p, n, S));
324   for (i = 0; i < n; i++) concat[i] = u_vec[i];
325   for (i = 0; i < p; i++) concat[n + i] = params[i];
326 
327   /* Propagate the appropriate seed matrix through the forward mode of AD */
328   fov_forward(tag, m, n + p, p, concat, S, NULL, J);
329   PetscCall(AdolcFree2(S));
330   PetscCall(PetscFree(concat));
331 
332   /* Set matrix values */
333   for (i = 0; i < m; i++) {
334     for (j = 0; j < p; j++) {
335       if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValues(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
336     }
337   }
338   PetscCall(AdolcFree2(J));
339   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
340   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
341   PetscFunctionReturn(0);
342 }
343 
344 /*
345   Compute local portion of Jacobian w.r.t a parameter for explicit TS.
346 
347   Input parameters:
348   tag    - tape identifier
349   u_vec  - vector at which to evaluate Jacobian
350   params - the parameters w.r.t. which we differentiate
351   ctx    - ADOL-C context, as defined above
352 
353   Output parameter:
354   A      - Mat object corresponding to Jacobian
355 */
356 PetscErrorCode PetscAdolcComputeRHSJacobianPLocal(PetscInt tag, Mat A, const PetscScalar *u_vec, PetscScalar *params, void *ctx) {
357   AdolcCtx     *adctx = (AdolcCtx *)ctx;
358   PetscInt      i, j = 0, m = adctx->m, n = adctx->n, p = adctx->num_params;
359   PetscScalar **J, *concat, **S;
360 
361   PetscFunctionBegin;
362 
363   /* Allocate memory and concatenate independent variable values with parameter */
364   PetscCall(AdolcMalloc2(m, p, &J));
365   PetscCall(PetscMalloc1(n + p, &concat));
366   PetscCall(AdolcMalloc2(n + p, p, &S));
367   PetscCall(Subidentity(p, n, S));
368   for (i = 0; i < n; i++) concat[i] = u_vec[i];
369   for (i = 0; i < p; i++) concat[n + i] = params[i];
370 
371   /* Propagate the appropriate seed matrix through the forward mode of AD */
372   fov_forward(tag, m, n + p, p, concat, S, NULL, J);
373   PetscCall(AdolcFree2(S));
374   PetscCall(PetscFree(concat));
375 
376   /* Set matrix values */
377   for (i = 0; i < m; i++) {
378     for (j = 0; j < p; j++) {
379       if (fabs(J[i][j]) > 1.e-16) PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &J[i][j], INSERT_VALUES));
380     }
381   }
382   PetscCall(AdolcFree2(J));
383   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
384   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
385   PetscFunctionReturn(0);
386 }
387 
388 /* --------------------------------------------------------------------------------
389    Drivers for Jacobian diagonal
390    ----------------------------------------------------------------------------- */
391 
392 /*
393   Compute local portion of Jacobian diagonal for implicit TS in compressed format and recover
394   from this, using precomputed seed matrix and recovery vector.
395 
396   Input parameters:
397   tag1  - tape identifier for df/dx part
398   tag2  - tape identifier for df/d(xdot) part
399   u_vec - vector at which to evaluate Jacobian
400   ctx   - ADOL-C context, as defined above
401 
402   Output parameter:
403   diag  - Vec object corresponding to Jacobian diagonal
404 */
405 PetscErrorCode PetscAdolcComputeIJacobianAndDiagonalLocal(PetscInt tag1, PetscInt tag2, Vec diag, PetscScalar *u_vec, PetscReal a, void *ctx) {
406   AdolcCtx     *adctx = (AdolcCtx *)ctx;
407   PetscInt      i, m = adctx->m, n = adctx->n, p = adctx->p;
408   PetscScalar **J;
409 
410   PetscFunctionBegin;
411   PetscCall(AdolcMalloc2(m, p, &J));
412 
413   /* dF/dx part */
414   if (adctx->Seed) fov_forward(tag1, m, n, p, u_vec, adctx->Seed, NULL, J);
415   else jacobian(tag1, m, n, u_vec, J);
416   if (adctx->sparse) {
417     PetscCall(RecoverDiagonalLocal(diag, INSERT_VALUES, m, adctx->rec, J, NULL));
418   } else {
419     for (i = 0; i < m; i++) {
420       if (fabs(J[i][i]) > 1.e-16) PetscCall(VecSetValuesLocal(diag, 1, &i, &J[i][i], INSERT_VALUES));
421     }
422   }
423   PetscCall(VecAssemblyBegin(diag));
424   PetscCall(VecAssemblyEnd(diag));
425 
426   /* a * dF/d(xdot) part */
427   if (adctx->Seed) fov_forward(tag2, m, n, p, u_vec, adctx->Seed, NULL, J);
428   else jacobian(tag2, m, n, u_vec, J);
429   if (adctx->sparse) {
430     PetscCall(RecoverDiagonalLocal(diag, ADD_VALUES, m, adctx->rec, J, NULL));
431   } else {
432     for (i = 0; i < m; i++) {
433       if (fabs(J[i][i]) > 1.e-16) {
434         J[i][i] *= a;
435         PetscCall(VecSetValuesLocal(diag, 1, &i, &J[i][i], ADD_VALUES));
436       }
437     }
438   }
439   PetscCall(VecAssemblyBegin(diag));
440   PetscCall(VecAssemblyEnd(diag));
441   PetscCall(AdolcFree2(J));
442   PetscFunctionReturn(0);
443 }
444