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