xref: /petsc/src/ts/tutorials/hybrid/ex1adj.c (revision 030f984af8d8bb4c203755d35bded3c05b3d83ce)
1 static char help[] = "Adjoint sensitivity of a hybrid system with state-dependent switchings.\n";
2 
3 /*
4   The dynamics is described by the ODE
5                   u_t = A_i u
6 
7   where A_1 = [ 1  -100
8                 10  1  ],
9         A_2 = [ 1    10
10                -100  1 ].
11   The index i changes from 1 to 2 when u[1]=2.75u[0] and from 2 to 1 when u[1]=0.36u[0].
12   Initially u=[0 1]^T and i=1.
13 
14   References:
15   H. Zhang, S. Abhyankar, E. Constantinescu, M. Mihai, Discrete Adjoint Sensitivity Analysis of Hybrid Dynamical Systems With Switching, IEEE Transactions on Circuits and Systems I: Regular Papers, 64(5), May 2017
16   I. A. Hiskens, M.A. Pai, Trajectory Sensitivity Analysis of Hybrid Systems, IEEE Transactions on Circuits and Systems, Vol 47, No 2, February 2000
17 */
18 
19 #include <petscts.h>
20 
21 typedef struct {
22   PetscScalar lambda1;
23   PetscScalar lambda2;
24   PetscInt    mode;  /* mode flag*/
25 } AppCtx;
26 
27 PetscErrorCode EventFunction(TS ts,PetscReal t,Vec U,PetscScalar *fvalue,void *ctx)
28 {
29   AppCtx            *actx=(AppCtx*)ctx;
30   PetscErrorCode    ierr;
31   const PetscScalar *u;
32 
33   PetscFunctionBegin;
34   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
35   if (actx->mode == 1) {
36     fvalue[0] = u[1]-actx->lambda1*u[0];
37   }else if (actx->mode == 2) {
38     fvalue[0] = u[1]-actx->lambda2*u[0];
39   }
40   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
41   PetscFunctionReturn(0);
42 }
43 
44 PetscErrorCode ShiftGradients(TS ts,Vec U,AppCtx *actx)
45 {
46   Vec               *lambda,*mu;
47   PetscScalar       *x,*y;
48   const PetscScalar *u;
49   PetscErrorCode    ierr;
50   PetscScalar       tmp[2],A1[2][2],A2[2],denorm;
51   PetscInt          numcost;
52 
53   PetscFunctionBegin;
54   ierr = TSGetCostGradients(ts,&numcost,&lambda,&mu);CHKERRQ(ierr);
55   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
56 
57   if (actx->mode==2) {
58     denorm = -actx->lambda1*(u[0]-100.*u[1])+1.*(10.*u[0]+u[1]);
59     A1[0][0] = 110.*u[1]*(-actx->lambda1)/denorm+1.;
60     A1[0][1] = -110.*u[0]*(-actx->lambda1)/denorm;
61     A1[1][0] = 110.*u[1]*1./denorm;
62     A1[1][1] = -110.*u[0]*1./denorm+1.;
63 
64     A2[0] = 110.*u[1]*(-u[0])/denorm;
65     A2[1] = -110.*u[0]*(-u[0])/denorm;
66   } else {
67     denorm = -actx->lambda2*(u[0]+10.*u[1])+1.*(-100.*u[0]+u[1]);
68     A1[0][0] = 110.*u[1]*(actx->lambda2)/denorm+1;
69     A1[0][1] = -110.*u[0]*(actx->lambda2)/denorm;
70     A1[1][0] = -110.*u[1]*1./denorm;
71     A1[1][1] = 110.*u[0]*1./denorm+1.;
72 
73     A2[0] = 0;
74     A2[1] = 0;
75   }
76 
77   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
78 
79   ierr   = VecGetArray(lambda[0],&x);CHKERRQ(ierr);
80   ierr   = VecGetArray(mu[0],&y);CHKERRQ(ierr);
81   tmp[0] = A1[0][0]*x[0]+A1[0][1]*x[1];
82   tmp[1] = A1[1][0]*x[0]+A1[1][1]*x[1];
83   y[0]   = y[0] + A2[0]*x[0]+A2[1]*x[1];
84   x[0]   = tmp[0];
85   x[1]   = tmp[1];
86   ierr   = VecRestoreArray(mu[0],&y);CHKERRQ(ierr);
87   ierr   = VecRestoreArray(lambda[0],&x);CHKERRQ(ierr);
88 
89   ierr   = VecGetArray(lambda[1],&x);CHKERRQ(ierr);
90   ierr   = VecGetArray(mu[1],&y);CHKERRQ(ierr);
91   tmp[0] = A1[0][0]*x[0]+A1[0][1]*x[1];
92   tmp[1] = A1[1][0]*x[0]+A1[1][1]*x[1];
93   y[0]   = y[0] + A2[0]*x[0]+A2[1]*x[1];
94   x[0]   = tmp[0];
95   x[1]   = tmp[1];
96   ierr   = VecRestoreArray(mu[1],&y);CHKERRQ(ierr);
97   ierr   = VecRestoreArray(lambda[1],&x);CHKERRQ(ierr);
98   PetscFunctionReturn(0);
99 }
100 
101 PetscErrorCode PostEventFunction(TS ts,PetscInt nevents,PetscInt event_list[],PetscReal t,Vec U,PetscBool forwardsolve,void* ctx)
102 {
103   AppCtx         *actx=(AppCtx*)ctx;
104   PetscErrorCode ierr;
105 
106   PetscFunctionBegin;
107   /* ierr = VecView(U,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr); */
108   if (!forwardsolve) {
109     ierr = ShiftGradients(ts,U,actx);CHKERRQ(ierr);
110   }
111   if (actx->mode == 1) {
112     actx->mode = 2;
113     /* ierr = PetscPrintf(PETSC_COMM_SELF,"Change from mode 1 to 2 at t = %f \n",(double)t);CHKERRQ(ierr); */
114   } else if (actx->mode == 2) {
115     actx->mode = 1;
116     /* ierr = PetscPrintf(PETSC_COMM_SELF,"Change from mode 2 to 1 at t = %f \n",(double)t);CHKERRQ(ierr); */
117   }
118   PetscFunctionReturn(0);
119 }
120 
121 /*
122      Defines the ODE passed to the ODE solver
123 */
124 static PetscErrorCode IFunction(TS ts,PetscReal t,Vec U,Vec Udot,Vec F,void *ctx)
125 {
126   AppCtx            *actx=(AppCtx*)ctx;
127   PetscErrorCode    ierr;
128   PetscScalar       *f;
129   const PetscScalar *u,*udot;
130 
131   PetscFunctionBegin;
132   /*  The next three lines allow us to access the entries of the vectors directly */
133   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
134   ierr = VecGetArrayRead(Udot,&udot);CHKERRQ(ierr);
135   ierr = VecGetArray(F,&f);CHKERRQ(ierr);
136 
137   if (actx->mode == 1) {
138     f[0] = udot[0]-u[0]+100*u[1];
139     f[1] = udot[1]-10*u[0]-u[1];
140   } else if (actx->mode == 2) {
141     f[0] = udot[0]-u[0]-10*u[1];
142     f[1] = udot[1]+100*u[0]-u[1];
143   }
144 
145   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
146   ierr = VecRestoreArrayRead(Udot,&udot);CHKERRQ(ierr);
147   ierr = VecRestoreArray(F,&f);CHKERRQ(ierr);
148   PetscFunctionReturn(0);
149 }
150 
151 /*
152      Defines the Jacobian of the ODE passed to the ODE solver. See TSSetIJacobian() for the meaning of a and the Jacobian.
153 */
154 static PetscErrorCode IJacobian(TS ts,PetscReal t,Vec U,Vec Udot,PetscReal a,Mat A,Mat B,void *ctx)
155 {
156   AppCtx            *actx=(AppCtx*)ctx;
157   PetscErrorCode    ierr;
158   PetscInt          rowcol[] = {0,1};
159   PetscScalar       J[2][2];
160   const PetscScalar *u,*udot;
161 
162   PetscFunctionBegin;
163   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
164   ierr = VecGetArrayRead(Udot,&udot);CHKERRQ(ierr);
165 
166   if (actx->mode == 1) {
167     J[0][0] = a-1;                       J[0][1] = 100;
168     J[1][0] = -10;                       J[1][1] = a-1;
169   } else if (actx->mode == 2) {
170     J[0][0] = a-1;                       J[0][1] = -10;
171     J[1][0] = 100;                       J[1][1] = a-1;
172   }
173   ierr = MatSetValues(B,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);CHKERRQ(ierr);
174 
175   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
176   ierr = VecRestoreArrayRead(Udot,&udot);CHKERRQ(ierr);
177 
178   ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
179   ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
180   if (A != B) {
181     ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
182     ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
183   }
184   PetscFunctionReturn(0);
185 }
186 
187 /* Matrix JacobianP is constant so that it only needs to be evaluated once */
188 static PetscErrorCode RHSJacobianP(TS ts,PetscReal t,Vec X,Mat A, void *ctx)
189 {
190   PetscFunctionBeginUser;
191   PetscFunctionReturn(0);
192 }
193 
194 int main(int argc,char **argv)
195 {
196   TS             ts;            /* ODE integrator */
197   Vec            U;             /* solution will be stored here */
198   Mat            A;             /* Jacobian matrix */
199   Mat            Ap;            /* dfdp */
200   PetscErrorCode ierr;
201   PetscMPIInt    size;
202   PetscInt       n = 2;
203   PetscScalar    *u,*v;
204   AppCtx         app;
205   PetscInt       direction[1];
206   PetscBool      terminate[1];
207   Vec            lambda[2],mu[2];
208   PetscReal      tend;
209 
210   FILE           *f;
211   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
212      Initialize program
213      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
214   ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;
215   ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRMPI(ierr);
216   if (size > 1) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Only for sequential runs");
217   app.mode = 1;
218   app.lambda1 = 2.75;
219   app.lambda2 = 0.36;
220   tend = 0.125;
221   ierr = PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"ex1adj options","");CHKERRQ(ierr);
222   {
223     ierr = PetscOptionsReal("-lambda1","","",app.lambda1,&app.lambda1,NULL);CHKERRQ(ierr);
224     ierr = PetscOptionsReal("-lambda2","","",app.lambda2,&app.lambda2,NULL);CHKERRQ(ierr);
225     ierr = PetscOptionsReal("-tend","","",tend,&tend,NULL);CHKERRQ(ierr);
226   }
227   ierr = PetscOptionsEnd();CHKERRQ(ierr);
228 
229   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
230     Create necessary matrix and vectors
231     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
232   ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
233   ierr = MatSetSizes(A,n,n,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
234   ierr = MatSetType(A,MATDENSE);CHKERRQ(ierr);
235   ierr = MatSetFromOptions(A);CHKERRQ(ierr);
236   ierr = MatSetUp(A);CHKERRQ(ierr);
237 
238   ierr = MatCreateVecs(A,&U,NULL);CHKERRQ(ierr);
239 
240   ierr = MatCreate(PETSC_COMM_WORLD,&Ap);CHKERRQ(ierr);
241   ierr = MatSetSizes(Ap,n,1,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
242   ierr = MatSetType(Ap,MATDENSE);CHKERRQ(ierr);
243   ierr = MatSetFromOptions(Ap);CHKERRQ(ierr);
244   ierr = MatSetUp(Ap);CHKERRQ(ierr);
245   ierr = MatZeroEntries(Ap);CHKERRQ(ierr); /* initialize to zeros */
246 
247   ierr = VecGetArray(U,&u);CHKERRQ(ierr);
248   u[0] = 0;
249   u[1] = 1;
250   ierr = VecRestoreArray(U,&u);CHKERRQ(ierr);
251   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
252      Create timestepping solver context
253      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
254   ierr = TSCreate(PETSC_COMM_WORLD,&ts);CHKERRQ(ierr);
255   ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
256   ierr = TSSetType(ts,TSCN);CHKERRQ(ierr);
257   ierr = TSSetIFunction(ts,NULL,(TSIFunction)IFunction,&app);CHKERRQ(ierr);
258   ierr = TSSetIJacobian(ts,A,A,(TSIJacobian)IJacobian,&app);CHKERRQ(ierr);
259   ierr = TSSetRHSJacobianP(ts,Ap,RHSJacobianP,&app);CHKERRQ(ierr);
260 
261   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
262      Set initial conditions
263    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
264   ierr = TSSetSolution(ts,U);CHKERRQ(ierr);
265 
266   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
267     Save trajectory of solution so that TSAdjointSolve() may be used
268    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
269   ierr = TSSetSaveTrajectory(ts);CHKERRQ(ierr);
270 
271   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
272      Set solver options
273    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
274   ierr = TSSetMaxTime(ts,tend);CHKERRQ(ierr);
275   ierr = TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
276   ierr = TSSetTimeStep(ts,1./256.);CHKERRQ(ierr);
277   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
278 
279   /* Set directions and terminate flags for the two events */
280   direction[0] = 0;
281   terminate[0] = PETSC_FALSE;
282   ierr = TSSetEventHandler(ts,1,direction,terminate,EventFunction,PostEventFunction,(void*)&app);CHKERRQ(ierr);
283 
284   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
285      Run timestepping solver
286      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
287   ierr = TSSolve(ts,U);CHKERRQ(ierr);
288 
289   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
290      Adjoint model starts here
291      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
292   ierr = MatCreateVecs(A,&lambda[0],NULL);CHKERRQ(ierr);
293   ierr = MatCreateVecs(A,&lambda[1],NULL);CHKERRQ(ierr);
294   /*   Set initial conditions for the adjoint integration */
295   ierr = VecZeroEntries(lambda[0]);CHKERRQ(ierr);
296   ierr = VecZeroEntries(lambda[1]);CHKERRQ(ierr);
297   ierr = VecGetArray(lambda[0],&u);CHKERRQ(ierr);
298   u[0] = 1.;
299   ierr = VecRestoreArray(lambda[0],&u);CHKERRQ(ierr);
300   ierr = VecGetArray(lambda[1],&u);CHKERRQ(ierr);
301   u[1] = 1.;
302   ierr = VecRestoreArray(lambda[1],&u);CHKERRQ(ierr);
303 
304   ierr = MatCreateVecs(Ap,&mu[0],NULL);CHKERRQ(ierr);
305   ierr = MatCreateVecs(Ap,&mu[1],NULL);CHKERRQ(ierr);
306   ierr = VecZeroEntries(mu[0]);CHKERRQ(ierr);
307   ierr = VecZeroEntries(mu[1]);CHKERRQ(ierr);
308   ierr = TSSetCostGradients(ts,2,lambda,mu);CHKERRQ(ierr);
309 
310   ierr = TSAdjointSolve(ts);CHKERRQ(ierr);
311 
312   /*
313   ierr = VecView(lambda[0],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
314   ierr = VecView(lambda[1],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
315   ierr = VecView(mu[0],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
316   ierr = VecView(mu[1],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
317   */
318   ierr = VecGetArray(mu[0],&u);CHKERRQ(ierr);
319   ierr = VecGetArray(mu[1],&v);CHKERRQ(ierr);
320   f = fopen("adj_mu.out", "a");
321   ierr = PetscFPrintf(PETSC_COMM_WORLD,f,"%20.15lf %20.15lf %20.15lf\n",tend,u[0],v[0]);CHKERRQ(ierr);
322   ierr = VecRestoreArray(mu[0],&u);CHKERRQ(ierr);
323   ierr = VecRestoreArray(mu[1],&v);CHKERRQ(ierr);
324   fclose(f);
325   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
326      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
327    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
328   ierr = MatDestroy(&A);CHKERRQ(ierr);
329   ierr = VecDestroy(&U);CHKERRQ(ierr);
330   ierr = TSDestroy(&ts);CHKERRQ(ierr);
331 
332   ierr = MatDestroy(&Ap);CHKERRQ(ierr);
333   ierr = VecDestroy(&lambda[0]);CHKERRQ(ierr);
334   ierr = VecDestroy(&lambda[1]);CHKERRQ(ierr);
335   ierr = VecDestroy(&mu[0]);CHKERRQ(ierr);
336   ierr = VecDestroy(&mu[1]);CHKERRQ(ierr);
337   ierr = PetscFinalize();
338   return ierr;
339 }
340 
341 /*TEST
342 
343    build:
344       requires: !complex
345 
346    test:
347       args: -ts_monitor -ts_adjoint_monitor
348 
349 TEST*/
350