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