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