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