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 initial 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 ODE integrators 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 the TS ODE integrators. Note that this file
48: automatically includes:
49: petscsys.h - base PETSc routines
50: petscis.h - index sets
51: petscvec.h - vectors
52: petscmat.h - matrices
53: petscpc.h - preconditioners
54: petscksp.h - linear solvers
55: petscsnes.h - nonlinear solvers
56: */
57: #include <petscts.h>
59: extern PetscErrorCode RHSFunction(TS, PetscReal, Vec, Vec, void *);
60: extern PetscErrorCode RHSJacobian(TS, PetscReal, Vec, Mat, Mat, void *);
61: extern PetscErrorCode RHSJacobianP_track(TS, PetscReal, Vec, Mat, void *);
62: extern PetscErrorCode RHSJacobianP_global(TS, PetscReal, Vec, Mat, void *);
63: extern PetscErrorCode Monitor(TS, PetscInt, PetscReal, Vec, void *);
64: extern PetscErrorCode AdjointMonitor(TS, PetscInt, PetscReal, Vec, PetscInt, Vec *, Vec *, void *);
66: /*
67: User-defined application context - contains data needed by the
68: application-provided call-back routines.
69: */
71: typedef struct {
72: /*------------- Forward solve data structures --------------*/
73: PetscInt max_steps; /* number of steps to run ts for */
74: PetscReal final_time; /* final time to integrate to*/
75: PetscReal time_step; /* ts integration time step */
76: Vec mu1; /* time dependent params */
77: Vec mu2; /* time dependent params */
78: Vec U; /* solution vector */
79: Mat A; /* Jacobian matrix */
81: /*------------- Adjoint solve data structures --------------*/
82: Mat Jacp; /* JacobianP matrix */
83: Vec lambda; /* adjoint variable */
84: Vec mup; /* adjoint variable */
86: /*------------- Global accumation vecs for monitor based tracking --------------*/
87: Vec sens_mu1; /* global sensitivity accumulation */
88: Vec sens_mu2; /* global sensitivity accumulation */
89: PetscInt adj_idx; /* to keep track of adjoint solve index */
90: } AppCtx;
92: typedef enum {
93: SA_TRACK,
94: SA_GLOBAL
95: } SAMethod;
96: static const char *const SAMethods[] = {"TRACK", "GLOBAL", "SAMethod", "SA_", NULL};
98: /* ----------------------- Explicit form of the ODE -------------------- */
100: PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, PetscCtx ctx)
101: {
102: AppCtx *user = (AppCtx *)ctx;
103: PetscScalar *f;
104: PetscInt curr_step;
105: const PetscScalar *u;
106: const PetscScalar *mu1;
107: const PetscScalar *mu2;
109: PetscFunctionBeginUser;
110: PetscCall(TSGetStepNumber(ts, &curr_step));
111: PetscCall(VecGetArrayRead(U, &u));
112: PetscCall(VecGetArrayRead(user->mu1, &mu1));
113: PetscCall(VecGetArrayRead(user->mu2, &mu2));
114: PetscCall(VecGetArray(F, &f));
115: f[0] = mu1[curr_step] * u[1];
116: f[1] = mu2[curr_step] * ((1. - u[0] * u[0]) * u[1] - u[0]);
117: PetscCall(VecRestoreArrayRead(U, &u));
118: PetscCall(VecRestoreArrayRead(user->mu1, &mu1));
119: PetscCall(VecRestoreArrayRead(user->mu2, &mu2));
120: PetscCall(VecRestoreArray(F, &f));
121: PetscFunctionReturn(PETSC_SUCCESS);
122: }
124: PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat A, Mat B, PetscCtx ctx)
125: {
126: AppCtx *user = (AppCtx *)ctx;
127: PetscInt rowcol[] = {0, 1};
128: PetscScalar J[2][2];
129: PetscInt curr_step;
130: const PetscScalar *u;
131: const PetscScalar *mu1;
132: const PetscScalar *mu2;
134: PetscFunctionBeginUser;
135: PetscCall(TSGetStepNumber(ts, &curr_step));
136: PetscCall(VecGetArrayRead(user->mu1, &mu1));
137: PetscCall(VecGetArrayRead(user->mu2, &mu2));
138: PetscCall(VecGetArrayRead(U, &u));
139: J[0][0] = 0;
140: J[1][0] = -mu2[curr_step] * (2.0 * u[1] * u[0] + 1.);
141: J[0][1] = mu1[curr_step];
142: J[1][1] = mu2[curr_step] * (1.0 - u[0] * u[0]);
143: PetscCall(MatSetValues(A, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
144: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
145: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
146: PetscCall(VecRestoreArrayRead(U, &u));
147: PetscCall(VecRestoreArrayRead(user->mu1, &mu1));
148: PetscCall(VecRestoreArrayRead(user->mu2, &mu2));
149: PetscFunctionReturn(PETSC_SUCCESS);
150: }
152: /* ------------------ Jacobian wrt parameters for tracking method ------------------ */
154: PetscErrorCode RHSJacobianP_track(TS ts, PetscReal t, Vec U, Mat A, PetscCtx ctx)
155: {
156: PetscInt row[] = {0, 1}, col[] = {0, 1};
157: PetscScalar J[2][2];
158: const PetscScalar *u;
160: PetscFunctionBeginUser;
161: PetscCall(VecGetArrayRead(U, &u));
162: J[0][0] = u[1];
163: J[1][0] = 0;
164: J[0][1] = 0;
165: J[1][1] = (1. - u[0] * u[0]) * u[1] - u[0];
166: PetscCall(MatSetValues(A, 2, row, 2, col, &J[0][0], INSERT_VALUES));
167: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
168: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
169: PetscCall(VecRestoreArrayRead(U, &u));
170: PetscFunctionReturn(PETSC_SUCCESS);
171: }
173: /* ------------------ Jacobian wrt parameters for global method ------------------ */
175: PetscErrorCode RHSJacobianP_global(TS ts, PetscReal t, Vec U, Mat A, PetscCtx ctx)
176: {
177: PetscInt row[] = {0, 1}, col[] = {0, 1};
178: PetscScalar J[2][2];
179: const PetscScalar *u;
180: PetscInt curr_step;
182: PetscFunctionBeginUser;
183: PetscCall(TSGetStepNumber(ts, &curr_step));
184: PetscCall(VecGetArrayRead(U, &u));
185: J[0][0] = u[1];
186: J[1][0] = 0;
187: J[0][1] = 0;
188: J[1][1] = (1. - u[0] * u[0]) * u[1] - u[0];
189: col[0] = curr_step * 2;
190: col[1] = curr_step * 2 + 1;
191: PetscCall(MatSetValues(A, 2, row, 2, col, &J[0][0], INSERT_VALUES));
192: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
193: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
194: PetscCall(VecRestoreArrayRead(U, &u));
195: PetscFunctionReturn(PETSC_SUCCESS);
196: }
198: /* Dump solution to console if called */
199: PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec U, PetscCtx ctx)
200: {
201: PetscFunctionBeginUser;
202: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Solution at time %e is \n", (double)t));
203: PetscCall(VecView(U, PETSC_VIEWER_STDOUT_WORLD));
204: PetscFunctionReturn(PETSC_SUCCESS);
205: }
207: /* Customized adjoint monitor to keep track of local
208: sensitivities by storing them in a global sensitivity array.
209: Note : This routine is only used for the tracking method. */
210: PetscErrorCode AdjointMonitor(TS ts, PetscInt steps, PetscReal time, Vec u, PetscInt numcost, Vec *lambda, Vec *mu, PetscCtx ctx)
211: {
212: AppCtx *user = (AppCtx *)ctx;
213: PetscInt curr_step;
214: PetscScalar *sensmu1_glob;
215: PetscScalar *sensmu2_glob;
216: const PetscScalar *sensmu_loc;
218: PetscFunctionBeginUser;
219: PetscCall(TSGetStepNumber(ts, &curr_step));
220: /* Note that we skip the first call to the monitor in the adjoint
221: solve since the sensitivities are already set (during
222: initialization of adjoint vectors).
223: We also note that each indvidial TSAdjointSolve calls the monitor
224: twice, once at the step it is integrating from and once at the step
225: it integrates to. Only the second call is useful for transferring
226: local sensitivities to the global array. */
227: if (curr_step == user->adj_idx) {
228: PetscFunctionReturn(PETSC_SUCCESS);
229: } else {
230: PetscCall(VecGetArrayRead(*mu, &sensmu_loc));
231: PetscCall(VecGetArray(user->sens_mu1, &sensmu1_glob));
232: PetscCall(VecGetArray(user->sens_mu2, &sensmu2_glob));
233: sensmu1_glob[curr_step] = sensmu_loc[0];
234: sensmu2_glob[curr_step] = sensmu_loc[1];
235: PetscCall(VecRestoreArray(user->sens_mu1, &sensmu1_glob));
236: PetscCall(VecRestoreArray(user->sens_mu2, &sensmu2_glob));
237: PetscCall(VecRestoreArrayRead(*mu, &sensmu_loc));
238: PetscFunctionReturn(PETSC_SUCCESS);
239: }
240: }
242: int main(int argc, char **argv)
243: {
244: TS ts;
245: AppCtx user;
246: PetscScalar *x_ptr, *y_ptr, *u_ptr;
247: PetscMPIInt size;
248: PetscBool monitor = PETSC_FALSE;
249: SAMethod sa = SA_GLOBAL;
251: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
252: Initialize program
253: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
254: PetscFunctionBeginUser;
255: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
256: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
257: PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
259: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
260: Set runtime options
261: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
262: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "SA analysis options.", "");
263: {
264: PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL));
265: PetscCall(PetscOptionsEnum("-sa_method", "Sensitivity analysis method (track or global)", "", SAMethods, (PetscEnum)sa, (PetscEnum *)&sa, NULL));
266: }
267: PetscOptionsEnd();
269: user.final_time = 0.1;
270: user.max_steps = 5;
271: user.time_step = user.final_time / user.max_steps;
273: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
274: Create necessary matrix and vectors for forward solve.
275: Create Jacp matrix for adjoint solve.
276: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
277: PetscCall(VecCreateSeq(PETSC_COMM_WORLD, user.max_steps, &user.mu1));
278: PetscCall(VecCreateSeq(PETSC_COMM_WORLD, user.max_steps, &user.mu2));
279: PetscCall(VecSet(user.mu1, 1.25));
280: PetscCall(VecSet(user.mu2, 1.0e2));
282: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
283: For tracking method : create the global sensitivity array to
284: accumulate sensitivity with respect to parameters at each step.
285: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
286: if (sa == SA_TRACK) {
287: PetscCall(VecCreateSeq(PETSC_COMM_WORLD, user.max_steps, &user.sens_mu1));
288: PetscCall(VecCreateSeq(PETSC_COMM_WORLD, user.max_steps, &user.sens_mu2));
289: }
291: PetscCall(MatCreate(PETSC_COMM_WORLD, &user.A));
292: PetscCall(MatSetSizes(user.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
293: PetscCall(MatSetFromOptions(user.A));
294: PetscCall(MatSetUp(user.A));
295: PetscCall(MatCreateVecs(user.A, &user.U, NULL));
297: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
298: Note that the dimensions of the Jacp matrix depend upon the
299: sensitivity analysis method being used !
300: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
301: PetscCall(MatCreate(PETSC_COMM_WORLD, &user.Jacp));
302: if (sa == SA_TRACK) PetscCall(MatSetSizes(user.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
303: if (sa == SA_GLOBAL) PetscCall(MatSetSizes(user.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, user.max_steps * 2));
304: PetscCall(MatSetFromOptions(user.Jacp));
305: PetscCall(MatSetUp(user.Jacp));
307: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
308: Create timestepping solver context
309: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
310: PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
311: PetscCall(TSSetEquationType(ts, TS_EQ_ODE_EXPLICIT));
312: PetscCall(TSSetType(ts, TSCN));
314: PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, &user));
315: PetscCall(TSSetRHSJacobian(ts, user.A, user.A, RHSJacobian, &user));
316: if (sa == SA_TRACK) PetscCall(TSSetRHSJacobianP(ts, user.Jacp, RHSJacobianP_track, &user));
317: if (sa == SA_GLOBAL) PetscCall(TSSetRHSJacobianP(ts, user.Jacp, RHSJacobianP_global, &user));
319: PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
320: PetscCall(TSSetMaxTime(ts, user.final_time));
321: PetscCall(TSSetTimeStep(ts, user.final_time / user.max_steps));
323: if (monitor) PetscCall(TSMonitorSet(ts, Monitor, &user, NULL));
324: if (sa == SA_TRACK) PetscCall(TSAdjointMonitorSet(ts, AdjointMonitor, &user, NULL));
326: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
327: Set initial conditions
328: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
329: PetscCall(VecGetArray(user.U, &x_ptr));
330: x_ptr[0] = 2.0;
331: x_ptr[1] = -2.0 / 3.0;
332: PetscCall(VecRestoreArray(user.U, &x_ptr));
334: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
335: Save trajectory of solution so that TSAdjointSolve() may be used
336: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
337: PetscCall(TSSetSaveTrajectory(ts));
339: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
340: Set runtime options
341: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
342: PetscCall(TSSetFromOptions(ts));
344: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
345: Execute forward model and print solution.
346: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
347: PetscCall(TSSolve(ts, user.U));
348: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Solution of forward TS :\n"));
349: PetscCall(VecView(user.U, PETSC_VIEWER_STDOUT_WORLD));
350: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Forward TS solve successful! Adjoint run begins!\n"));
352: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
353: Adjoint model starts here! Create adjoint vectors.
354: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
355: PetscCall(MatCreateVecs(user.A, &user.lambda, NULL));
356: PetscCall(MatCreateVecs(user.Jacp, &user.mup, NULL));
358: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
359: Set initial conditions for the adjoint vector
360: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
361: PetscCall(VecGetArray(user.U, &u_ptr));
362: PetscCall(VecGetArray(user.lambda, &y_ptr));
363: y_ptr[0] = 2 * (u_ptr[0] - 1.5967);
364: y_ptr[1] = 2 * (u_ptr[1] - -(1.02969));
365: PetscCall(VecRestoreArray(user.lambda, &y_ptr));
366: PetscCall(VecRestoreArray(user.U, &y_ptr));
367: PetscCall(VecSet(user.mup, 0));
369: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
370: Set number of cost functions.
371: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
372: PetscCall(TSSetCostGradients(ts, 1, &user.lambda, &user.mup));
374: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
375: The adjoint vector mup has to be reset for each adjoint step when
376: using the tracking method as we want to treat the parameters at each
377: time step one at a time and prevent accumulation of the sensitivities
378: from parameters at previous time steps.
379: This is not necessary for the global method as each time dependent
380: parameter is treated as an independent parameter.
381: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
382: if (sa == SA_TRACK) {
383: for (user.adj_idx = user.max_steps; user.adj_idx > 0; user.adj_idx--) {
384: PetscCall(VecSet(user.mup, 0));
385: PetscCall(TSAdjointSetSteps(ts, 1));
386: PetscCall(TSAdjointSolve(ts));
387: }
388: }
389: if (sa == SA_GLOBAL) PetscCall(TSAdjointSolve(ts));
391: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
392: Display adjoint sensitivities wrt parameters and initial conditions
393: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
394: if (sa == SA_TRACK) {
395: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n sensitivity wrt mu1: d[cost]/d[mu1]\n"));
396: PetscCall(VecView(user.sens_mu1, PETSC_VIEWER_STDOUT_WORLD));
397: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n sensitivity wrt mu2: d[cost]/d[mu2]\n"));
398: PetscCall(VecView(user.sens_mu2, PETSC_VIEWER_STDOUT_WORLD));
399: }
401: if (sa == SA_GLOBAL) {
402: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n sensitivity wrt params: d[cost]/d[p], where p refers to \nthe interlaced vector made by combining mu1,mu2\n"));
403: PetscCall(VecView(user.mup, PETSC_VIEWER_STDOUT_WORLD));
404: }
406: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n sensitivity wrt initial conditions: d[cost]/d[u(t=0)]\n"));
407: PetscCall(VecView(user.lambda, PETSC_VIEWER_STDOUT_WORLD));
409: /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
410: Free work space!
411: All PETSc objects should be destroyed when they are no longer needed.
412: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
413: PetscCall(MatDestroy(&user.A));
414: PetscCall(MatDestroy(&user.Jacp));
415: PetscCall(VecDestroy(&user.U));
416: PetscCall(VecDestroy(&user.lambda));
417: PetscCall(VecDestroy(&user.mup));
418: PetscCall(VecDestroy(&user.mu1));
419: PetscCall(VecDestroy(&user.mu2));
420: if (sa == SA_TRACK) {
421: PetscCall(VecDestroy(&user.sens_mu1));
422: PetscCall(VecDestroy(&user.sens_mu2));
423: }
424: PetscCall(TSDestroy(&ts));
425: PetscCall(PetscFinalize());
426: return 0;
427: }
429: /*TEST
431: test:
432: requires: !complex
433: suffix : track
434: args : -sa_method track
436: test:
437: requires: !complex
438: suffix : global
439: args : -sa_method global
441: TEST*/