Actual source code: tssen.c

  1: #include <petsc/private/tsimpl.h>
  2: #include <petscdraw.h>

  4: PetscLogEvent TS_AdjointStep, TS_ForwardStep, TS_JacobianPEval;

  6: /* #define TSADJOINT_STAGE */

  8: /* ------------------------ Sensitivity Context ---------------------------*/

 10: /*@
 11:   TSSetRHSJacobianP - Sets the function that computes the Jacobian of $G$ w.r.t. the parameters $p$ where $U_t = G(U,p,t)$, as well as the location to store the matrix.

 13:   Logically Collective

 15:   Input Parameters:
 16: + ts   - `TS` context obtained from `TSCreate()`
 17: . Amat - JacobianP matrix
 18: . func - function
 19: - ctx  - [optional] function context

 21:   Level: intermediate

 23:   Notes:
 24:   `Amat` has the same number of rows and the same row parallel layout as `u`, `Amat` has the same number of columns and parallel layout as `p`

 26:   When `TSSetIJacobianP()` is also called, the two must be given different matrices since each holds a separate term of the
 27:   parameter Jacobian; sharing one is an error.

 29: .seealso: [](ch_ts), `TS`, `TSRHSJacobianPFn`, `TSGetRHSJacobianP()`, `TSSetIJacobianP()`
 30: @*/
 31: PetscErrorCode TSSetRHSJacobianP(TS ts, Mat Amat, TSRHSJacobianPFn *func, PetscCtx ctx)
 32: {
 33:   PetscFunctionBegin;
 36:   /* ts->Jacp may legitimately alias ts->Jacprhs after TSSetUp() when only this routine was called, so a shared matrix is
 37:      rejected only once an IJacobianP exists whose separate term would be stored in it */
 38:   PetscCheck(!ts->ijacobianp || Amat != ts->Jacp, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_WRONGSTATE, "TSSetIJacobianP() and TSSetRHSJacobianP() must be given different matrices");

 40:   ts->rhsjacobianp    = func;
 41:   ts->rhsjacobianpctx = ctx;
 42:   if (Amat) {
 43:     PetscCall(PetscObjectReference((PetscObject)Amat));
 44:     PetscCall(MatDestroy(&ts->Jacprhs));
 45:     ts->Jacprhs = Amat;
 46:   }
 47:   PetscFunctionReturn(PETSC_SUCCESS);
 48: }

 50: /*@
 51:   TSGetRHSJacobianP - Gets the function that computes the Jacobian of $G $ w.r.t. the parameters $p$ where $ U_t = G(U,p,t)$, as well as the location to store the matrix.

 53:   Logically Collective

 55:   Input Parameter:
 56: . ts - `TS` context obtained from `TSCreate()`

 58:   Output Parameters:
 59: + Amat - JacobianP matrix
 60: . func - function
 61: - ctx  - [optional] function context

 63:   Level: intermediate

 65:   Note:
 66:   `Amat` has the same number of rows and the same row parallel layout as `u`, `Amat` has the same number of columns and parallel layout as `p`

 68: .seealso: [](ch_ts), `TSSetRHSJacobianP()`, `TS`, `TSRHSJacobianPFn`
 69: @*/
 70: PetscErrorCode TSGetRHSJacobianP(TS ts, Mat *Amat, TSRHSJacobianPFn **func, PetscCtxRt ctx)
 71: {
 72:   PetscFunctionBegin;
 73:   if (func) *func = ts->rhsjacobianp;
 74:   if (ctx) *(void **)ctx = ts->rhsjacobianpctx;
 75:   if (Amat) *Amat = ts->Jacprhs;
 76:   PetscFunctionReturn(PETSC_SUCCESS);
 77: }

 79: /*@
 80:   TSComputeRHSJacobianP - Runs the user-defined JacobianP function.

 82:   Collective

 84:   Input Parameters:
 85: + ts - The `TS` context obtained from `TSCreate()`
 86: . t  - the time
 87: - U  - the solution at which to compute the Jacobian

 89:   Output Parameter:
 90: . Amat - the computed Jacobian

 92:   Level: developer

 94: .seealso: [](ch_ts), `TSSetRHSJacobianP()`, `TS`
 95: @*/
 96: PetscErrorCode TSComputeRHSJacobianP(TS ts, PetscReal t, Vec U, Mat Amat)
 97: {
 98:   PetscFunctionBegin;
 99:   if (!Amat) PetscFunctionReturn(PETSC_SUCCESS);

103:   if (ts->rhsjacobianp) PetscCallBack("TS callback JacobianP for sensitivity analysis", (*ts->rhsjacobianp)(ts, t, U, Amat, ts->rhsjacobianpctx));
104:   else {
105:     PetscBool assembled;
106:     PetscCall(MatZeroEntries(Amat));
107:     PetscCall(MatAssembled(Amat, &assembled));
108:     if (!assembled) {
109:       PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
110:       PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
111:     }
112:   }
113:   PetscFunctionReturn(PETSC_SUCCESS);
114: }

116: /*@
117:   TSSetIJacobianP - Sets the function that computes the Jacobian of $F$ w.r.t. the parameters $p$ where $F(Udot,U,p,t) = G(U,p,t)$, as well as the location to store the matrix.

119:   Logically Collective

121:   Input Parameters:
122: + ts   - `TS` context obtained from `TSCreate()`
123: . Amat - JacobianP matrix
124: . func - function
125: - ctx  - [optional] function context

127:   Calling sequence of `func`:
128: + ts    - the `TS` context
129: . t     - current timestep
130: . U     - input vector (current ODE solution)
131: . Udot  - time derivative of state vector
132: . shift - shift to apply, see the note in `TSSetIJacobian()`
133: . A     - output matrix
134: - ctx   - [optional] function context

136:   Level: intermediate

138:   Notes:
139:   `Amat` has the same number of rows and the same row parallel layout as `u`, `Amat` has the same number of columns and parallel layout as `p`

141:   When `TSSetRHSJacobianP()` is also called, the two must be given different matrices since each holds a separate term of the
142:   parameter Jacobian; sharing one is an error.

144: .seealso: [](ch_ts), `TSSetRHSJacobianP()`, `TS`
145: @*/
146: PetscErrorCode TSSetIJacobianP(TS ts, Mat Amat, PetscErrorCode (*func)(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal shift, Mat A, PetscCtx ctx), PetscCtx ctx)
147: {
148:   PetscFunctionBegin;
151:   /* ts->Jacprhs is only ever the RHSJacobianP matrix, so reusing it here means the two parameter Jacobian terms would
152:      clobber each other, unless no callback is registered and there is no second term to store */
153:   PetscCheck(!func || Amat != ts->Jacprhs, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_WRONGSTATE, "TSSetIJacobianP() and TSSetRHSJacobianP() must be given different matrices");

155:   ts->ijacobianp    = func;
156:   ts->ijacobianpctx = ctx;
157:   if (Amat) {
158:     PetscCall(PetscObjectReference((PetscObject)Amat));
159:     PetscCall(MatDestroy(&ts->Jacp));
160:     ts->Jacp = Amat;
161:   }
162:   PetscFunctionReturn(PETSC_SUCCESS);
163: }

165: /*@
166:   TSGetIJacobianP - Gets the function that computes the Jacobian of $ F$ w.r.t. the parameters $p$ where $F(Udot,U,p,t) = G(U,p,t) $, as well as the location to store the matrix.

168:   Logically Collective

170:   Input Parameter:
171: . ts - `TS` context obtained from `TSCreate()`

173:   Output Parameters:
174: + Amat - JacobianP matrix
175: . func - the function that computes the JacobianP
176: - ctx  - [optional] function context

178:   Calling sequence of `func`:
179: + ts    - the `TS` context
180: . t     - current timestep
181: . U     - input vector (current ODE solution)
182: . Udot  - time derivative of state vector
183: . shift - shift to apply, see the note in `TSSetIJacobian()`
184: . A     - output matrix
185: - ctx   - [optional] function context

187:   Level: intermediate

189:   Note:
190:   `Amat` has the same number of rows and the same row parallel layout as `u`, `Amat` has the same number of columns and parallel layout as `p`

192: .seealso: [](ch_ts), `TSSetRHSJacobianP()`, `TS`, `TSSetIJacobianP()`, `TSGetRHSJacobianP()`
193: @*/
194: PetscErrorCode TSGetIJacobianP(TS ts, Mat *Amat, PetscErrorCode (**func)(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal shift, Mat A, PetscCtx ctx), PetscCtxRt ctx)
195: {
196:   PetscFunctionBegin;

199:   if (func) *func = ts->ijacobianp;
200:   if (ctx) *(void **)ctx = ts->ijacobianpctx;
201:   if (Amat) *Amat = ts->Jacp;
202:   PetscFunctionReturn(PETSC_SUCCESS);
203: }

205: /*@
206:   TSComputeIJacobianP - Runs the user-defined IJacobianP function.

208:   Collective

210:   Input Parameters:
211: + ts    - the `TS` context
212: . t     - current timestep
213: . U     - state vector
214: . Udot  - time derivative of state vector
215: . shift - shift to apply, see note below
216: - imex  - flag indicates if the method is IMEX so that the `RHSJacobianP` should be kept separate

218:   Output Parameter:
219: . Amat - Jacobian matrix

221:   Level: developer

223: .seealso: [](ch_ts), `TS`, `TSSetIJacobianP()`
224: @*/
225: PetscErrorCode TSComputeIJacobianP(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal shift, Mat Amat, PetscBool imex)
226: {
227:   PetscFunctionBegin;
228:   if (!Amat) PetscFunctionReturn(PETSC_SUCCESS);

233:   PetscCall(PetscLogEventBegin(TS_JacobianPEval, ts, U, Amat, 0));
234:   if (ts->ijacobianp) PetscCallBack("TS callback JacobianP for sensitivity analysis", (*ts->ijacobianp)(ts, t, U, Udot, shift, Amat, ts->ijacobianpctx));
235:   else { /* system was written as Udot = G(t,U), so the implicit term is zero; Amat must still be cleared because it can
236:             hold values left by an earlier registration or by the previous call on shared RHS storage */
237:     PetscBool assembled;

239:     PetscCall(MatZeroEntries(Amat));
240:     PetscCall(MatAssembled(Amat, &assembled));
241:     if (!assembled) {
242:       PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
243:       PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
244:     }
245:   }
246:   if (!imex) {
247:     if (ts->rhsjacobianp) PetscCall(TSComputeRHSJacobianP(ts, t, U, ts->Jacprhs));
248:     if (ts->Jacprhs == Amat) { /* No IJacobian, so we only have the RHS matrix */
249:       PetscCall(MatScale(Amat, -1));
250:     } else if (ts->Jacprhs) { /* Both IJacobian and RHSJacobian */
251:       MatStructure axpy = DIFFERENT_NONZERO_PATTERN;

253:       PetscCall(MatAXPY(Amat, -1, ts->Jacprhs, axpy));
254:     }
255:   }
256:   PetscCall(PetscLogEventEnd(TS_JacobianPEval, ts, U, Amat, 0));
257:   PetscFunctionReturn(PETSC_SUCCESS);
258: }

