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