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