260: /*@
261:   TSSetCostIntegrand - Sets the routine for evaluating the integral term in one or more cost functions

263:   Logically Collective

265:   Input Parameters:
266: + ts           - the `TS` context obtained from `TSCreate()`
267: . numcost      - number of gradients to be computed, this is the number of cost functions
268: . costintegral - vector that stores the integral values
269: . rf           - routine for evaluating the integrand function
270: . drduf        - function that computes the gradients of the `r` with respect to `u`
271: . drdpf        - function that computes the gradients of the `r` with respect to p, can be `NULL` if parametric sensitivity is not desired (`mu` = `NULL`)
272: . fwd          - flag indicating whether to evaluate cost integral in the forward run or the adjoint run
273: - ctx          - [optional] application context for private data for the function evaluation routine (may be `NULL`)

275:   Calling sequence of `rf`:
276: + ts  - the integrator
277: . t   - the time
278: . U   - the solution
279: . F   - the computed value of the function
280: - ctx - the application context

282:   Calling sequence of `drduf`:
283: + ts   - the integrator
284: . t    - the time
285: . U    - the solution
286: . dRdU - the computed gradients of the `r` with respect to `u`
287: - ctx  - the application context

289:   Calling sequence of `drdpf`:
290: + ts   - the integrator
291: . t    - the time
292: . U    - the solution
293: . dRdP - the computed gradients of the `r` with respect to `p`
294: - ctx  - the application context

296:   Level: deprecated

298:   Notes:
299:   For optimization there is usually a single cost function (numcost = 1). For sensitivities there may be multiple cost functions

301:   Use `TSCreateQuadratureTS()` and `TSForwardSetSensitivities()` instead

303: .seealso: [](ch_ts), `TS`, `TSSetRHSJacobianP()`, `TSGetCostGradients()`, `TSSetCostGradients()`,
304:           `TSCreateQuadratureTS()`, `TSForwardSetSensitivities()`
305: @*/
306: PetscErrorCode TSSetCostIntegrand(TS ts, PetscInt numcost, Vec costintegral, PetscErrorCode (*rf)(TS ts, PetscReal t, Vec U, Vec F, PetscCtx ctx), PetscErrorCode (*drduf)(TS ts, PetscReal t, Vec U, Vec *dRdU, PetscCtx ctx), PetscErrorCode (*drdpf)(TS ts, PetscReal t, Vec U, Vec *dRdP, PetscCtx ctx), PetscBool fwd, PetscCtx ctx)
307: {
308:   PetscFunctionBegin;
311:   PetscCheck(!ts->numcost || ts->numcost == numcost, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "The number of cost functions (2nd parameter of TSSetCostIntegrand()) is inconsistent with the one set by TSSetCostGradients() or TSForwardSetIntegralGradients()");
312:   if (!ts->numcost) ts->numcost = numcost;

314:   if (costintegral) {
315:     PetscCall(PetscObjectReference((PetscObject)costintegral));
316:     PetscCall(VecDestroy(&ts->vec_costintegral));
317:     ts->vec_costintegral = costintegral;
318:   } else {
319:     if (!ts->vec_costintegral) { /* Create a seq vec if user does not provide one */
320:       PetscCall(VecCreateSeq(PETSC_COMM_SELF, numcost, &ts->vec_costintegral));
321:     } else {
322:       PetscCall(VecSet(ts->vec_costintegral, 0.0));
323:     }
324:   }
325:   if (!ts->vec_costintegrand) {
326:     PetscCall(VecDuplicate(ts->vec_costintegral, &ts->vec_costintegrand));
327:   } else {
328:     PetscCall(VecSet(ts->vec_costintegrand, 0.0));
329:   }
330:   ts->costintegralfwd  = fwd; /* Evaluate the cost integral in forward run if fwd is true */
331:   ts->costintegrand    = rf;
332:   ts->costintegrandctx = ctx;
333:   ts->drdufunction     = drduf;
334:   ts->drdpfunction     = drdpf;
335:   PetscFunctionReturn(PETSC_SUCCESS);
336: }

338: /*@
339:   TSGetCostIntegral - Returns the values of the integral term in the cost functions.
340:   It is valid to call the routine after a backward run.

342:   Not Collective

344:   Input Parameter:
345: . ts - the `TS` context obtained from `TSCreate()`

347:   Output Parameter:
348: . v - the vector containing the integrals for each cost function

350:   Level: intermediate

352: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSSetCostIntegrand()`
353: @*/
354: PetscErrorCode TSGetCostIntegral(TS ts, Vec *v)
355: {
356:   TS quadts;

358:   PetscFunctionBegin;
360:   PetscAssertPointer(v, 2);
361:   PetscCall(TSGetQuadratureTS(ts, NULL, &quadts));
362:   *v = quadts->vec_sol;
363:   PetscFunctionReturn(PETSC_SUCCESS);
364: }

366: /*@
367:   TSComputeCostIntegrand - Evaluates the integral function in the cost functions.

369:   Input Parameters:
370: + ts - the `TS` context
371: . t  - current time
372: - U  - state vector, i.e. current solution

374:   Output Parameter:
375: . Q - vector of size numcost to hold the outputs

377:   Level: deprecated

379:   Note:
380:   Most users should not need to explicitly call this routine, as it
381:   is used internally within the sensitivity analysis context.

383: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSSetCostIntegrand()`
384: @*/
385: PetscErrorCode TSComputeCostIntegrand(TS ts, PetscReal t, Vec U, Vec Q)
386: {
387:   PetscFunctionBegin;

392:   PetscCall(PetscLogEventBegin(TS_FunctionEval, ts, U, Q, 0));
393:   if (ts->costintegrand) PetscCallBack("TS callback integrand in the cost function", (*ts->costintegrand)(ts, t, U, Q, ts->costintegrandctx));
394:   else PetscCall(VecZeroEntries(Q));
395:   PetscCall(PetscLogEventEnd(TS_FunctionEval, ts, U, Q, 0));
396:   PetscFunctionReturn(PETSC_SUCCESS);
397: }

399: // PetscClangLinter pragma disable: -fdoc-*
400: /*@
401:   TSComputeDRDUFunction - Deprecated, use `TSGetQuadratureTS()` then `TSComputeRHSJacobian()`

403:   Level: deprecated

405: @*/
406: PetscErrorCode TSComputeDRDUFunction(TS ts, PetscReal t, Vec U, Vec *DRDU)
407: {
408:   PetscFunctionBegin;
409:   if (!DRDU) PetscFunctionReturn(PETSC_SUCCESS);

413:   PetscCallBack("TS callback DRDU for sensitivity analysis", (*ts->drdufunction)(ts, t, U, DRDU, ts->costintegrandctx));
414:   PetscFunctionReturn(PETSC_SUCCESS);
415: }

417: // PetscClangLinter pragma disable: -fdoc-*
418: /*@
419:   TSComputeDRDPFunction - Deprecated, use `TSGetQuadratureTS()` then `TSComputeRHSJacobianP()`

421:   Level: deprecated

423: @*/
424: PetscErrorCode TSComputeDRDPFunction(TS ts, PetscReal t, Vec U, Vec *DRDP)
425: {
426:   PetscFunctionBegin;
427:   if (!DRDP) PetscFunctionReturn(PETSC_SUCCESS);

431:   PetscCallBack("TS callback DRDP for sensitivity analysis", (*ts->drdpfunction)(ts, t, U, DRDP, ts->costintegrandctx));
432:   PetscFunctionReturn(PETSC_SUCCESS);
433: }

435: // PetscClangLinter pragma disable: -fdoc-param-list-func-parameter-documentation
436: // PetscClangLinter pragma disable: -fdoc-section-header-unknown
437: /*@
438:   TSSetIHessianProduct - Sets the function that computes the vector-Hessian-vector product. The Hessian is the second-order derivative of `F` (IFunction) w.r.t. the state variable.

440:   Logically Collective

442:   Input Parameters:
443: + ts                   - `TS` context obtained from `TSCreate()`
444: . ihp1                 - an array of vectors storing the result of vector-Hessian-vector product for $F_{UU}$
445: . ihessianproductfunc1 - vector-Hessian-vector product function for $F_{UU}$
446: . ihp2                 - an array of vectors storing the result of vector-Hessian-vector product for $F_{UP}$
447: . ihessianproductfunc2 - vector-Hessian-vector product function for $F_{UP}$
448: . ihp3                 - an array of vectors storing the result of vector-Hessian-vector product for $F_{PU}$
449: . ihessianproductfunc3 - vector-Hessian-vector product function for $F_{PU}$
450: . ihp4                 - an array of vectors storing the result of vector-Hessian-vector product for $F_{PP}$
451: . ihessianproductfunc4 - vector-Hessian-vector product function for $F_{PP}$
452: - ctx                  - [optional] function context

454:   Calling sequence of `ihessianproductfunc1`:
455: + ts  - the `TS` context
456: . t   - current timestep
457: . U   - input vector (current ODE solution)
458: . Vl  - an array of input vectors to be left-multiplied with the Hessian
459: . Vr  - input vector to be right-multiplied with the Hessian
460: . VHV - an array of output vectors for vector-Hessian-vector product
461: - ctx - [optional] function context

463:   Level: intermediate

465:   Notes:
466:   All other functions have the same calling sequence as `ihessianproductfunc1`, so their
467:   descriptions are omitted for brevity.

469:   The first Hessian function and the working array are required.
470:   As an example to implement the callback functions, the second callback function calculates the vector-Hessian-vector product
471:   $Vl_n^T*F_UP*Vr$
472:   where the vector $Vl_n$ (n-th element in the array `Vl`) and `Vr` are of size `N` and `M` respectively, and the Hessian $F_{UP}$ is of size $N x N x M.$
473:   Each entry of $F_{UP}$ corresponds to the derivative
474:   $ F_UP[i][j][k] = \frac{\partial^2 F[i]}{\partial U[j] \partial P[k]}.$
475:   The result of the vector-Hessian-vector product for $Vl_n$ needs to be stored in vector $VHV_n$ with the j-th entry being
476:   $ VHV_n[j] = \sum_i \sum_k {Vl_n[i] * F_UP[i][j][k] * Vr[k]}$
477:   If the cost function is a scalar, there will be only one vector in `Vl` and `VHV`.

479: .seealso: [](ch_ts), `TS`
480: @*/
481: PetscErrorCode TSSetIHessianProduct(TS ts, Vec ihp1[], PetscErrorCode (*ihessianproductfunc1)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec ihp2[], PetscErrorCode (*ihessianproductfunc2)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec ihp3[], PetscErrorCode (*ihessianproductfunc3)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec ihp4[], PetscErrorCode (*ihessianproductfunc4)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), PetscCtx ctx)
482: {
483:   PetscFunctionBegin;
485:   PetscAssertPointer(ihp1, 2);

487:   ts->ihessianproductctx = ctx;
488:   if (ihp1) ts->vecs_fuu = ihp1;
489:   if (ihp2) ts->vecs_fup = ihp2;
490:   if (ihp3) ts->vecs_fpu = ihp3;
491:   if (ihp4) ts->vecs_fpp = ihp4;
492:   ts->ihessianproduct_fuu = ihessianproductfunc1;
493:   ts->ihessianproduct_fup = ihessianproductfunc2;
494:   ts->ihessianproduct_fpu = ihessianproductfunc3;
495:   ts->ihessianproduct_fpp = ihessianproductfunc4;
496:   PetscFunctionReturn(PETSC_SUCCESS);
497: }

