Actual source code: ex20td.c
petsc-3.14.6 2021-03-30
1: static char help[] = "Performs adjoint sensitivity analysis for a van der Pol like \n\
2: equation with time dependent parameters using two approaches : \n\
3: track : track only local sensitivities at each adjoint step \n\
4: and accumulate them in a global array \n\
5: global : track parameters at all timesteps together \n\
6: Choose one of the two at runtime by -sa_method {track,global}. \n";
8: /*
9: Concepts: TS^adjoint for time dependent parameters
10: Concepts: Customized adjoint monitor based sensitivity tracking
11: Concepts: All at once approach to sensitivity tracking
12: Processors: 1
13: */
16: /*
17: Simple example to demonstrate TSAdjoint capabilities for time dependent params
18: without integral cost terms using either a tracking or global method.
20: Modify the Van Der Pol Eq to :
21: [u1'] = [mu1(t)*u1]
22: [u2'] = [mu2(t)*((1-u1^2)*u2-u1)]
23: (with initial conditions & params independent)
25: Define uref to be solution with initail conditions (2,-2/3), mu=(1,1e3)
26: - u_ref : (1.5967,-1.02969)
28: Define const function as cost = 2-norm(u - u_ref);
30: Initialization for the adjoint TS :
31: - dcost/dy|final_time = 2*(u-u_ref)|final_time
32: - dcost/dp|final_time = 0
34: The tracking method only tracks local sensitivity at each time step
35: and accumulates these sensitivities in a global array. Since the structure
36: of the equations being solved at each time step does not change, the jacobian
37: wrt parameters is defined analogous to constant RHSJacobian for a liner
38: TSSolve and the size of the jacP is independent of the number of time
39: steps. Enable this mode of adjoint analysis by -sa_method track.
41: The global method combines the parameters at all timesteps and tracks them
42: together. Thus, the columns of the jacP matrix are filled dependent upon the
43: time step. Also, the dimensions of the jacP matrix now depend upon the number
44: of time steps. Enable this mode of adjoint analysis by -sa_method global.
46: Since the equations here have parameters at predefined time steps, this
47: example should be run with non adaptive time stepping solvers only. This
48: can be ensured by -ts_adapt_type none (which is the default behavior only
49: for certain TS solvers like TSCN. If using an explicit TS like TSRK,
50: please be sure to add the aforementioned option to disable adaptive
51: timestepping.)
52: */
54: /*
55: Include "petscts.h" so that we can use TS solvers. Note that this file
56: automatically includes:
57: petscsys.h - base PETSc routines petscvec.h - vectors
58: petscmat.h - matrices
59: petscis.h - index sets petscksp.h - Krylov subspace methods
60: petscviewer.h - viewers petscpc.h - preconditioners
61: petscksp.h - linear solvers petscsnes.h - nonlinear solvers
62: */
63: #include <petscts.h>
65: extern PetscErrorCode RHSFunction(TS ,PetscReal ,Vec ,Vec ,void *);
66: extern PetscErrorCode RHSJacobian(TS ,PetscReal ,Vec ,Mat ,Mat ,void *);
67: extern PetscErrorCode RHSJacobianP_track(TS ,PetscReal ,Vec ,Mat ,void *);
68: extern PetscErrorCode RHSJacobianP_global(TS ,PetscReal ,Vec ,Mat ,void *);
69: extern PetscErrorCode Monitor(TS ,PetscInt ,PetscReal ,Vec ,void *);
70: extern PetscErrorCode AdjointMonitor(TS ,PetscInt ,PetscReal ,Vec ,PetscInt ,Vec *, Vec *,void *);
72: /*
73: User-defined application context - contains data needed by the
74: application-provided call-back routines.
75: */
77: typedef struct {
78: /*------------- Forward solve data structures --------------*/
79: PetscInt max_steps; /* number of steps to run ts for */
80: PetscReal final_time; /* final time to integrate to*/
81: PetscReal time_step; /* ts integration time step */
82: Vec mu1; /* time dependent params */
83: Vec mu2; /* time dependent params */
84: Vec U; /* solution vector */
85: Mat A; /* Jacobian matrix */
87: /*------------- Adjoint solve data structures --------------*/
88: Mat Jacp; /* JacobianP matrix */
89: Vec lambda; /* adjoint variable */
90: Vec mup; /* adjoint variable */
92: /*------------- Global accumation vecs for monitor based tracking --------------*/
93: Vec sens_mu1; /* global sensitivity accumulation */
94: Vec sens_mu2; /* global sensitivity accumulation */
95: PetscInt adj_idx; /* to keep track of adjoint solve index */
96: } AppCtx;
98: typedef enum {SA_TRACK, SA_GLOBAL} SAMethod;
99: static const char *const SAMethods[] = {"TRACK","GLOBAL","SAMethod","SA_",0};
101: /* ----------------------- Explicit form of the ODE -------------------- */
103: PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec U,Vec F,void *ctx)
104: {
105: PetscErrorCode ierr;
106: AppCtx *user = (AppCtx*) ctx;
107: PetscScalar *f;
108: PetscInt curr_step;
109: const PetscScalar *u;
110: const PetscScalar *mu1;
111: const PetscScalar *mu2;
114: TSGetStepNumber(ts,&curr_step);
115: VecGetArrayRead(U,&u);
116: VecGetArrayRead(user->mu1,&mu1);
117: VecGetArrayRead(user->mu2,&mu2);
118: VecGetArray(F,&f);
119: f[0] = mu1[curr_step]*u[1];
120: f[1] = mu2[curr_step]*((1.-u[0]*u[0])*u[1]-u[0]);
121: VecRestoreArrayRead(U,&u);
122: VecRestoreArrayRead(user->mu1,&mu1);
123: VecRestoreArrayRead(user->mu2,&mu2);
124: VecRestoreArray(F,&f);
125: return(0);
126: }
128: PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,void *ctx)
129: {
130: PetscErrorCode ierr;
131: AppCtx *user = (AppCtx*) ctx;
132: PetscInt rowcol[] = {0,1};
133: PetscScalar J[2][2];
134: PetscInt curr_step;
135: const PetscScalar *u;
136: const PetscScalar *mu1;
137: const PetscScalar *mu2;
140: TSGetStepNumber(ts,&curr_step);
141: VecGetArrayRead(user->mu1,&mu1);
142: VecGetArrayRead(user->mu2,&mu2);
143: VecGetArrayRead(U,&u);
144: J[0][0] = 0;
145: J[1][0] = -mu2[curr_step]*(2.0*u[1]*u[0]+1.);
146: J[0][1] = mu1[curr_step];
147: J[1][1] = mu2[curr_step]*(1.0-u[0]*u[0]);
148: MatSetValues(A,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);
149: MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
150: MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
151: VecRestoreArrayRead(U,&u);
152: VecRestoreArrayRead(user->mu1,&mu1);
153: VecRestoreArrayRead(user->mu2,&mu2);
154: return(0);
155: }
157: /* ------------------ Jacobian wrt parameters for tracking method ------------------ */
159: PetscErrorCode RHSJacobianP_track(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
160: {
161: PetscErrorCode ierr;
162: PetscInt row[] = {0,1},col[] = {0,1};
163: PetscScalar J[2][2];
164: const PetscScalar *u;
167: VecGetArrayRead(U,&u);
168: J[0][0] = u[1];
169: J[1][0] = 0;
170: J[0][1] = 0;
171: J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
172: MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
173: MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
174: MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
175: VecRestoreArrayRead(U,&u);
176: return(0);
177: }
179: /* ------------------ Jacobian wrt parameters for global method ------------------ */
181: PetscErrorCode RHSJacobianP_global(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
182: {
183: PetscErrorCode ierr;
184: PetscInt row[] = {0,1},col[] = {0,1};
185: PetscScalar J[2][2];
186: const PetscScalar *u;
187: PetscInt curr_step;
190: TSGetStepNumber(ts,&curr_step);
191: VecGetArrayRead(U,&u);
192: J[0][0] = u[1];
193: J[1][0] = 0;
194: J[0][1] = 0;
195: J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
196: col[0] = (curr_step)*2;
197: col[1] = (curr_step)*2+1;
198: MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
199: MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
200: MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
201: VecRestoreArrayRead(U,&u);
202: return(0);
203: }
205: /* Dump solution to console if called */
206: PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec U,void *ctx)
207: {
208: PetscErrorCode ierr;
211: PetscPrintf(PETSC_COMM_WORLD,"\n Solution at time %e is \n", t);
212: VecView(U,PETSC_VIEWER_STDOUT_WORLD);
213: return(0);
214: }
216: /* Customized adjoint monitor to keep track of local
217: sensitivities by storing them in a global sensitivity array.
218: Note : This routine is only used for the tracking method. */
219: PetscErrorCode AdjointMonitor(TS ts,PetscInt steps,PetscReal time,Vec u,PetscInt numcost,Vec *lambda, Vec *mu,void *ctx)
220: {
221: PetscErrorCode ierr;
222: AppCtx *user = (AppCtx*) ctx;
223: PetscInt curr_step;
224: PetscScalar *sensmu1_glob;
225: PetscScalar *sensmu2_glob;
226: const PetscScalar *sensmu_loc;
229: TSGetStepNumber(ts,&curr_step);
230: /* Note that we skip the first call to the monitor in the adjoint
231: solve since the sensitivities are already set (during
232: initialization of adjoint vectors).
233: We also note that each indvidial TSAdjointSolve calls the monitor
234: twice, once at the step it is integrating from and once at the step
235: it integrates to. Only the second call is useful for transferring
236: local sensitivities to the global array. */
237: if (curr_step == user->adj_idx) {
238: return(0);
239: } else {
240: VecGetArrayRead(*mu,&sensmu_loc);
241: VecGetArray(user->sens_mu1,&sensmu1_glob);
242: VecGetArray(user->sens_mu2,&sensmu2_glob);
243: sensmu1_glob[curr_step] = sensmu_loc[0];
244: sensmu2_glob[curr_step] = sensmu_loc[1];
245: VecRestoreArray(user->sens_mu1,&sensmu1_glob);
246: VecRestoreArray(user->sens_mu2,&sensmu2_glob);
247: VecRestoreArrayRead(*mu,&sensmu_loc);
248: return(0);
249: }
250: }
253: int main(int argc,char **argv)
254: {
255: TS ts;
256: AppCtx user;
257: PetscScalar *x_ptr,*y_ptr,*u_ptr;
258: PetscMPIInt size;
259: PetscBool monitor = PETSC_FALSE;
260: SAMethod sa = SA_GLOBAL;
263: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
264: Initialize program
265: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
266: PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;
267: MPI_Comm_size(PETSC_COMM_WORLD,&size);
268: if (size != 1) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
270: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
271: Set runtime options
272: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
273: PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"SA analysis options.","");{
274: PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL);
275: PetscOptionsEnum("-sa_method","Sensitivity analysis method (track or global)","",SAMethods,(PetscEnum)sa,(PetscEnum*)&sa,NULL);
276: }
277: PetscOptionsEnd();
279: user.final_time = 0.1;
280: user.max_steps = 5;
281: user.time_step = user.final_time/user.max_steps;
283: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
284: Create necessary matrix and vectors for forward solve.
285: Create Jacp matrix for adjoint solve.
286: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
287: VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu1);
288: VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu2);
289: VecSet(user.mu1,1.25);
290: VecSet(user.mu2,1.0e2);
292: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
293: For tracking method : create the global sensitivity array to
294: accumulate sensitivity with respect to parameters at each step.
295: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
296: if (sa == SA_TRACK) {
297: VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu1);
298: VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu2);
299: }
301: MatCreate(PETSC_COMM_WORLD,&user.A);
302: MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2);
303: MatSetFromOptions(user.A);
304: MatSetUp(user.A);
305: MatCreateVecs(user.A,&user.U,NULL);
307: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
308: Note that the dimensions of the Jacp matrix depend upon the
309: sensitivity analysis method being used !
310: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
311: MatCreate(PETSC_COMM_WORLD,&user.Jacp);
312: if (sa == SA_TRACK) {
313: MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2);
314: }
315: if (sa == SA_GLOBAL) {
316: MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,user.max_steps*2);
317: }
318: MatSetFromOptions(user.Jacp);
319: MatSetUp(user.Jacp);
321: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
322: Create timestepping solver context
323: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
324: TSCreate(PETSC_COMM_WORLD,&ts);
325: TSSetEquationType(ts,TS_EQ_ODE_EXPLICIT);
326: TSSetType(ts,TSCN);
328: TSSetRHSFunction(ts,NULL,RHSFunction,&user);
329: TSSetRHSJacobian(ts,user.A,user.A,RHSJacobian,&user);
330: if (sa == SA_TRACK) {
331: TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_track,&user);
332: }
333: if (sa == SA_GLOBAL) {
334: TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_global,&user);
335: }
337: TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
338: TSSetMaxTime(ts,user.final_time);
339: TSSetTimeStep(ts,user.final_time/user.max_steps);
341: if (monitor) {
342: TSMonitorSet(ts,Monitor,&user,NULL);
343: }
344: if (sa == SA_TRACK) {
345: TSAdjointMonitorSet(ts,AdjointMonitor,&user,NULL);
346: }
348: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
349: Set initial conditions
350: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
351: VecGetArray(user.U,&x_ptr);
352: x_ptr[0] = 2.0;
353: x_ptr[1] = -2.0/3.0;
354: VecRestoreArray(user.U,&x_ptr);
356: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
357: Save trajectory of solution so that TSAdjointSolve() may be used
358: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
359: TSSetSaveTrajectory(ts);
361: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
362: Set runtime options
363: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
364: TSSetFromOptions(ts);
366: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
367: Execute forward model and print solution.
368: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
369: TSSolve(ts,user.U);
370: PetscPrintf(PETSC_COMM_WORLD,"\n Solution of forward TS :\n");
371: VecView(user.U,PETSC_VIEWER_STDOUT_WORLD);
372: PetscPrintf(PETSC_COMM_WORLD,"\n Forward TS solve successfull! Adjoint run begins!\n");
374: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
375: Adjoint model starts here! Create adjoint vectors.
376: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
377: MatCreateVecs(user.A,&user.lambda,NULL);
378: MatCreateVecs(user.Jacp,&user.mup,NULL);
380: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
381: Set initial conditions for the adjoint vector
382: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
383: VecGetArray(user.U,&u_ptr);
384: VecGetArray(user.lambda,&y_ptr);
385: y_ptr[0] = 2*(u_ptr[0] - 1.5967);
386: y_ptr[1] = 2*(u_ptr[1] - -(1.02969));
387: VecRestoreArray(user.lambda,&y_ptr);
388: VecRestoreArray(user.U,&y_ptr);
389: VecSet(user.mup,0);
391: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
392: Set number of cost functions.
393: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
394: TSSetCostGradients(ts,1,&user.lambda,&user.mup);
396: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
397: The adjoint vector mup has to be reset for each adjoint step when
398: using the tracking method as we want to treat the parameters at each
399: time step one at a time and prevent accumulation of the sensitivities
400: from parameters at previous time steps.
401: This is not necessary for the global method as each time dependent
402: parameter is treated as an independent parameter.
403: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
404: if (sa == SA_TRACK) {
405: for (user.adj_idx=user.max_steps; user.adj_idx>0; user.adj_idx--) {
406: VecSet(user.mup,0);
407: TSAdjointSetSteps(ts, 1);
408: TSAdjointSolve(ts);
409: }
410: }
411: if (sa == SA_GLOBAL) {
412: TSAdjointSolve(ts);
413: }
415: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
416: Dispaly adjoint sensitivities wrt parameters and initial conditions
417: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
418: if (sa == SA_TRACK) {
419: PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt mu1: d[cost]/d[mu1]\n");
420: VecView(user.sens_mu1,PETSC_VIEWER_STDOUT_WORLD);
421: PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt mu2: d[cost]/d[mu2]\n");
422: VecView(user.sens_mu2,PETSC_VIEWER_STDOUT_WORLD);
423: }
425: if (sa == SA_GLOBAL) {
426: PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt params: d[cost]/d[p], where p refers to \n\
427: the interlaced vector made by combining mu1,mu2\n");
428: VecView(user.mup,PETSC_VIEWER_STDOUT_WORLD);
429: }
431: PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt initial conditions: d[cost]/d[u(t=0)]\n");
432: VecView(user.lambda,PETSC_VIEWER_STDOUT_WORLD);
434: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
435: Free work space!
436: All PETSc objects should be destroyed when they are no longer needed.
437: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
438: MatDestroy(&user.A);
439: MatDestroy(&user.Jacp);
440: VecDestroy(&user.U);
441: VecDestroy(&user.lambda);
442: VecDestroy(&user.mup);
443: VecDestroy(&user.mu1);
444: VecDestroy(&user.mu2);
445: if (sa == SA_TRACK) {
446: VecDestroy(&user.sens_mu1);
447: VecDestroy(&user.sens_mu2);
448: }
449: TSDestroy(&ts);
450: PetscFinalize();
451: return(ierr);
452: }
455: /*TEST
457: test:
458: requires: !complex
459: suffix : track
460: args : -sa_method track
462: test:
463: requires: !complex
464: suffix : global
465: args : -sa_method global
467: TEST*/