xref: /petsc/src/ts/tutorials/power_grid/ex3opt.c (revision 66af8762ec03dbef0e079729eb2a1734a35ed7ff)
1 static char help[] = "Finds optimal parameter P_m for the generator system while maintaining generator stability.\n";
2 
3 /*F
4 
5 \begin{eqnarray}
6                  \frac{d \theta}{dt} = \omega_b (\omega - \omega_s)
7                  \frac{2 H}{\omega_s}\frac{d \omega}{dt} & = & P_m - P_max \sin(\theta) -D(\omega - \omega_s)\\
8 \end{eqnarray}
9 
10 F*/
11 
12 /*
13   This code demonstrates how to solve a ODE-constrained optimization problem with TAO, TSEvent, TSAdjoint and TS.
14   The problem features discontinuities and a cost function in integral form.
15   The gradient is computed with the discrete adjoint of an implicit theta method, see ex3adj.c for details.
16 */
17 
18 #include <petsctao.h>
19 #include <petscts.h>
20 #include "ex3.h"
21 
22 PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
23 
24 PetscErrorCode monitor(Tao tao, AppCtx *ctx)
25 {
26   FILE              *fp;
27   PetscInt           iterate;
28   PetscReal          f, gnorm, cnorm, xdiff;
29   TaoConvergedReason reason;
30 
31   PetscFunctionBeginUser;
32   PetscCall(TaoGetSolutionStatus(tao, &iterate, &f, &gnorm, &cnorm, &xdiff, &reason));
33 
34   fp = fopen("ex3opt_conv.out", "a");
35   PetscCall(PetscFPrintf(PETSC_COMM_WORLD, fp, "%" PetscInt_FMT " %g\n", iterate, (double)gnorm));
36   fclose(fp);
37   PetscFunctionReturn(PETSC_SUCCESS);
38 }
39 
40 int main(int argc, char **argv)
41 {
42   Vec          p;
43   PetscScalar *x_ptr;
44   PetscMPIInt  size;
45   AppCtx       ctx;
46   Tao          tao;
47   KSP          ksp;
48   PC           pc;
49   Vec          lambda[1], mu[1], lowerb, upperb;
50   PetscBool    printtofile;
51   PetscInt     direction[2];
52   PetscBool    terminate[2];
53   Mat          qgrad; /* Forward sesivitiy */
54   Mat          sp;    /* Forward sensitivity matrix */
55 
56   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
57      Initialize program
58      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
59   PetscFunctionBeginUser;
60   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
61   PetscFunctionBeginUser;
62   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
63   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
64 
65   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
66     Set runtime options
67     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
68   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Swing equation options", "");
69   {
70     ctx.beta    = 2;
71     ctx.c       = 10000.0;
72     ctx.u_s     = 1.0;
73     ctx.omega_s = 1.0;
74     ctx.omega_b = 120.0 * PETSC_PI;
75     ctx.H       = 5.0;
76     PetscCall(PetscOptionsScalar("-Inertia", "", "", ctx.H, &ctx.H, NULL));
77     ctx.D = 5.0;
78     PetscCall(PetscOptionsScalar("-D", "", "", ctx.D, &ctx.D, NULL));
79     ctx.E        = 1.1378;
80     ctx.V        = 1.0;
81     ctx.X        = 0.545;
82     ctx.Pmax     = ctx.E * ctx.V / ctx.X;
83     ctx.Pmax_ini = ctx.Pmax;
84     PetscCall(PetscOptionsScalar("-Pmax", "", "", ctx.Pmax, &ctx.Pmax, NULL));
85     ctx.Pm = 1.06;
86     PetscCall(PetscOptionsScalar("-Pm", "", "", ctx.Pm, &ctx.Pm, NULL));
87     ctx.tf  = 0.1;
88     ctx.tcl = 0.2;
89     PetscCall(PetscOptionsReal("-tf", "Time to start fault", "", ctx.tf, &ctx.tf, NULL));
90     PetscCall(PetscOptionsReal("-tcl", "Time to end fault", "", ctx.tcl, &ctx.tcl, NULL));
91     printtofile = PETSC_FALSE;
92     PetscCall(PetscOptionsBool("-printtofile", "Print convergence results to file", "", printtofile, &printtofile, NULL));
93     ctx.sa = SA_ADJ;
94     PetscCall(PetscOptionsEnum("-sa_method", "Sensitivity analysis method (adj or tlm)", "", SAMethods, (PetscEnum)ctx.sa, (PetscEnum *)&ctx.sa, NULL));
95   }
96   PetscOptionsEnd();
97 
98   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
99     Create necessary matrix and vectors
100     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
101   PetscCall(MatCreate(PETSC_COMM_WORLD, &ctx.Jac));
102   PetscCall(MatSetSizes(ctx.Jac, 2, 2, PETSC_DETERMINE, PETSC_DETERMINE));
103   PetscCall(MatSetType(ctx.Jac, MATDENSE));
104   PetscCall(MatSetFromOptions(ctx.Jac));
105   PetscCall(MatSetUp(ctx.Jac));
106   PetscCall(MatCreate(PETSC_COMM_WORLD, &ctx.Jacp));
107   PetscCall(MatSetSizes(ctx.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 1));
108   PetscCall(MatSetFromOptions(ctx.Jacp));
109   PetscCall(MatSetUp(ctx.Jacp));
110   PetscCall(MatCreateVecs(ctx.Jac, &ctx.U, NULL));
111   PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, &ctx.DRDP));
112   PetscCall(MatSetUp(ctx.DRDP));
113   PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, &ctx.DRDU));
114   PetscCall(MatSetUp(ctx.DRDU));
115 
116   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
117      Create timestepping solver context
118      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
119   PetscCall(TSCreate(PETSC_COMM_WORLD, &ctx.ts));
120   PetscCall(TSSetProblemType(ctx.ts, TS_NONLINEAR));
121   PetscCall(TSSetType(ctx.ts, TSCN));
122   PetscCall(TSSetRHSFunction(ctx.ts, NULL, (TSRHSFunction)RHSFunction, &ctx));
123   PetscCall(TSSetRHSJacobian(ctx.ts, ctx.Jac, ctx.Jac, (TSRHSJacobian)RHSJacobian, &ctx));
124   PetscCall(TSSetRHSJacobianP(ctx.ts, ctx.Jacp, RHSJacobianP, &ctx));
125 
126   if (ctx.sa == SA_ADJ) {
127     PetscCall(MatCreateVecs(ctx.Jac, &lambda[0], NULL));
128     PetscCall(MatCreateVecs(ctx.Jacp, &mu[0], NULL));
129     PetscCall(TSSetSaveTrajectory(ctx.ts));
130     PetscCall(TSSetCostGradients(ctx.ts, 1, lambda, mu));
131     PetscCall(TSCreateQuadratureTS(ctx.ts, PETSC_FALSE, &ctx.quadts));
132     PetscCall(TSSetRHSFunction(ctx.quadts, NULL, (TSRHSFunction)CostIntegrand, &ctx));
133     PetscCall(TSSetRHSJacobian(ctx.quadts, ctx.DRDU, ctx.DRDU, (TSRHSJacobian)DRDUJacobianTranspose, &ctx));
134     PetscCall(TSSetRHSJacobianP(ctx.quadts, ctx.DRDP, DRDPJacobianTranspose, &ctx));
135   }
136   if (ctx.sa == SA_TLM) {
137     PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, &qgrad));
138     PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, &sp));
139     PetscCall(TSForwardSetSensitivities(ctx.ts, 1, sp));
140     PetscCall(TSCreateQuadratureTS(ctx.ts, PETSC_TRUE, &ctx.quadts));
141     PetscCall(TSForwardSetSensitivities(ctx.quadts, 1, qgrad));
142     PetscCall(TSSetRHSFunction(ctx.quadts, NULL, (TSRHSFunction)CostIntegrand, &ctx));
143     PetscCall(TSSetRHSJacobian(ctx.quadts, ctx.DRDU, ctx.DRDU, (TSRHSJacobian)DRDUJacobianTranspose, &ctx));
144     PetscCall(TSSetRHSJacobianP(ctx.quadts, ctx.DRDP, DRDPJacobianTranspose, &ctx));
145   }
146 
147   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
148      Set solver options
149    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
150   PetscCall(TSSetMaxTime(ctx.ts, 1.0));
151   PetscCall(TSSetExactFinalTime(ctx.ts, TS_EXACTFINALTIME_MATCHSTEP));
152   PetscCall(TSSetTimeStep(ctx.ts, 0.03125));
153   PetscCall(TSSetFromOptions(ctx.ts));
154 
155   direction[0] = direction[1] = 1;
156   terminate[0] = terminate[1] = PETSC_FALSE;
157   PetscCall(TSSetEventHandler(ctx.ts, 2, direction, terminate, EventFunction, PostEventFunction, &ctx));
158 
159   /* Create TAO solver and set desired solution method */
160   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
161   PetscCall(TaoSetType(tao, TAOBLMVM));
162   if (printtofile) PetscCall(TaoSetMonitor(tao, (PetscErrorCode(*)(Tao, void *))monitor, (void *)&ctx, NULL));
163   /*
164      Optimization starts
165   */
166   /* Set initial solution guess */
167   PetscCall(VecCreateSeq(PETSC_COMM_WORLD, 1, &p));
168   PetscCall(VecGetArray(p, &x_ptr));
169   x_ptr[0] = ctx.Pm;
170   PetscCall(VecRestoreArray(p, &x_ptr));
171 
172   PetscCall(TaoSetSolution(tao, p));
173   /* Set routine for function and gradient evaluation */
174   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&ctx));
175 
176   /* Set bounds for the optimization */
177   PetscCall(VecDuplicate(p, &lowerb));
178   PetscCall(VecDuplicate(p, &upperb));
179   PetscCall(VecGetArray(lowerb, &x_ptr));
180   x_ptr[0] = 0.;
181   PetscCall(VecRestoreArray(lowerb, &x_ptr));
182   PetscCall(VecGetArray(upperb, &x_ptr));
183   x_ptr[0] = 1.1;
184   PetscCall(VecRestoreArray(upperb, &x_ptr));
185   PetscCall(TaoSetVariableBounds(tao, lowerb, upperb));
186 
187   /* Check for any TAO command line options */
188   PetscCall(TaoSetFromOptions(tao));
189   PetscCall(TaoGetKSP(tao, &ksp));
190   if (ksp) {
191     PetscCall(KSPGetPC(ksp, &pc));
192     PetscCall(PCSetType(pc, PCNONE));
193   }
194 
195   /* SOLVE THE APPLICATION */
196   PetscCall(TaoSolve(tao));
197 
198   PetscCall(VecView(p, PETSC_VIEWER_STDOUT_WORLD));
199 
200   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
201      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
202    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
203   PetscCall(MatDestroy(&ctx.Jac));
204   PetscCall(MatDestroy(&ctx.Jacp));
205   PetscCall(MatDestroy(&ctx.DRDU));
206   PetscCall(MatDestroy(&ctx.DRDP));
207   PetscCall(VecDestroy(&ctx.U));
208   if (ctx.sa == SA_ADJ) {
209     PetscCall(VecDestroy(&lambda[0]));
210     PetscCall(VecDestroy(&mu[0]));
211   }
212   if (ctx.sa == SA_TLM) {
213     PetscCall(MatDestroy(&qgrad));
214     PetscCall(MatDestroy(&sp));
215   }
216   PetscCall(TSDestroy(&ctx.ts));
217   PetscCall(VecDestroy(&p));
218   PetscCall(VecDestroy(&lowerb));
219   PetscCall(VecDestroy(&upperb));
220   PetscCall(TaoDestroy(&tao));
221   PetscCall(PetscFinalize());
222   return 0;
223 }
224 
225 /* ------------------------------------------------------------------ */
226 /*
227    FormFunctionGradient - Evaluates the function and corresponding gradient.
228 
229    Input Parameters:
230    tao - the Tao context
231    X   - the input vector
232    ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()
233 
234    Output Parameters:
235    f   - the newly evaluated function
236    G   - the newly evaluated gradient
237 */
238 PetscErrorCode FormFunctionGradient(Tao tao, Vec P, PetscReal *f, Vec G, void *ctx0)
239 {
240   AppCtx      *ctx = (AppCtx *)ctx0;
241   PetscInt     nadj;
242   PetscReal    ftime;
243   PetscInt     steps;
244   PetscScalar *u;
245   PetscScalar *x_ptr, *y_ptr;
246   Vec          q;
247   Mat          qgrad;
248 
249   PetscFunctionBeginUser;
250   PetscCall(VecGetArrayRead(P, (const PetscScalar **)&x_ptr));
251   ctx->Pm = x_ptr[0];
252   PetscCall(VecRestoreArrayRead(P, (const PetscScalar **)&x_ptr));
253 
254   /* reinitialize the solution vector */
255   PetscCall(VecGetArray(ctx->U, &u));
256   u[0] = PetscAsinScalar(ctx->Pm / ctx->Pmax);
257   u[1] = 1.0;
258   PetscCall(VecRestoreArray(ctx->U, &u));
259   PetscCall(TSSetSolution(ctx->ts, ctx->U));
260 
261   /* reset time */
262   PetscCall(TSSetTime(ctx->ts, 0.0));
263 
264   /* reset step counter, this is critical for adjoint solver */
265   PetscCall(TSSetStepNumber(ctx->ts, 0));
266 
267   /* reset step size, the step size becomes negative after TSAdjointSolve */
268   PetscCall(TSSetTimeStep(ctx->ts, 0.03125));
269 
270   /* reinitialize the integral value */
271   PetscCall(TSGetQuadratureTS(ctx->ts, NULL, &ctx->quadts));
272   PetscCall(TSGetSolution(ctx->quadts, &q));
273   PetscCall(VecSet(q, 0.0));
274 
275   if (ctx->sa == SA_TLM) { /* reset the forward sensitivities */
276     TS             quadts;
277     Mat            sp;
278     PetscScalar    val[2];
279     const PetscInt row[] = {0, 1}, col[] = {0};
280 
281     PetscCall(TSGetQuadratureTS(ctx->ts, NULL, &quadts));
282     PetscCall(TSForwardGetSensitivities(quadts, NULL, &qgrad));
283     PetscCall(MatZeroEntries(qgrad));
284     PetscCall(TSForwardGetSensitivities(ctx->ts, NULL, &sp));
285     val[0] = 1. / PetscSqrtScalar(1. - (ctx->Pm / ctx->Pmax) * (ctx->Pm / ctx->Pmax)) / ctx->Pmax;
286     val[1] = 0.0;
287     PetscCall(MatSetValues(sp, 2, row, 1, col, val, INSERT_VALUES));
288     PetscCall(MatAssemblyBegin(sp, MAT_FINAL_ASSEMBLY));
289     PetscCall(MatAssemblyEnd(sp, MAT_FINAL_ASSEMBLY));
290   }
291 
292   /* solve the ODE */
293   PetscCall(TSSolve(ctx->ts, ctx->U));
294   PetscCall(TSGetSolveTime(ctx->ts, &ftime));
295   PetscCall(TSGetStepNumber(ctx->ts, &steps));
296 
297   if (ctx->sa == SA_ADJ) {
298     Vec *lambda, *mu;
299     /* reset the terminal condition for adjoint */
300     PetscCall(TSGetCostGradients(ctx->ts, &nadj, &lambda, &mu));
301     PetscCall(VecGetArray(lambda[0], &y_ptr));
302     y_ptr[0] = 0.0;
303     y_ptr[1] = 0.0;
304     PetscCall(VecRestoreArray(lambda[0], &y_ptr));
305     PetscCall(VecGetArray(mu[0], &x_ptr));
306     x_ptr[0] = -1.0;
307     PetscCall(VecRestoreArray(mu[0], &x_ptr));
308 
309     /* solve the adjont */
310     PetscCall(TSAdjointSolve(ctx->ts));
311 
312     PetscCall(ComputeSensiP(lambda[0], mu[0], ctx));
313     PetscCall(VecCopy(mu[0], G));
314   }
315 
316   if (ctx->sa == SA_TLM) {
317     PetscCall(VecGetArray(G, &x_ptr));
318     PetscCall(MatDenseGetArray(qgrad, &y_ptr));
319     x_ptr[0] = y_ptr[0] - 1.;
320     PetscCall(MatDenseRestoreArray(qgrad, &y_ptr));
321     PetscCall(VecRestoreArray(G, &x_ptr));
322   }
323 
324   PetscCall(TSGetSolution(ctx->quadts, &q));
325   PetscCall(VecGetArray(q, &x_ptr));
326   *f = -ctx->Pm + x_ptr[0];
327   PetscCall(VecRestoreArray(q, &x_ptr));
328   PetscFunctionReturn(PETSC_SUCCESS);
329 }
330 
331 /*TEST
332 
333    build:
334       requires: !complex !single
335 
336    test:
337       args: -viewer_binary_skip_info -ts_type cn -pc_type lu -tao_monitor
338 
339    test:
340       suffix: 2
341       output_file: output/ex3opt_1.out
342       args: -sa_method tlm -ts_type cn -pc_type lu -tao_monitor
343 TEST*/
344