499: /*@
500:   TSComputeIHessianProductFunctionUU - Runs the user-defined vector-Hessian-vector product function for Fuu.

502:   Collective

504:   Input Parameters:
505: + ts - The `TS` context obtained from `TSCreate()`
506: . t  - the time
507: . U  - the solution at which to compute the Hessian product
508: . Vl - the array of input vectors to be multiplied with the Hessian from the left
509: - Vr - the input vector to be multiplied with the Hessian from the right

511:   Output Parameter:
512: . VHV - the array of output vectors that store the Hessian product

514:   Level: developer

516:   Note:
517:   `TSComputeIHessianProductFunctionUU()` is typically used for sensitivity implementation,
518:   so most users would not generally call this routine themselves.

520: .seealso: [](ch_ts), `TSSetIHessianProduct()`
521: @*/
522: PetscErrorCode TSComputeIHessianProductFunctionUU(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
523: {
524:   PetscFunctionBegin;
525:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

529:   if (ts->ihessianproduct_fuu) PetscCallBack("TS callback IHessianProduct 1 for sensitivity analysis", (*ts->ihessianproduct_fuu)(ts, t, U, Vl, Vr, VHV, ts->ihessianproductctx));

531:   /* does not consider IMEX for now, so either IHessian or RHSHessian will be calculated, using the same output VHV */
532:   if (ts->rhshessianproduct_guu) {
533:     PetscInt nadj;
534:     PetscCall(TSComputeRHSHessianProductFunctionUU(ts, t, U, Vl, Vr, VHV));
535:     for (nadj = 0; nadj < ts->numcost; nadj++) PetscCall(VecScale(VHV[nadj], -1));
536:   }
537:   PetscFunctionReturn(PETSC_SUCCESS);
538: }

540: /*@
541:   TSComputeIHessianProductFunctionUP - Runs the user-defined vector-Hessian-vector product function for Fup.

543:   Collective

545:   Input Parameters:
546: + ts - The `TS` context obtained from `TSCreate()`
547: . t  - the time
548: . U  - the solution at which to compute the Hessian product
549: . Vl - the array of input vectors to be multiplied with the Hessian from the left
550: - Vr - the input vector to be multiplied with the Hessian from the right

552:   Output Parameter:
553: . VHV - the array of output vectors that store the Hessian product

555:   Level: developer

557:   Note:
558:   `TSComputeIHessianProductFunctionUP()` is typically used for sensitivity implementation,
559:   so most users would not generally call this routine themselves.

561: .seealso: [](ch_ts), `TSSetIHessianProduct()`
562: @*/
563: PetscErrorCode TSComputeIHessianProductFunctionUP(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
564: {
565:   PetscFunctionBegin;
566:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

570:   if (ts->ihessianproduct_fup) PetscCallBack("TS callback IHessianProduct 2 for sensitivity analysis", (*ts->ihessianproduct_fup)(ts, t, U, Vl, Vr, VHV, ts->ihessianproductctx));

572:   /* does not consider IMEX for now, so either IHessian or RHSHessian will be calculated, using the same output VHV */
573:   if (ts->rhshessianproduct_gup) {
574:     PetscCall(TSComputeRHSHessianProductFunctionUP(ts, t, U, Vl, Vr, VHV));
575:     for (PetscInt nadj = 0; nadj < ts->numcost; nadj++) PetscCall(VecScale(VHV[nadj], -1));
576:   }
577:   PetscFunctionReturn(PETSC_SUCCESS);
578: }

580: /*@
581:   TSComputeIHessianProductFunctionPU - Runs the user-defined vector-Hessian-vector product function for Fpu.

583:   Collective

585:   Input Parameters:
586: + ts - The `TS` context obtained from `TSCreate()`
587: . t  - the time
588: . U  - the solution at which to compute the Hessian product
589: . Vl - the array of input vectors to be multiplied with the Hessian from the left
590: - Vr - the input vector to be multiplied with the Hessian from the right

592:   Output Parameter:
593: . VHV - the array of output vectors that store the Hessian product

595:   Level: developer

597:   Note:
598:   `TSComputeIHessianProductFunctionPU()` is typically used for sensitivity implementation,
599:   so most users would not generally call this routine themselves.

601: .seealso: [](ch_ts), `TSSetIHessianProduct()`
602: @*/
603: PetscErrorCode TSComputeIHessianProductFunctionPU(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
604: {
605:   PetscFunctionBegin;
606:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

610:   if (ts->ihessianproduct_fpu) PetscCallBack("TS callback IHessianProduct 3 for sensitivity analysis", (*ts->ihessianproduct_fpu)(ts, t, U, Vl, Vr, VHV, ts->ihessianproductctx));

612:   /* does not consider IMEX for now, so either IHessian or RHSHessian will be calculated, using the same output VHV */
613:   if (ts->rhshessianproduct_gpu) {
614:     PetscCall(TSComputeRHSHessianProductFunctionPU(ts, t, U, Vl, Vr, VHV));
615:     for (PetscInt nadj = 0; nadj < ts->numcost; nadj++) PetscCall(VecScale(VHV[nadj], -1));
616:   }
617:   PetscFunctionReturn(PETSC_SUCCESS);
618: }

620: /*@
621:   TSComputeIHessianProductFunctionPP - Runs the user-defined vector-Hessian-vector product function for Fpp.

623:   Collective

625:   Input Parameters:
626: + ts - The `TS` context obtained from `TSCreate()`
627: . t  - the time
628: . U  - the solution at which to compute the Hessian product
629: . Vl - the array of input vectors to be multiplied with the Hessian from the left
630: - Vr - the input vector to be multiplied with the Hessian from the right

632:   Output Parameter:
633: . VHV - the array of output vectors that store the Hessian product

635:   Level: developer

637:   Note:
638:   `TSComputeIHessianProductFunctionPP()` is typically used for sensitivity implementation,
639:   so most users would not generally call this routine themselves.

641: .seealso: [](ch_ts), `TSSetIHessianProduct()`
642: @*/
643: PetscErrorCode TSComputeIHessianProductFunctionPP(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
644: {
645:   PetscFunctionBegin;
646:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

650:   if (ts->ihessianproduct_fpp) PetscCallBack("TS callback IHessianProduct 3 for sensitivity analysis", (*ts->ihessianproduct_fpp)(ts, t, U, Vl, Vr, VHV, ts->ihessianproductctx));

652:   /* does not consider IMEX for now, so either IHessian or RHSHessian will be calculated, using the same output VHV */
653:   if (ts->rhshessianproduct_gpp) {
654:     PetscCall(TSComputeRHSHessianProductFunctionPP(ts, t, U, Vl, Vr, VHV));
655:     for (PetscInt nadj = 0; nadj < ts->numcost; nadj++) PetscCall(VecScale(VHV[nadj], -1));
656:   }
657:   PetscFunctionReturn(PETSC_SUCCESS);
658: }

660: // PetscClangLinter pragma disable: -fdoc-param-list-func-parameter-documentation
661: // PetscClangLinter pragma disable: -fdoc-section-header-unknown
662: /*@
663:   TSSetRHSHessianProduct - Sets the function that computes the vector-Hessian-vector
664:   product. The Hessian is the second-order derivative of `G` (RHSFunction) w.r.t. the state
665:   variable.

667:   Logically Collective

669:   Input Parameters:
670: + ts                     - `TS` context obtained from `TSCreate()`
671: . rhshp1                 - an array of vectors storing the result of vector-Hessian-vector product for $G_{UU}$
672: . rhshessianproductfunc1 - vector-Hessian-vector product function for $G_{UU}$
673: . rhshp2                 - an array of vectors storing the result of vector-Hessian-vector product for $G_{UP}$
674: . rhshessianproductfunc2 - vector-Hessian-vector product function for $G_{UP}$
675: . rhshp3                 - an array of vectors storing the result of vector-Hessian-vector product for $G_{PU}$
676: . rhshessianproductfunc3 - vector-Hessian-vector product function for $G_{PU}$
677: . rhshp4                 - an array of vectors storing the result of vector-Hessian-vector product for $G_{PP}$
678: . rhshessianproductfunc4 - vector-Hessian-vector product function for $G_{PP}$
679: - ctx                    - [optional] function context

681:   Calling sequence of `rhshessianproductfunc1`:
682: + ts  - the `TS` context
683: . t   - current timestep
684: . U   - input vector (current ODE solution)
685: . Vl  - an array of input vectors to be left-multiplied with the Hessian
686: . Vr  - input vector to be right-multiplied with the Hessian
687: . VHV - an array of output vectors for vector-Hessian-vector product
688: - ctx - [optional] function context

690:   Level: intermediate

692:   Notes:
693:   All other functions have the same calling sequence as `rhshessianproductfunc1`, so their
694:   descriptions are omitted for brevity.

696:   The first Hessian function and the working array are required.

698:   As an example to implement the callback functions, the second callback function calculates the vector-Hessian-vector product
699:   $ Vl_n^T*G_UP*Vr$
700:   where the vector $Vl_n$ (n-th element in the array $Vl$) and $Vr$ are of size $N$ and $M$ respectively, and the Hessian $G_{UP}$ is of size $N x N x M$.
701:   Each entry of $G_{UP}$ corresponds to the derivative
702:   $ G_UP[i][j][k] = \frac{\partial^2 G[i]}{\partial U[j] \partial P[k]}.$
703:   The result of the vector-Hessian-vector product for $Vl_n$ needs to be stored in vector $VHV_n$ with j-th entry being
704:   $ VHV_n[j] = \sum_i \sum_k {Vl_n[i] * G_UP[i][j][k] * Vr[k]}$
705:   If the cost function is a scalar, there will be only one vector in $Vl$ and $VHV$.

707: .seealso: `TS`, `TSAdjoint`
708: @*/
709: PetscErrorCode TSSetRHSHessianProduct(TS ts, Vec rhshp1[], PetscErrorCode (*rhshessianproductfunc1)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec rhshp2[], PetscErrorCode (*rhshessianproductfunc2)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec rhshp3[], PetscErrorCode (*rhshessianproductfunc3)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), Vec rhshp4[], PetscErrorCode (*rhshessianproductfunc4)(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[], PetscCtx ctx), PetscCtx ctx)
710: {
711:   PetscFunctionBegin;
713:   PetscAssertPointer(rhshp1, 2);

715:   ts->rhshessianproductctx = ctx;
716:   if (rhshp1) ts->vecs_guu = rhshp1;
717:   if (rhshp2) ts->vecs_gup = rhshp2;
718:   if (rhshp3) ts->vecs_gpu = rhshp3;
719:   if (rhshp4) ts->vecs_gpp = rhshp4;
720:   ts->rhshessianproduct_guu = rhshessianproductfunc1;
721:   ts->rhshessianproduct_gup = rhshessianproductfunc2;
722:   ts->rhshessianproduct_gpu = rhshessianproductfunc3;
723:   ts->rhshessianproduct_gpp = rhshessianproductfunc4;
724:   PetscFunctionReturn(PETSC_SUCCESS);
725: }

727: /*@
728:   TSComputeRHSHessianProductFunctionUU - Runs the user-defined vector-Hessian-vector product function for $G_{uu}$.

730:   Collective

732:   Input Parameters:
733: + ts - The `TS` context obtained from `TSCreate()`
734: . t  - the time
735: . U  - the solution at which to compute the Hessian product
736: . Vl - the array of input vectors to be multiplied with the Hessian from the left
737: - Vr - the input vector to be multiplied with the Hessian from the right

739:   Output Parameter:
740: . VHV - the array of output vectors that store the Hessian product

742:   Level: developer

744:   Note:
745:   `TSComputeRHSHessianProductFunctionUU()` is typically used for sensitivity implementation,
746:   so most users would not generally call this routine themselves.

748: .seealso: [](ch_ts), `TS`, `TSSetRHSHessianProduct()`
749: @*/
750: PetscErrorCode TSComputeRHSHessianProductFunctionUU(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
751: {
752:   PetscFunctionBegin;
753:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

757:   PetscCallBack("TS callback RHSHessianProduct 1 for sensitivity analysis", (*ts->rhshessianproduct_guu)(ts, t, U, Vl, Vr, VHV, ts->rhshessianproductctx));
758:   PetscFunctionReturn(PETSC_SUCCESS);
759: }

761: /*@
762:   TSComputeRHSHessianProductFunctionUP - Runs the user-defined vector-Hessian-vector product function for $G_{up}$.

764:   Collective

766:   Input Parameters:
767: + ts - The `TS` context obtained from `TSCreate()`
768: . t  - the time
769: . U  - the solution at which to compute the Hessian product
770: . Vl - the array of input vectors to be multiplied with the Hessian from the left
771: - Vr - the input vector to be multiplied with the Hessian from the right

773:   Output Parameter:
774: . VHV - the array of output vectors that store the Hessian product

776:   Level: developer

778:   Note:
779:   `TSComputeRHSHessianProductFunctionUP()` is typically used for sensitivity implementation,
780:   so most users would not generally call this routine themselves.

782: .seealso: [](ch_ts), `TS`, `TSSetRHSHessianProduct()`
783: @*/
784: PetscErrorCode TSComputeRHSHessianProductFunctionUP(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
785: {
786:   PetscFunctionBegin;
787:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

791:   PetscCallBack("TS callback RHSHessianProduct 2 for sensitivity analysis", (*ts->rhshessianproduct_gup)(ts, t, U, Vl, Vr, VHV, ts->rhshessianproductctx));
792:   PetscFunctionReturn(PETSC_SUCCESS);
793: }

795: /*@
796:   TSComputeRHSHessianProductFunctionPU - Runs the user-defined vector-Hessian-vector product function for $G_{pu}.$

798:   Collective

800:   Input Parameters:
801: + ts - The `TS` context obtained from `TSCreate()`
802: . t  - the time
803: . U  - the solution at which to compute the Hessian product
804: . Vl - the array of input vectors to be multiplied with the Hessian from the left
805: - Vr - the input vector to be multiplied with the Hessian from the right

807:   Output Parameter:
808: . VHV - the array of output vectors that store the Hessian product

810:   Level: developer

812:   Note:
813:   `TSComputeRHSHessianProductFunctionPU()` is typically used for sensitivity implementation,
814:   so most users would not generally call this routine themselves.

816: .seealso: [](ch_ts), `TSSetRHSHessianProduct()`
817: @*/
818: PetscErrorCode TSComputeRHSHessianProductFunctionPU(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
819: {
820:   PetscFunctionBegin;
821:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

825:   PetscCallBack("TS callback RHSHessianProduct 3 for sensitivity analysis", (*ts->rhshessianproduct_gpu)(ts, t, U, Vl, Vr, VHV, ts->rhshessianproductctx));
826:   PetscFunctionReturn(PETSC_SUCCESS);
827: }

829: /*@
830:   TSComputeRHSHessianProductFunctionPP - Runs the user-defined vector-Hessian-vector product function for $G_{pp}$.

832:   Collective

834:   Input Parameters:
835: + ts - The `TS` context obtained from `TSCreate()`
836: . t  - the time
837: . U  - the solution at which to compute the Hessian product
838: . Vl - the array of input vectors to be multiplied with the Hessian from the left
839: - Vr - the input vector to be multiplied with the Hessian from the right

841:   Output Parameter:
842: . VHV - the array of output vectors that store the Hessian product

844:   Level: developer

846:   Note:
847:   `TSComputeRHSHessianProductFunctionPP()` is typically used for sensitivity implementation,
848:   so most users would not generally call this routine themselves.

850: .seealso: [](ch_ts), `TSSetRHSHessianProduct()`
851: @*/
852: PetscErrorCode TSComputeRHSHessianProductFunctionPP(TS ts, PetscReal t, Vec U, Vec Vl[], Vec Vr, Vec VHV[])
853: {
854:   PetscFunctionBegin;
855:   if (!VHV) PetscFunctionReturn(PETSC_SUCCESS);

859:   PetscCallBack("TS callback RHSHessianProduct 3 for sensitivity analysis", (*ts->rhshessianproduct_gpp)(ts, t, U, Vl, Vr, VHV, ts->rhshessianproductctx));
860:   PetscFunctionReturn(PETSC_SUCCESS);
861: }

863: /* --------------------------- Adjoint sensitivity ---------------------------*/

865: /*@
866:   TSSetCostGradients - Sets the initial value of the gradients of the cost function w.r.t. initial values and w.r.t. the problem parameters
867:   for use by the `TS` adjoint routines.

869:   Logically Collective

871:   Input Parameters:
872: + ts      - the `TS` context obtained from `TSCreate()`
873: . numcost - number of gradients to be computed, this is the number of cost functions
874: . lambda  - gradients with respect to the initial condition variables, the dimension and parallel layout of these vectors is the same as the ODE solution vector
875: - mu      - gradients with respect to the parameters, the number of entries in these vectors is the same as the number of parameters

877:   Level: beginner

879:   Notes:
880:   the entries in these vectors must be correctly initialized with the values lambda_i = df/dy|finaltime  mu_i = df/dp|finaltime

882:   After `TSAdjointSolve()` is called the lambda and the mu contain the computed sensitivities

884: .seealso: `TS`, `TSAdjointSolve()`, `TSGetCostGradients()`
885: @*/
886: PetscErrorCode TSSetCostGradients(TS ts, PetscInt numcost, Vec lambda[], Vec mu[])
887: {
888:   PetscFunctionBegin;
890:   PetscAssertPointer(lambda, 3);
891:   ts->vecs_sensi  = lambda;
892:   ts->vecs_sensip = mu;
893:   PetscCheck(!ts->numcost || ts->numcost == numcost, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "The number of cost functions (2nd parameter of TSSetCostIntegrand()) is inconsistent with the one set by TSSetCostIntegrand");
894:   ts->numcost = numcost;
895:   PetscFunctionReturn(PETSC_SUCCESS);
896: }

898: /*@
899:   TSGetCostGradients - Returns the gradients from the `TSAdjointSolve()`

901:   Not Collective, but the vectors returned are parallel if `TS` is parallel

903:   Input Parameter:
904: . ts - the `TS` context obtained from `TSCreate()`

906:   Output Parameters:
907: + numcost - size of returned arrays
908: . lambda  - vectors containing the gradients of the cost functions with respect to the ODE/DAE solution variables
909: - mu      - vectors containing the gradients of the cost functions with respect to the problem parameters

911:   Level: intermediate

913: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSSetCostGradients()`
914: @*/
915: PetscErrorCode TSGetCostGradients(TS ts, PetscInt *numcost, Vec *lambda[], Vec *mu[])
916: {
917:   PetscFunctionBegin;
919:   if (numcost) *numcost = ts->numcost;
920:   if (lambda) *lambda = ts->vecs_sensi;
921:   if (mu) *mu = ts->vecs_sensip;
922:   PetscFunctionReturn(PETSC_SUCCESS);
923: }

925: /*@
926:   TSSetCostHessianProducts - Sets the initial value of the Hessian-vector products of the cost function w.r.t. initial values and w.r.t. the problem parameters
927:   for use by the `TS` adjoint routines.

929:   Logically Collective

931:   Input Parameters:
932: + ts      - the `TS` context obtained from `TSCreate()`
933: . numcost - number of cost functions
934: . lambda2 - Hessian-vector product with respect to the initial condition variables, the dimension and parallel layout of these vectors is the same as the ODE solution vector
935: . mu2     - Hessian-vector product with respect to the parameters, the number of entries in these vectors is the same as the number of parameters
936: - dir     - the direction vector that are multiplied with the Hessian of the cost functions

938:   Level: beginner

940:   Notes:
941:   Hessian of the cost function is completely different from Hessian of the ODE/DAE system

943:   For second-order adjoint, one needs to call this function and then `TSAdjointSetForward()` before `TSSolve()`.

945:   After `TSAdjointSolve()` is called, the lambda2 and the mu2 will contain the computed second-order adjoint sensitivities, and can be used to produce Hessian-vector product (not the full Hessian matrix). Users must provide a direction vector; it is usually generated by an optimization solver.

947:   Passing `NULL` for `lambda2` disables the second-order calculation.

949: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSAdjointSetForward()`
950: @*/
951: PetscErrorCode TSSetCostHessianProducts(TS ts, PetscInt numcost, Vec lambda2[], Vec mu2[], Vec dir)
952: {
953:   PetscFunctionBegin;
955:   PetscCheck(!ts->numcost || ts->numcost == numcost, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "The number of cost functions (2nd parameter of TSSetCostIntegrand()) is inconsistent with the one set by TSSetCostIntegrand");
956:   ts->numcost      = numcost;
957:   ts->vecs_sensi2  = lambda2;
958:   ts->vecs_sensi2p = mu2;
959:   ts->vec_dir      = dir;
960:   PetscFunctionReturn(PETSC_SUCCESS);
961: }

963: /*@
964:   TSGetCostHessianProducts - Returns the gradients from the `TSAdjointSolve()`

966:   Not Collective, but vectors returned are parallel if `TS` is parallel

968:   Input Parameter:
969: . ts - the `TS` context obtained from `TSCreate()`

971:   Output Parameters:
972: + numcost - number of cost functions
973: . lambda2 - Hessian-vector product with respect to the initial condition variables, the dimension and parallel layout of these vectors is the same as the ODE solution vector
974: . mu2     - Hessian-vector product with respect to the parameters, the number of entries in these vectors is the same as the number of parameters
975: - dir     - the direction vector that are multiplied with the Hessian of the cost functions

977:   Level: intermediate

979: .seealso: [](ch_ts), `TSAdjointSolve()`, `TSSetCostHessianProducts()`
980: @*/
981: PetscErrorCode TSGetCostHessianProducts(TS ts, PetscInt *numcost, Vec *lambda2[], Vec *mu2[], Vec *dir)
982: {
983:   PetscFunctionBegin;
985:   if (numcost) *numcost = ts->numcost;
986:   if (lambda2) *lambda2 = ts->vecs_sensi2;
987:   if (mu2) *mu2 = ts->vecs_sensi2p;
988:   if (dir) *dir = ts->vec_dir;
989:   PetscFunctionReturn(PETSC_SUCCESS);
990: }

992: /*@
993:   TSAdjointSetForward - Trigger the tangent linear solver and initialize the forward sensitivities

995:   Logically Collective

997:   Input Parameters:
998: + ts   - the `TS` context obtained from `TSCreate()`
999: - didp - the derivative of initial values w.r.t. parameters

1001:   Level: intermediate

1003:   Notes:
1004:   When computing sensitivities w.r.t. initial condition, set didp to `NULL` so that the solver will take it as an identity matrix mathematically.
1005:   `TSAdjoint` does not reset the tangent linear solver automatically, `TSAdjointResetForward()` should be called to reset the tangent linear solver.

1007: .seealso: [](ch_ts), `TSAdjointSolve()`, `TSSetCostHessianProducts()`, `TSAdjointResetForward()`
1008: @*/
1009: PetscErrorCode TSAdjointSetForward(TS ts, Mat didp)
1010: {
1011:   Mat          A;
1012:   Vec          sp;
1013:   PetscScalar *xarr;
1014:   PetscInt     lsize;

1016:   PetscFunctionBegin;
1017:   ts->forward_solve = PETSC_TRUE; /* turn on tangent linear mode */
1018:   PetscCheck(ts->vecs_sensi2, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "Must call TSSetCostHessianProducts() first");
1019:   PetscCheck(ts->vec_dir, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "Directional vector is missing. Call TSSetCostHessianProducts() to set it.");
1020:   /* create a single-column dense matrix */
1021:   PetscCall(VecGetLocalSize(ts->vec_sol, &lsize));
1022:   PetscCall(MatCreateDense(PetscObjectComm((PetscObject)ts), lsize, PETSC_DECIDE, PETSC_DECIDE, 1, NULL, &A));

1024:   PetscCall(VecDuplicate(ts->vec_sol, &sp));
1025:   PetscCall(MatDenseGetColumn(A, 0, &xarr));
1026:   PetscCall(VecPlaceArray(sp, xarr));
1027:   if (ts->vecs_sensi2p) { /* tangent linear variable initialized as 2*dIdP*dir */
1028:     if (didp) {
1029:       PetscCall(MatMult(didp, ts->vec_dir, sp));
1030:       PetscCall(VecScale(sp, 2.));
1031:     } else {
1032:       PetscCall(VecZeroEntries(sp));
1033:     }
1034:   } else { /* tangent linear variable initialized as dir */
1035:     PetscCall(VecCopy(ts->vec_dir, sp));
1036:   }
1037:   PetscCall(VecResetArray(sp));
1038:   PetscCall(MatDenseRestoreColumn(A, &xarr));
1039:   PetscCall(VecDestroy(&sp));

1041:   PetscCall(TSForwardSetInitialSensitivities(ts, A)); /* if didp is NULL, identity matrix is assumed */

1043:   PetscCall(MatDestroy(&A));
1044:   PetscFunctionReturn(PETSC_SUCCESS);
1045: }

1047: /*@
1048:   TSAdjointResetForward - Reset the tangent linear solver and destroy the tangent linear context

1050:   Logically Collective

1052:   Input Parameter:
1053: . ts - the `TS` context obtained from `TSCreate()`

1055:   Level: intermediate

1057: .seealso: [](ch_ts), `TSAdjointSetForward()`
1058: @*/
1059: PetscErrorCode TSAdjointResetForward(TS ts)
1060: {
1061:   PetscFunctionBegin;
1062:   ts->forward_solve = PETSC_FALSE; /* turn off tangent linear mode */
1063:   PetscCall(TSForwardReset(ts));
1064:   PetscFunctionReturn(PETSC_SUCCESS);
1065: }

1067: /*@
1068:   TSAdjointSetUp - Sets up the internal data structures for the later use
1069:   of an adjoint solver

1071:   Collective

1073:   Input Parameter:
1074: . ts - the `TS` context obtained from `TSCreate()`

1076:   Level: advanced

1078: .seealso: [](ch_ts), `TSCreate()`, `TSAdjointStep()`, `TSSetCostGradients()`
1079: @*/
1080: PetscErrorCode TSAdjointSetUp(TS ts)
1081: {
1082:   TSTrajectory tj;
1083:   PetscBool    match;

1085:   PetscFunctionBegin;
1087:   if (ts->adjointsetupcalled) PetscFunctionReturn(PETSC_SUCCESS);
1088:   PetscCheck(ts->vecs_sensi, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_WRONGSTATE, "Must call TSSetCostGradients() first");
1089:   PetscCheck(!ts->vecs_sensip || ts->Jacp || ts->Jacprhs, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_WRONGSTATE, "Must call TSSetRHSJacobianP() or TSSetIJacobianP() first");
1090:   PetscCall(TSGetTrajectory(ts, &tj));
1091:   PetscCall(PetscObjectTypeCompare((PetscObject)tj, TSTRAJECTORYBASIC, &match));
1092:   if (match) {
1093:     PetscBool solution_only;
1094:     PetscCall(TSTrajectoryGetSolutionOnly(tj, &solution_only));
1095:     PetscCheck(!solution_only, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "TSAdjoint cannot use the solution-only mode when choosing the Basic TSTrajectory type. Turn it off with -ts_trajectory_solution_only 0");
1096:   }
1097:   PetscCall(TSTrajectorySetUseHistory(tj, PETSC_FALSE)); /* not use TSHistory */

1099:   if (ts->quadraturets) { /* if there is integral in the cost function */
1100:     PetscCall(VecDuplicate(ts->vecs_sensi[0], &ts->vec_drdu_col));
1101:     if (ts->vecs_sensip) PetscCall(VecDuplicate(ts->vecs_sensip[0], &ts->vec_drdp_col));
1102:   }

1104:   PetscTryTypeMethod(ts, adjointsetup);
1105:   ts->adjointsetupcalled = PETSC_TRUE;
1106:   PetscFunctionReturn(PETSC_SUCCESS);
1107: }

1109: /*@
1110:   TSAdjointReset - Resets a `TS` adjoint context and removes any allocated `Vec`s and `Mat`s.

1112:   Collective

1114:   Input Parameter:
1115: . ts - the `TS` context obtained from `TSCreate()`

1117:   Level: beginner

1119: .seealso: [](ch_ts), `TSCreate()`, `TSAdjointSetUp()`, `TSDestroy()`
1120: @*/
1121: PetscErrorCode TSAdjointReset(TS ts)
1122: {
1123:   PetscFunctionBegin;
1125:   PetscTryTypeMethod(ts, adjointreset);
1126:   if (ts->quadraturets) { /* if there is integral in the cost function */
1127:     PetscCall(VecDestroy(&ts->vec_drdu_col));
1128:     if (ts->vecs_sensip) PetscCall(VecDestroy(&ts->vec_drdp_col));
1129:   }
1130:   ts->vecs_sensi         = NULL;
1131:   ts->vecs_sensip        = NULL;
1132:   ts->vecs_sensi2        = NULL;
1133:   ts->vecs_sensi2p       = NULL;
1134:   ts->vec_dir            = NULL;
1135:   ts->adjointsetupcalled = PETSC_FALSE;
1136:   PetscFunctionReturn(PETSC_SUCCESS);
1137: }

1139: /*@
1140:   TSAdjointSetSteps - Sets the number of steps the adjoint solver should take backward in time

1142:   Logically Collective

1144:   Input Parameters:
1145: + ts    - the `TS` context obtained from `TSCreate()`
1146: - steps - number of steps to use

1148:   Level: intermediate

1150:   Notes:
1151:   Normally one does not call this and `TSAdjointSolve()` integrates back to the original timestep. One can call this
1152:   so as to integrate back to less than the original timestep

1154: .seealso: [](ch_ts), `TSAdjointSolve()`, `TS`, `TSSetExactFinalTime()`
1155: @*/
1156: PetscErrorCode TSAdjointSetSteps(TS ts, PetscInt steps)
1157: {
1158:   PetscFunctionBegin;
1161:   PetscCheck(steps >= 0, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Cannot step back a negative number of steps");
1162:   PetscCheck(steps <= ts->steps, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Cannot step back more than the total number of forward steps");
1163:   ts->adjoint_max_steps = steps;
1164:   PetscFunctionReturn(PETSC_SUCCESS);
1165: }

1167: // PetscClangLinter pragma disable: -fdoc-*
1168: /*@
1169:   TSAdjointSetRHSJacobian - Deprecated, use `TSSetRHSJacobianP()`

1171:   Level: deprecated
1172: @*/
1173: PetscErrorCode TSAdjointSetRHSJacobian(TS ts, Mat Amat, PetscErrorCode (*func)(TS, PetscReal, Vec, Mat, void *), PetscCtx ctx)
1174: {
1175:   PetscFunctionBegin;

1179:   ts->rhsjacobianp    = func;
1180:   ts->rhsjacobianpctx = ctx;
1181:   if (Amat) {
1182:     PetscCall(PetscObjectReference((PetscObject)Amat));
1183:     PetscCall(MatDestroy(&ts->Jacp));
1184:     ts->Jacp = Amat;
1185:   }
1186:   PetscFunctionReturn(PETSC_SUCCESS);
1187: }

1189: // PetscClangLinter pragma disable: -fdoc-*
1190: /*@
1191:   TSAdjointComputeRHSJacobian - Deprecated, use `TSComputeRHSJacobianP()`

1193:   Level: deprecated
1194: @*/
1195: PetscErrorCode TSAdjointComputeRHSJacobian(TS ts, PetscReal t, Vec U, Mat Amat)
1196: {
1197:   PetscFunctionBegin;

1202:   PetscCallBack("TS callback JacobianP for sensitivity analysis", (*ts->rhsjacobianp)(ts, t, U, Amat, ts->rhsjacobianpctx));
1203:   PetscFunctionReturn(PETSC_SUCCESS);
1204: }

1206: // PetscClangLinter pragma disable: -fdoc-*
1207: /*@
1208:   TSAdjointComputeDRDYFunction - Deprecated, use `TSGetQuadratureTS()` then `TSComputeRHSJacobian()`

1210:   Level: deprecated
1211: @*/
1212: PetscErrorCode TSAdjointComputeDRDYFunction(TS ts, PetscReal t, Vec U, Vec *DRDU)
1213: {
1214:   PetscFunctionBegin;

1218:   PetscCallBack("TS callback DRDY for sensitivity analysis", (*ts->drdufunction)(ts, t, U, DRDU, ts->costintegrandctx));
1219:   PetscFunctionReturn(PETSC_SUCCESS);
1220: }

1222: // PetscClangLinter pragma disable: -fdoc-*
1223: /*@
1224:   TSAdjointComputeDRDPFunction - Deprecated, use `TSGetQuadratureTS()` then `TSComputeRHSJacobianP()`

1226:   Level: deprecated
1227: @*/
1228: PetscErrorCode TSAdjointComputeDRDPFunction(TS ts, PetscReal t, Vec U, Vec *DRDP)
1229: {
1230:   PetscFunctionBegin;

1234:   PetscCallBack("TS callback DRDP for sensitivity analysis", (*ts->drdpfunction)(ts, t, U, DRDP, ts->costintegrandctx));
1235:   PetscFunctionReturn(PETSC_SUCCESS);
1236: }

1238: // PetscClangLinter pragma disable: -fdoc-param-list-func-parameter-documentation
1239: /*@
1240:   TSAdjointMonitorSensi - monitors the first lambda sensitivity

1242:   Level: intermediate

1244: .seealso: [](ch_ts), `TSAdjointMonitorSet()`
1245: @*/
1246: static PetscErrorCode TSAdjointMonitorSensi(TS ts, PetscInt step, PetscReal ptime, Vec v, PetscInt numcost, Vec *lambda, Vec *mu, PetscViewerAndFormat *vf)
1247: {
1248:   PetscViewer viewer = vf->viewer;

1250:   PetscFunctionBegin;
1252:   PetscCall(PetscViewerPushFormat(viewer, vf->format));
1253:   PetscCall(VecView(lambda[0], viewer));
1254:   PetscCall(PetscViewerPopFormat(viewer));
1255:   PetscFunctionReturn(PETSC_SUCCESS);
1256: }

1258: /*@
1259:   TSAdjointMonitorSetFromOptions - Sets a monitor function and viewer appropriate for the type indicated by the user

1261:   Collective

1263:   Input Parameters:
1264: + ts           - `TS` object you wish to monitor
1265: . name         - the monitor type one is seeking
1266: . help         - message indicating what monitoring is done
1267: . manual       - manual page for the monitor
1268: . monitor      - the monitor function, its context must be a `PetscViewerAndFormat`
1269: - monitorsetup - a function that is called once ONLY if the user selected this monitor that may set additional features of the `TS` or `PetscViewer` objects

1271:   Calling sequence of `monitor`:
1272: + ts      - the `TS` context
1273: . step    - iteration number (after the final time step the monitor routine is called with
1274:                 a step of -1, this is at the final time which may have been interpolated to)
1275: . time    - current time
1276: . u       - current iterate
1277: . numcost - number of cost functions
1278: . lambda  - sensitivities to initial conditions
1279: . mu      - sensitivities to parameters
1280: - vf      - the `PetscViewer` and format the monitor is using

1282:   Calling sequence of `monitorsetup`:
1283: + ts - the `TS` object being monitored
1284: - vf - the `PetscViewer` and format the monitor is using

1286:   Level: developer

1288: .seealso: [](ch_ts), `PetscOptionsCreateViewer()`, `PetscOptionsGetReal()`, `PetscOptionsHasName()`, `PetscOptionsGetString()`,
1289:           `PetscOptionsGetIntArray()`, `PetscOptionsGetRealArray()`, `PetscOptionsBool()`,
1290:           `PetscOptionsInt()`, `PetscOptionsString()`, `PetscOptionsReal()`,
1291:           `PetscOptionsName()`, `PetscOptionsBegin()`, `PetscOptionsEnd()`, `PetscOptionsHeadBegin()`,
1292:           `PetscOptionsStringArray()`, `PetscOptionsRealArray()`, `PetscOptionsScalar()`,
1293:           `PetscOptionsBoolGroupBegin()`, `PetscOptionsBoolGroup()`, `PetscOptionsBoolGroupEnd()`,
1294:           `PetscOptionsFList()`, `PetscOptionsEList()`, `PetscViewerAndFormat`
1295: @*/
1296: PetscErrorCode TSAdjointMonitorSetFromOptions(TS ts, const char name[], const char help[], const char manual[], PetscErrorCode (*monitor)(TS ts, PetscInt step, PetscReal time, Vec u, PetscInt numcost, Vec *lambda, Vec *mu, PetscViewerAndFormat *vf), PetscErrorCode (*monitorsetup)(TS ts, PetscViewerAndFormat *vf))
1297: {
1298:   PetscViewer       viewer;
1299:   PetscViewerFormat format;
1300:   PetscBool         flg;

1302:   PetscFunctionBegin;
1303:   PetscCall(PetscOptionsCreateViewer(PetscObjectComm((PetscObject)ts), ((PetscObject)ts)->options, ((PetscObject)ts)->prefix, name, &viewer, &format, &flg));
1304:   if (flg) {
1305:     PetscViewerAndFormat *vf;
1306:     PetscCall(PetscViewerAndFormatCreate(viewer, format, &vf));
1307:     PetscCall(PetscViewerDestroy(&viewer));
1308:     if (monitorsetup) PetscCall((*monitorsetup)(ts, vf));
1309:     PetscCall(TSAdjointMonitorSet(ts, (PetscErrorCode (*)(TS, PetscInt, PetscReal, Vec, PetscInt, Vec *, Vec *, PetscCtx))monitor, vf, (PetscCtxDestroyFn *)PetscViewerAndFormatDestroy));
1310:   }
1311:   PetscFunctionReturn(PETSC_SUCCESS);
1312: }

1314: /*@
1315:   TSAdjointMonitorSet - Sets an ADDITIONAL function that is to be used at every
1316:   timestep to display the iteration's  progress.

1318:   Logically Collective

1320:   Input Parameters:
1321: + ts              - the `TS` context obtained from `TSCreate()`
1322: . adjointmonitor  - monitoring routine
1323: . adjointmctx     - [optional] context for private data for the monitor routine (use `NULL` if no context is desired)
1324: - adjointmdestroy - [optional] routine that frees monitor context (may be `NULL`), see `PetscCtxDestroyFn` for its calling sequence

1326:   Calling sequence of `adjointmonitor`:
1327: + ts          - the `TS` context
1328: . steps       - iteration number (after the final time step the monitor routine is called with
1329:                 a step of -1, this is at the final time which may have been interpolated to)
1330: . time        - current time
1331: . u           - current iterate
1332: . numcost     - number of cost functions
1333: . lambda      - sensitivities to initial conditions
1334: . mu          - sensitivities to parameters
1335: - adjointmctx - [optional] adjoint monitoring context

1337:   Level: intermediate

1339:   Note:
1340:   This routine adds an additional monitor to the list of monitors that
1341:   already has been loaded.

1343:   Fortran Notes:
1344:   Only a single monitor function can be set for each `TS` object

1346: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSAdjointMonitorCancel()`, `PetscCtxDestroyFn`
1347: @*/
1348: PetscErrorCode TSAdjointMonitorSet(TS ts, PetscErrorCode (*adjointmonitor)(TS ts, PetscInt steps, PetscReal time, Vec u, PetscInt numcost, Vec *lambda, Vec *mu, PetscCtx adjointmctx), PetscCtx adjointmctx, PetscCtxDestroyFn *adjointmdestroy)
1349: {
1350:   PetscFunctionBegin;
1352:   for (PetscInt i = 0; i < ts->numbermonitors; i++) {
1353:     PetscBool identical;

1355:     PetscCall(PetscMonitorCompare((PetscErrorCode (*)(void))(PetscVoidFn *)adjointmonitor, adjointmctx, adjointmdestroy, (PetscErrorCode (*)(void))(PetscVoidFn *)ts->adjointmonitor[i], ts->adjointmonitorcontext[i], ts->adjointmonitordestroy[i], &identical));
1356:     if (identical) PetscFunctionReturn(PETSC_SUCCESS);
1357:   }
1358:   PetscCheck(ts->numberadjointmonitors < MAXTSMONITORS, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Too many adjoint monitors set");
1359:   ts->adjointmonitor[ts->numberadjointmonitors]          = adjointmonitor;
1360:   ts->adjointmonitordestroy[ts->numberadjointmonitors]   = adjointmdestroy;
1361:   ts->adjointmonitorcontext[ts->numberadjointmonitors++] = adjointmctx;
1362:   PetscFunctionReturn(PETSC_SUCCESS);
1363: }

1365: /*@
1366:   TSAdjointMonitorCancel - Clears all the adjoint monitors that have been set on a time-step object.

1368:   Logically Collective

1370:   Input Parameter:
1371: . ts - the `TS` context obtained from `TSCreate()`

1373:   Notes:
1374:   There is no way to remove a single, specific monitor.

1376:   Level: intermediate

1378: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSAdjointMonitorSet()`
1379: @*/
1380: PetscErrorCode TSAdjointMonitorCancel(TS ts)
1381: {
1382:   PetscFunctionBegin;
1384:   for (PetscInt i = 0; i < ts->numberadjointmonitors; i++) {
1385:     if (ts->adjointmonitordestroy[i]) PetscCall((*ts->adjointmonitordestroy[i])(&ts->adjointmonitorcontext[i]));
1386:   }
1387:   ts->numberadjointmonitors = 0;
1388:   PetscFunctionReturn(PETSC_SUCCESS);
1389: }

1391: /*@
1392:   TSAdjointMonitorDefault - the default monitor of adjoint computations

1394:   Input Parameters:
1395: + ts      - the `TS` context
1396: . step    - iteration number (after the final time step the monitor routine is called with a
1397: step of -1, this is at the final time which may have been interpolated to)
1398: . time    - current time
1399: . v       - current iterate
1400: . numcost - number of cost functions
1401: . lambda  - sensitivities to initial conditions
1402: . mu      - sensitivities to parameters
1403: - vf      - the viewer and format

1405:   Level: intermediate

1407: .seealso: [](ch_ts), `TS`, `TSAdjointSolve()`, `TSAdjointMonitorSet()`
1408: @*/
1409: PetscErrorCode TSAdjointMonitorDefault(TS ts, PetscInt step, PetscReal time, Vec v, PetscInt numcost, Vec lambda[], Vec mu[], PetscViewerAndFormat *vf)
1410: {
1411:   PetscViewer viewer = vf->viewer;

1413:   PetscFunctionBegin;
1414:   (void)v;
1415:   (void)numcost;
1416:   (void)lambda;
1417:   (void)mu;
1419:   PetscCall(PetscViewerPushFormat(viewer, vf->format));
1420:   PetscCall(PetscViewerASCIIAddTab(viewer, ((PetscObject)ts)->tablevel));
1421:   PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscInt_FMT " TS dt %g time %g%s", step, (double)ts->time_step, (double)time, ts->steprollback ? " (r)\n" : "\n"));
1422:   PetscCall(PetscViewerASCIISubtractTab(viewer, ((PetscObject)ts)->tablevel));
1423:   PetscCall(PetscViewerPopFormat(viewer));
1424:   PetscFunctionReturn(PETSC_SUCCESS);
1425: }

1427: /*@
1428:   TSAdjointMonitorDrawSensi - Monitors progress of the adjoint `TS` solvers by calling
1429:   `VecView()` for the sensitivities to initial states at each timestep

1431:   Collective

1433:   Input Parameters:
1434: + ts      - the `TS` context
1435: . step    - current time-step
1436: . ptime   - current time
1437: . u       - current state
1438: . numcost - number of cost functions
1439: . lambda  - sensitivities to initial conditions
1440: . mu      - sensitivities to parameters
1441: - dummy   - either a viewer or `NULL`

1443:   Level: intermediate

1445: .seealso: [](ch_ts), `TSAdjointSolve()`, `TSAdjointMonitorSet()`, `TSAdjointMonitorDefault()`, `VecView()`
1446: @*/
1447: PetscErrorCode TSAdjointMonitorDrawSensi(TS ts, PetscInt step, PetscReal ptime, Vec u, PetscInt numcost, Vec lambda[], Vec mu[], void *dummy)
1448: {
1449:   TSMonitorDrawCtx ictx = (TSMonitorDrawCtx)dummy;
1450:   PetscDraw        draw;
1451:   PetscReal        xl, yl, xr, yr, h;
1452:   char             time[32];

1454:   PetscFunctionBegin;
1455:   if (!(((ictx->howoften > 0) && (!(step % ictx->howoften))) || ((ictx->howoften == -1) && ts->reason))) PetscFunctionReturn(PETSC_SUCCESS);

1457:   PetscCall(VecView(lambda[0], ictx->viewer));
1458:   PetscCall(PetscViewerDrawGetDraw(ictx->viewer, 0, &draw));
1459:   PetscCall(PetscSNPrintf(time, 32, "Timestep %" PetscInt_FMT " Time %g", step, (double)ptime));
1460:   PetscCall(PetscDrawGetCoordinates(draw, &xl, &yl, &xr, &yr));
1461:   h = yl + .95 * (yr - yl);
1462:   PetscCall(PetscDrawStringCentered(draw, .5 * (xl + xr), h, PETSC_DRAW_BLACK, time));
1463:   PetscCall(PetscDrawFlush(draw));
1464:   PetscFunctionReturn(PETSC_SUCCESS);
1465: }

1467: /*@
1468:   TSAdjointSetFromOptions - Sets various `TS` adjoint parameters from options database.

1470:   Collective

1472:   Input Parameters:
1473: + ts                 - the `TS` context
1474: - PetscOptionsObject - the options context

1476:   Options Database Keys:
1477: + -ts_adjoint_solve (yes|no)     - After solving the ODE/DAE solve the adjoint problem (requires `-ts_save_trajectory`)
1478: . -ts_adjoint_monitor            - print information at each adjoint time step
1479: - -ts_adjoint_monitor_draw_sensi - monitor the sensitivity of the first cost function wrt initial conditions (lambda[0]) graphically

1481:   Level: developer

1483:   Note:
1484:   This is not normally called directly by users

1486: .seealso: [](ch_ts), `TSSetSaveTrajectory()`, `TSTrajectorySetUp()`
1487: @*/
1488: PetscErrorCode TSAdjointSetFromOptions(TS ts, PetscOptionItems PetscOptionsObject)
1489: {
1490:   PetscBool tflg, opt;

1492:   PetscFunctionBegin;
1494:   PetscOptionsHeadBegin(PetscOptionsObject, "TS Adjoint options");
1495:   tflg = ts->adjoint_solve ? PETSC_TRUE : PETSC_FALSE;
1496:   PetscCall(PetscOptionsBool("-ts_adjoint_solve", "Solve the adjoint problem immediately after solving the forward problem", "", tflg, &tflg, &opt));
1497:   if (opt) {
1498:     PetscCall(TSSetSaveTrajectory(ts));
1499:     ts->adjoint_solve = tflg;
1500:   }
1501:   PetscCall(TSAdjointMonitorSetFromOptions(ts, "-ts_adjoint_monitor", "Monitor adjoint timestep size", "TSAdjointMonitorDefault", TSAdjointMonitorDefault, NULL));
1502:   PetscCall(TSAdjointMonitorSetFromOptions(ts, "-ts_adjoint_monitor_sensi", "Monitor sensitivity in the adjoint computation", "TSAdjointMonitorSensi", TSAdjointMonitorSensi, NULL));
1503:   opt = PETSC_FALSE;
1504:   PetscCall(PetscOptionsName("-ts_adjoint_monitor_draw_sensi", "Monitor adjoint sensitivities (lambda only) graphically", "TSAdjointMonitorDrawSensi", &opt));
1505:   if (opt) {
1506:     TSMonitorDrawCtx ctx;
1507:     PetscInt         howoften = 1;

1509:     PetscCall(PetscOptionsInt("-ts_adjoint_monitor_draw_sensi", "Monitor adjoint sensitivities (lambda only) graphically", "TSAdjointMonitorDrawSensi", howoften, &howoften, NULL));
1510:     PetscCall(TSMonitorDrawCtxCreate(PetscObjectComm((PetscObject)ts), NULL, NULL, PETSC_DECIDE, PETSC_DECIDE, 300, 300, howoften, &ctx));
1511:     PetscCall(TSAdjointMonitorSet(ts, TSAdjointMonitorDrawSensi, ctx, (PetscCtxDestroyFn *)TSMonitorDrawCtxDestroy));
1512:   }
1513:   PetscFunctionReturn(PETSC_SUCCESS);
1514: }

1516: /*@
1517:   TSAdjointStep - Steps one time step backward in the adjoint run

1519:   Collective

1521:   Input Parameter:
1522: . ts - the `TS` context obtained from `TSCreate()`

1524:   Level: intermediate

1526: .seealso: [](ch_ts), `TSAdjointSetUp()`, `TSAdjointSolve()`
1527: @*/
1528: PetscErrorCode TSAdjointStep(TS ts)
1529: {
1530:   DM dm;

1532:   PetscFunctionBegin;
1534:   PetscCall(TSGetDM(ts, &dm));
1535:   PetscCall(TSAdjointSetUp(ts));
1536:   ts->steps--; /* must decrease the step index before the adjoint step is taken. */

1538:   ts->reason     = TS_CONVERGED_ITERATING;
1539:   ts->ptime_prev = ts->ptime;
1540:   PetscCall(PetscLogEventBegin(TS_AdjointStep, ts, 0, 0, 0));
1541:   PetscUseTypeMethod(ts, adjointstep);
1542:   PetscCall(PetscLogEventEnd(TS_AdjointStep, ts, 0, 0, 0));
1543:   ts->adjoint_steps++;

1545:   if (ts->reason < 0) {
1546:     PetscCheck(!ts->errorifstepfailed, PetscObjectComm((PetscObject)ts), PETSC_ERR_NOT_CONVERGED, "TSAdjointStep has failed due to %s", TSConvergedReasons[ts->reason]);
1547:   } else if (!ts->reason) {
1548:     if (ts->adjoint_steps >= ts->adjoint_max_steps) ts->reason = TS_CONVERGED_ITS;
1549:   }
1550:   PetscFunctionReturn(PETSC_SUCCESS);
1551: }

1553: /*@
1554:   TSAdjointSolve - Solves the discrete ajoint problem for an ODE/DAE

1556:   Collective
1557:   `

1559:   Input Parameter:
1560: . ts - the `TS` context obtained from `TSCreate()`

1562:   Options Database Key:
1563: . -ts_adjoint_view_solution viewerinfo - views the first gradient with respect to the initial values

1565:   Level: intermediate

1567:   Notes:
1568:   This must be called after a call to `TSSolve()` that solves the forward problem

1570:   By default this will integrate back to the initial time, one can use `TSAdjointSetSteps()` to step back to a later time

1572: .seealso: [](ch_ts), `TSCreate()`, `TSSetCostGradients()`, `TSSetSolution()`, `TSAdjointStep()`
1573: @*/
1574: PetscErrorCode TSAdjointSolve(TS ts)
1575: {
1576:   static PetscBool cite = PETSC_FALSE;
1577: #if defined(TSADJOINT_STAGE)
1578:   PetscLogStage adjoint_stage;
1579: #endif

1581:   PetscFunctionBegin;
1583:   PetscCall(PetscCitationsRegister("@article{Zhang2022tsadjoint,\n"
1584:                                    "  title         = {{PETSc TSAdjoint: A Discrete Adjoint ODE Solver for First-Order and Second-Order Sensitivity Analysis}},\n"
1585:                                    "  author        = {Zhang, Hong and Constantinescu, Emil M.  and Smith, Barry F.},\n"
1586:                                    "  journal       = {SIAM Journal on Scientific Computing},\n"
1587:                                    "  volume        = {44},\n"
1588:                                    "  number        = {1},\n"
1589:                                    "  pages         = {C1-C24},\n"
1590:                                    "  doi           = {10.1137/21M140078X},\n"
1591:                                    "  year          = {2022}\n}\n",
1592:                                    &cite));
1593: #if defined(TSADJOINT_STAGE)
1594:   PetscCall(PetscLogStageRegister("TSAdjoint", &adjoint_stage));
1595:   PetscCall(PetscLogStagePush(adjoint_stage));
1596: #endif
1597:   PetscCall(TSAdjointSetUp(ts));

1599:   /* reset time step and iteration counters */
1600:   ts->adjoint_steps     = 0;
1601:   ts->ksp_its           = 0;
1602:   ts->snes_its          = 0;
1603:   ts->num_snes_failures = 0;
1604:   ts->reject            = 0;
1605:   ts->reason            = TS_CONVERGED_ITERATING;

1607:   if (!ts->adjoint_max_steps) ts->adjoint_max_steps = ts->steps;
1608:   if (ts->adjoint_steps >= ts->adjoint_max_steps) ts->reason = TS_CONVERGED_ITS;

1610:   while (!ts->reason) {
1611:     PetscCall(TSTrajectoryGet(ts->trajectory, ts, ts->steps, &ts->ptime));
1612:     PetscCall(TSAdjointMonitor(ts, ts->steps, ts->ptime, ts->vec_sol, ts->numcost, ts->vecs_sensi, ts->vecs_sensip));
1613:     PetscCall(TSAdjointEventHandler(ts));
1614:     PetscCall(TSAdjointStep(ts));
1615:     if ((ts->vec_costintegral || ts->quadraturets) && !ts->costintegralfwd) PetscCall(TSAdjointCostIntegral(ts));
1616:   }
1617:   if (!ts->steps) {
1618:     PetscCall(TSTrajectoryGet(ts->trajectory, ts, ts->steps, &ts->ptime));
1619:     PetscCall(TSAdjointMonitor(ts, ts->steps, ts->ptime, ts->vec_sol, ts->numcost, ts->vecs_sensi, ts->vecs_sensip));
1620:   }
1621:   ts->solvetime = ts->ptime;
1622:   PetscCall(TSTrajectoryViewFromOptions(ts->trajectory, NULL, "-ts_trajectory_view"));
1623:   PetscCall(VecViewFromOptions(ts->vecs_sensi[0], (PetscObject)ts, "-ts_adjoint_view_solution"));
1624:   ts->adjoint_max_steps = 0;
1625: #if defined(TSADJOINT_STAGE)
1626:   PetscCall(PetscLogStagePop());
1627: #endif
1628:   PetscFunctionReturn(PETSC_SUCCESS);
1629: }

1631: /*@
1632:   TSAdjointMonitor - Runs all user-provided adjoint monitor routines set using `TSAdjointMonitorSet()`

1634:   Collective

1636:   Input Parameters:
1637: + ts      - time stepping context obtained from `TSCreate()`
1638: . step    - step number that has just completed
1639: . ptime   - model time of the state
1640: . u       - state at the current model time
1641: . numcost - number of cost functions (dimension of lambda  or mu)
1642: . lambda  - vectors containing the gradients of the cost functions with respect to the ODE/DAE solution variables
1643: - mu      - vectors containing the gradients of the cost functions with respect to the problem parameters

1645:   Level: developer

1647:   Note:
1648:   `TSAdjointMonitor()` is typically used automatically within the time stepping implementations.
1649:   Users would almost never call this routine directly.

1651: .seealso: `TSAdjointMonitorSet()`, `TSAdjointSolve()`
1652: @*/
1653: PetscErrorCode TSAdjointMonitor(TS ts, PetscInt step, PetscReal ptime, Vec u, PetscInt numcost, Vec lambda[], Vec mu[])
1654: {
1655:   PetscInt i, n = ts->numberadjointmonitors;

1657:   PetscFunctionBegin;
1660:   PetscCall(VecLockReadPush(u));
1661:   for (i = 0; i < n; i++) PetscCall((*ts->adjointmonitor[i])(ts, step, ptime, u, numcost, lambda, mu, ts->adjointmonitorcontext[i]));
1662:   PetscCall(VecLockReadPop(u));
1663:   PetscFunctionReturn(PETSC_SUCCESS);
1664: }

1666: /*@
1667:   TSAdjointCostIntegral - Evaluate the cost integral in the adjoint run.

1669:   Collective

1671:   Input Parameter:
1672: . ts - time stepping context

1674:   Level: advanced

1676:   Notes:
1677:   This function cannot be called until `TSAdjointStep()` has been completed.

1679: .seealso: [](ch_ts), `TSAdjointSolve()`, `TSAdjointStep()`
1680:  @*/
1681: PetscErrorCode TSAdjointCostIntegral(TS ts)
1682: {
1683:   PetscFunctionBegin;
1685:   PetscUseTypeMethod(ts, adjointintegral);
1686:   PetscFunctionReturn(PETSC_SUCCESS);
1687: }

1689: /* ------------------ Forward (tangent linear) sensitivity  ------------------*/

1691: /*@
1692:   TSForwardSetUp - Sets up the internal data structures for the later use
1693:   of forward sensitivity analysis

1695:   Collective

1697:   Input Parameter:
1698: . ts - the `TS` context obtained from `TSCreate()`

1700:   Level: advanced

1702: .seealso: [](ch_ts), `TS`, `TSCreate()`, `TSDestroy()`, `TSSetUp()`
1703: @*/
1704: PetscErrorCode TSForwardSetUp(TS ts)
1705: {
1706:   PetscFunctionBegin;
1708:   if (ts->forwardsetupcalled) PetscFunctionReturn(PETSC_SUCCESS);
1709:   PetscTryTypeMethod(ts, forwardsetup);
1710:   PetscCall(VecDuplicate(ts->vec_sol, &ts->vec_sensip_col));
1711:   ts->forwardsetupcalled = PETSC_TRUE;
1712:   PetscFunctionReturn(PETSC_SUCCESS);
1713: }

1715: /*@
1716:   TSForwardReset - Reset the internal data structures used by forward sensitivity analysis

1718:   Collective

1720:   Input Parameter:
1721: . ts - the `TS` context obtained from `TSCreate()`

1723:   Level: advanced

1725: .seealso: [](ch_ts), `TSCreate()`, `TSDestroy()`, `TSForwardSetUp()`
1726: @*/
1727: PetscErrorCode TSForwardReset(TS ts)
1728: {
1729:   TS quadts = ts->quadraturets;

1731:   PetscFunctionBegin;
1733:   PetscTryTypeMethod(ts, forwardreset);
1734:   PetscCall(MatDestroy(&ts->mat_sensip));
1735:   if (quadts) PetscCall(MatDestroy(&quadts->mat_sensip));
1736:   PetscCall(VecDestroy(&ts->vec_sensip_col));
1737:   ts->forward_solve      = PETSC_FALSE;
1738:   ts->forwardsetupcalled = PETSC_FALSE;
1739:   PetscFunctionReturn(PETSC_SUCCESS);
1740: }

1742: /*@
1743:   TSForwardSetIntegralGradients - Set the vectors holding forward sensitivities of the integral term.

1745:   Input Parameters:
1746: + ts        - the `TS` context obtained from `TSCreate()`
1747: . numfwdint - number of integrals
1748: - vp        - the vectors containing the gradients for each integral w.r.t. parameters

1750:   Level: deprecated

1752: .seealso: [](ch_ts), `TSForwardGetSensitivities()`, `TSForwardGetIntegralGradients()`, `TSForwardStep()`
1753: @*/
1754: PetscErrorCode TSForwardSetIntegralGradients(TS ts, PetscInt numfwdint, Vec vp[])
1755: {
1756:   PetscFunctionBegin;
1758:   PetscCheck(!ts->numcost || ts->numcost == numfwdint, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "The number of cost functions (2nd parameter of TSSetCostIntegrand()) is inconsistent with the one set by TSSetCostIntegrand()");
1759:   if (!ts->numcost) ts->numcost = numfwdint;

1761:   ts->vecs_integral_sensip = vp;
1762:   PetscFunctionReturn(PETSC_SUCCESS);
1763: }

1765: /*@
1766:   TSForwardGetIntegralGradients - Returns the forward sensitivities of the integral term.

1768:   Input Parameter:
1769: . ts - the `TS` context obtained from `TSCreate()`

1771:   Output Parameters:
1772: + numfwdint - number of integrals
1773: - vp        - the vectors containing the gradients for each integral w.r.t. parameters

1775:   Level: deprecated

1777: .seealso: [](ch_ts), `TSForwardSetSensitivities()`, `TSForwardSetIntegralGradients()`, `TSForwardStep()`
1778: @*/
1779: PetscErrorCode TSForwardGetIntegralGradients(TS ts, PetscInt *numfwdint, Vec *vp[])
1780: {
1781:   PetscFunctionBegin;
1783:   PetscAssertPointer(vp, 3);
1784:   if (numfwdint) *numfwdint = ts->numcost;
1785:   if (vp) *vp = ts->vecs_integral_sensip;
1786:   PetscFunctionReturn(PETSC_SUCCESS);
1787: }

1789: /*@
1790:   TSForwardStep - Compute the forward sensitivity for one time step.

1792:   Collective

1794:   Input Parameter:
1795: . ts - time stepping context

1797:   Level: advanced

1799:   Notes:
1800:   This function cannot be called until `TSStep()` has been completed.

1802: .seealso: [](ch_ts), `TSForwardSetSensitivities()`, `TSForwardGetSensitivities()`, `TSForwardSetIntegralGradients()`, `TSForwardGetIntegralGradients()`, `TSForwardSetUp()`
1803: @*/
1804: PetscErrorCode TSForwardStep(TS ts)
1805: {
1806:   PetscFunctionBegin;
1808:   PetscCall(PetscLogEventBegin(TS_ForwardStep, ts, 0, 0, 0));
1809:   PetscUseTypeMethod(ts, forwardstep);
1810:   PetscCall(PetscLogEventEnd(TS_ForwardStep, ts, 0, 0, 0));
1811:   PetscCheck(ts->reason >= 0 || !ts->errorifstepfailed, PetscObjectComm((PetscObject)ts), PETSC_ERR_NOT_CONVERGED, "TSFowardStep has failed due to %s", TSConvergedReasons[ts->reason]);
1812:   PetscFunctionReturn(PETSC_SUCCESS);
1813: }

1815: /*@
1816:   TSForwardSetSensitivities - Sets the initial value of the trajectory sensitivities of solution  w.r.t. the problem parameters and initial values.

1818:   Logically Collective

1820:   Input Parameters:
1821: + ts   - the `TS` context obtained from `TSCreate()`
1822: . nump - number of parameters
1823: - Smat - sensitivities with respect to the parameters, the number of entries in these vectors is the same as the number of parameters

1825:   Level: beginner

1827:   Notes:
1828:   Use `PETSC_DETERMINE` to use the number of columns of `Smat` for `nump`

1830:   Forward sensitivity is also called 'trajectory sensitivity' in some fields such as power systems.
1831:   This function turns on a flag to trigger `TSSolve()` to compute forward sensitivities automatically.
1832:   You must call this function before `TSSolve()`.
1833:   The entries in the sensitivity matrix must be correctly initialized with the values S = dy/dp|startingtime.

1835: .seealso: [](ch_ts), `TSForwardGetSensitivities()`, `TSForwardSetIntegralGradients()`, `TSForwardGetIntegralGradients()`, `TSForwardStep()`
1836: @*/
1837: PetscErrorCode TSForwardSetSensitivities(TS ts, PetscInt nump, Mat Smat)
1838: {
1839:   PetscFunctionBegin;
1842:   ts->forward_solve = PETSC_TRUE;
1843:   if (nump == PETSC_DEFAULT || nump == PETSC_DETERMINE) PetscCall(MatGetSize(Smat, NULL, &ts->num_parameters));
1844:   else ts->num_parameters = nump;
1845:   PetscCall(PetscObjectReference((PetscObject)Smat));
1846:   PetscCall(MatDestroy(&ts->mat_sensip));
1847:   ts->mat_sensip = Smat;
1848:   PetscFunctionReturn(PETSC_SUCCESS);
1849: }

1851: /*@
1852:   TSForwardGetSensitivities - Returns the trajectory sensitivities

1854:   Not Collective, but Smat returned is parallel if ts is parallel

1856:   Output Parameters:
1857: + ts   - the `TS` context obtained from `TSCreate()`
1858: . nump - number of parameters
1859: - Smat - sensitivities with respect to the parameters, the number of entries in these vectors is the same as the number of parameters

1861:   Level: intermediate

1863: .seealso: [](ch_ts), `TSForwardSetSensitivities()`, `TSForwardSetIntegralGradients()`, `TSForwardGetIntegralGradients()`, `TSForwardStep()`
1864: @*/
1865: PetscErrorCode TSForwardGetSensitivities(TS ts, PetscInt *nump, Mat *Smat)
1866: {
1867:   PetscFunctionBegin;
1869:   if (nump) *nump = ts->num_parameters;
1870:   if (Smat) *Smat = ts->mat_sensip;
1871:   PetscFunctionReturn(PETSC_SUCCESS);
1872: }

1874: /*@
1875:   TSForwardCostIntegral - Evaluate the cost integral in the forward run.

1877:   Collective

1879:   Input Parameter:
1880: . ts - time stepping context

1882:   Level: advanced

1884:   Note:
1885:   This function cannot be called until `TSStep()` has been completed.

1887: .seealso: [](ch_ts), `TS`, `TSSolve()`, `TSAdjointCostIntegral()`
1888: @*/
1889: PetscErrorCode TSForwardCostIntegral(TS ts)
1890: {
1891:   PetscFunctionBegin;
1893:   PetscUseTypeMethod(ts, forwardintegral);
1894:   PetscFunctionReturn(PETSC_SUCCESS);
1895: }

1897: /*@
1898:   TSForwardSetInitialSensitivities - Set initial values for tangent linear sensitivities

1900:   Collective

1902:   Input Parameters:
1903: + ts   - the `TS` context obtained from `TSCreate()`
1904: - didp - parametric sensitivities of the initial condition

1906:   Level: intermediate

1908:   Notes:
1909:   `TSSolve()` allows users to pass the initial solution directly to `TS`. But the tangent linear variables cannot be initialized in this way.
1910:   This function is used to set initial values for tangent linear variables.

1912: .seealso: [](ch_ts), `TS`, `TSForwardSetSensitivities()`
1913: @*/
1914: PetscErrorCode TSForwardSetInitialSensitivities(TS ts, Mat didp)
1915: {
1916:   PetscFunctionBegin;
1919:   if (!ts->mat_sensip) PetscCall(TSForwardSetSensitivities(ts, PETSC_DETERMINE, didp));
1920:   PetscFunctionReturn(PETSC_SUCCESS);
1921: }

1923: /*@
1924:   TSForwardGetStages - Get the number of stages and the tangent linear sensitivities at the intermediate stages

1926:   Input Parameter:
1927: . ts - the `TS` context obtained from `TSCreate()`

1929:   Output Parameters:
1930: + ns - number of stages
1931: - S  - tangent linear sensitivities at the intermediate stages

1933:   Level: advanced

1935: .seealso: `TS`
1936: @*/
1937: PetscErrorCode TSForwardGetStages(TS ts, PetscInt *ns, Mat **S)
1938: {
1939:   PetscFunctionBegin;

1942:   if (!ts->ops->getstages) *S = NULL;
1943:   else PetscUseTypeMethod(ts, forwardgetstages, ns, S);
1944:   PetscFunctionReturn(PETSC_SUCCESS);
1945: }

1947: /*@
1948:   TSCreateQuadratureTS - Create a sub-`TS` that evaluates integrals over time

1950:   Input Parameters:
1951: + ts  - the `TS` context obtained from `TSCreate()`
1952: - fwd - flag indicating whether to evaluate cost integral in the forward run or the adjoint run

1954:   Output Parameter:
1955: . quadts - the child `TS` context

1957:   Level: intermediate

1959: .seealso: [](ch_ts), `TSGetQuadratureTS()`
1960: @*/
1961: PetscErrorCode TSCreateQuadratureTS(TS ts, PetscBool fwd, TS *quadts)
1962: {
1963:   char prefix[128];

1965:   PetscFunctionBegin;
1967:   PetscAssertPointer(quadts, 3);
1968:   PetscCall(TSDestroy(&ts->quadraturets));
1969:   PetscCall(TSCreate(PetscObjectComm((PetscObject)ts), &ts->quadraturets));
1970:   PetscCall(PetscObjectIncrementTabLevel((PetscObject)ts->quadraturets, (PetscObject)ts, 1));
1971:   PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "%squad_", ((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : ""));
1972:   PetscCall(TSSetOptionsPrefix(ts->quadraturets, prefix));
1973:   *quadts = ts->quadraturets;

1975:   if (ts->numcost) {
1976:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, ts->numcost, &(*quadts)->vec_sol));
1977:   } else {
1978:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, 1, &(*quadts)->vec_sol));
1979:   }
1980:   ts->costintegralfwd = fwd;
1981:   PetscFunctionReturn(PETSC_SUCCESS);
1982: }

1984: /*@
1985:   TSGetQuadratureTS - Return the sub-`TS` that evaluates integrals over time

1987:   Input Parameter:
1988: . ts - the `TS` context obtained from `TSCreate()`

1990:   Output Parameters:
1991: + fwd    - flag indicating whether to evaluate cost integral in the forward run or the adjoint run
1992: - quadts - the child `TS` context

1994:   Level: intermediate

1996: .seealso: [](ch_ts), `TSCreateQuadratureTS()`
1997: @*/
1998: PetscErrorCode TSGetQuadratureTS(TS ts, PetscBool *fwd, TS *quadts)
1999: {
2000:   PetscFunctionBegin;
2002:   if (fwd) *fwd = ts->costintegralfwd;
2003:   if (quadts) *quadts = ts->quadraturets;
2004:   PetscFunctionReturn(PETSC_SUCCESS);
2005: }

2007: /*@
2008:   TSComputeSNESJacobian - Compute the Jacobian needed for the `SNESSolve()` in `TS`

2010:   Collective

2012:   Input Parameters:
2013: + ts - the `TS` context obtained from `TSCreate()`
2014: - x  - state vector

2016:   Output Parameters:
2017: + J    - Jacobian matrix
2018: - Jpre - matrix used to compute the preconditioner for `J` (may be same as `J`)

2020:   Level: developer

2022:   Note:
2023:   Uses finite differencing when `TS` Jacobian is not available.

2025: .seealso: `SNES`, `TS`, `SNESSetJacobian()`, `TSSetRHSJacobian()`, `TSSetIJacobian()`
2026: @*/
2027: PetscErrorCode TSComputeSNESJacobian(TS ts, Vec x, Mat J, Mat Jpre)
2028: {
2029:   SNES snes                                          = ts->snes;
2030:   PetscErrorCode (*jac)(SNES, Vec, Mat, Mat, void *) = NULL;

2032:   PetscFunctionBegin;
2033:   /*
2034:     Unlike implicit methods, explicit methods do not have SNESMatFDColoring in the snes object
2035:     because SNESSolve() has not been called yet; so querying SNESMatFDColoring does not work for
2036:     explicit methods. Instead, we check the Jacobian compute function directly to determine if FD
2037:     coloring is used.
2038:   */
2039:   PetscCall(SNESGetJacobian(snes, NULL, NULL, &jac, NULL));
2040:   if (jac == SNESComputeJacobianDefaultColor) {
2041:     Vec f;
2042:     PetscCall(SNESSetSolution(snes, x));
2043:     PetscCall(SNESGetFunction(snes, &f, NULL, NULL));
2044:     /* Force MatFDColoringApply to evaluate the SNES residual function for the base vector */
2045:     PetscCall(SNESComputeFunction(snes, x, f));
2046:   }
2047:   PetscCall(SNESComputeJacobian(snes, x, J, Jpre));
2048:   PetscFunctionReturn(PETSC_SUCCESS);
2049: }