Actual source code: ex36.c
1: static char help[] = "Tests TSARKIMEX adjoint parameter sensitivities when the parameter appears in the IFunction, the RHSFunction, or both.\n\n";
3: /*
4: The scalar problem is
6: u_t + alpha u = -beta u, u(0) = 1
8: split so that alpha u is the implicit part F(t,u,u_t) = u_t + alpha u and -beta u is the
9: explicit part G(t,u). The single parameter p enters alpha, beta, or both depending on
10: -param_dependence, so that exactly one of TSSetIJacobianP() and TSSetRHSJacobianP() is
11: registered in the one-sided cases. With r the resulting decay rate, u(T) = exp(-r T) and
12: the cost J = u(T) has the analytic gradient dJ/dp = -T (dr/dp) exp(-r T). The computed gradient
13: is the gradient of the discrete cost, so it matches the analytic one only up to the time
14: discretization error.
15: */
17: #include <petscts.h>
19: typedef enum {
20: PARAM_IMPLICIT,
21: PARAM_EXPLICIT,
22: PARAM_BOTH
23: } ParamDependence;
24: static const char *const ParamDependences[] = {"implicit", "explicit", "both", "ParamDependence", "PARAM_", NULL};
26: typedef struct {
27: PetscReal p;
28: ParamDependence dep;
29: } AppCtx;
31: static PetscReal Alpha(AppCtx *user)
32: {
33: return user->dep == PARAM_EXPLICIT ? 1.0 : user->p;
34: }
36: static PetscReal Beta(AppCtx *user)
37: {
38: return user->dep == PARAM_IMPLICIT ? 1.0 : user->p;
39: }
41: /* G(t,u) = -beta u */
42: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
43: {
44: AppCtx *user = (AppCtx *)ctx;
46: PetscFunctionBeginUser;
47: PetscCall(VecCopy(U, G));
48: PetscCall(VecScale(G, -Beta(user)));
49: PetscFunctionReturn(PETSC_SUCCESS);
50: }
52: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
53: {
54: AppCtx *user = (AppCtx *)ctx;
56: PetscFunctionBeginUser;
57: PetscCall(MatSetValue(P, 0, 0, -Beta(user), INSERT_VALUES));
58: PetscCall(MatAssemblyBegin(P, MAT_FINAL_ASSEMBLY));
59: PetscCall(MatAssemblyEnd(P, MAT_FINAL_ASSEMBLY));
60: PetscFunctionReturn(PETSC_SUCCESS);
61: }
63: /* F(t,u,u_t) = u_t + alpha u */
64: static PetscErrorCode IFunction(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, void *ctx)
65: {
66: AppCtx *user = (AppCtx *)ctx;
68: PetscFunctionBeginUser;
69: PetscCall(VecCopy(U, F));
70: PetscCall(VecScale(F, Alpha(user)));
71: PetscCall(VecAXPY(F, 1.0, Udot));
72: PetscFunctionReturn(PETSC_SUCCESS);
73: }
75: static PetscErrorCode IJacobian(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal shift, Mat J, Mat P, void *ctx)
76: {
77: AppCtx *user = (AppCtx *)ctx;
79: PetscFunctionBeginUser;
80: PetscCall(MatSetValue(P, 0, 0, shift + Alpha(user), INSERT_VALUES));
81: PetscCall(MatAssemblyBegin(P, MAT_FINAL_ASSEMBLY));
82: PetscCall(MatAssemblyEnd(P, MAT_FINAL_ASSEMBLY));
83: PetscFunctionReturn(PETSC_SUCCESS);
84: }
86: /* dF/dp = u */
87: static PetscErrorCode IJacobianP(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal shift, Mat Jacp, void *ctx)
88: {
89: const PetscScalar *u;
91: PetscFunctionBeginUser;
92: PetscCall(VecGetArrayRead(U, &u));
93: PetscCall(MatSetValue(Jacp, 0, 0, u[0], INSERT_VALUES));
94: PetscCall(VecRestoreArrayRead(U, &u));
95: PetscCall(MatAssemblyBegin(Jacp, MAT_FINAL_ASSEMBLY));
96: PetscCall(MatAssemblyEnd(Jacp, MAT_FINAL_ASSEMBLY));
97: PetscFunctionReturn(PETSC_SUCCESS);
98: }
100: /* dG/dp = -u */
101: static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec U, Mat Jacp, void *ctx)
102: {
103: const PetscScalar *u;
105: PetscFunctionBeginUser;
106: PetscCall(VecGetArrayRead(U, &u));
107: PetscCall(MatSetValue(Jacp, 0, 0, -u[0], INSERT_VALUES));
108: PetscCall(VecRestoreArrayRead(U, &u));
109: PetscCall(MatAssemblyBegin(Jacp, MAT_FINAL_ASSEMBLY));
110: PetscCall(MatAssemblyEnd(Jacp, MAT_FINAL_ASSEMBLY));
111: PetscFunctionReturn(PETSC_SUCCESS);
112: }
114: int main(int argc, char **argv)
115: {
116: TS ts;
117: Vec U, lambda, mu;
118: Mat Jrhs, Ji, Jacp = NULL, Jacprhs = NULL;
119: AppCtx user;
120: PetscReal ftime = 0.1, rtol = 1e-2, rate, drate, analytic, gradient;
121: PetscBool nullijacobianp = PETSC_FALSE;
122: const PetscScalar *m;
124: PetscFunctionBeginUser;
125: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
126: user.p = 2.0;
127: user.dep = PARAM_IMPLICIT;
128: PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Adjoint parameter Jacobian options", "TS");
129: PetscCall(PetscOptionsEnum("-param_dependence", "Which part of the IMEX splitting the parameter appears in", NULL, ParamDependences, (PetscEnum)user.dep, (PetscEnum *)&user.dep, NULL));
130: PetscCall(PetscOptionsReal("-p", "Value of the parameter", NULL, user.p, &user.p, NULL));
131: PetscCall(PetscOptionsBool("-null_ijacobianp", "Register the IJacobianP matrix without a callback, as PETSC_NULL_FUNCTION does from Fortran", NULL, nullijacobianp, &nullijacobianp, NULL));
132: PetscOptionsEnd();
134: PetscCall(VecCreateSeq(PETSC_COMM_SELF, 1, &U));
135: PetscCall(VecSet(U, 1.0));
136: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, 1, 1, NULL, &Jrhs));
137: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, 1, 1, NULL, &Ji));
139: PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
140: PetscCall(TSSetType(ts, TSARKIMEX));
141: PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, &user));
142: PetscCall(TSSetRHSJacobian(ts, Jrhs, Jrhs, RHSJacobian, &user));
143: PetscCall(TSSetIFunction(ts, NULL, IFunction, &user));
144: PetscCall(TSSetIJacobian(ts, Ji, Ji, IJacobian, &user));
145: if (user.dep != PARAM_EXPLICIT) {
146: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, 1, 1, NULL, &Jacp));
147: if (nullijacobianp) {
148: /* seed the matrix so that a contribution taken from it, rather than from the absent callback, is unmistakable */
149: PetscCall(MatSetValue(Jacp, 0, 0, 1e3, INSERT_VALUES));
150: PetscCall(MatAssemblyBegin(Jacp, MAT_FINAL_ASSEMBLY));
151: PetscCall(MatAssemblyEnd(Jacp, MAT_FINAL_ASSEMBLY));
152: PetscCall(TSSetIJacobianP(ts, Jacp, NULL, NULL));
153: } else PetscCall(TSSetIJacobianP(ts, Jacp, IJacobianP, &user));
154: }
155: if (user.dep != PARAM_IMPLICIT) {
156: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, 1, 1, NULL, &Jacprhs));
157: PetscCall(TSSetRHSJacobianP(ts, Jacprhs, RHSJacobianP, &user));
158: }
160: PetscCall(TSSetSaveTrajectory(ts));
161: PetscCall(TSSetTime(ts, 0.0));
162: PetscCall(TSSetTimeStep(ts, 0.01));
163: PetscCall(TSSetMaxTime(ts, ftime));
164: PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
165: PetscCall(TSSetFromOptions(ts));
166: PetscCall(TSSolve(ts, U));
168: /* J = u(T), so lambda(T) = dJ/du(T) = 1 and mu(T) = 0 */
169: PetscCall(VecCreateSeq(PETSC_COMM_SELF, 1, &lambda));
170: PetscCall(VecSet(lambda, 1.0));
171: PetscCall(VecCreateSeq(PETSC_COMM_SELF, 1, &mu));
172: PetscCall(VecSet(mu, 0.0));
173: PetscCall(TSSetCostGradients(ts, 1, &lambda, &mu));
174: PetscCall(TSAdjointSolve(ts));
176: rate = Alpha(&user) + Beta(&user);
177: drate = user.dep == PARAM_BOTH ? 2.0 : 1.0;
178: /* an unregistered callback contributes nothing, so its term drops out of the expected gradient as well */
179: if (nullijacobianp && user.dep != PARAM_EXPLICIT) drate -= 1.0;
180: analytic = -ftime * drate * PetscExpReal(-rate * ftime);
181: PetscCall(VecGetArrayRead(mu, &m));
182: gradient = PetscRealPart(m[0]);
183: PetscCall(VecRestoreArrayRead(mu, &m));
184: /* petscdiff masks floating point numbers unless DIFF_NUMBERS=1, so report the verdict as text */
185: PetscCall(PetscPrintf(PETSC_COMM_SELF, "parameter dependence = %s\n", ParamDependences[user.dep]));
186: /* the agreement is the positive condition so that a NaN gradient, which compares false against everything, is reported as a failure */
187: if (PetscAbsReal(gradient - analytic) <= rtol * PetscAbsReal(analytic)) PetscCall(PetscPrintf(PETSC_COMM_SELF, " dJ/dp agrees with the analytic gradient to within the time discretization error\n"));
188: else PetscCall(PetscPrintf(PETSC_COMM_SELF, " dJ/dp = %g differs from the analytic gradient %g\n", (double)gradient, (double)analytic));
190: PetscCall(VecDestroy(&U));
191: PetscCall(VecDestroy(&lambda));
192: PetscCall(VecDestroy(&mu));
193: PetscCall(MatDestroy(&Jrhs));
194: PetscCall(MatDestroy(&Ji));
195: PetscCall(MatDestroy(&Jacp));
196: PetscCall(MatDestroy(&Jacprhs));
197: PetscCall(TSDestroy(&ts));
198: PetscCall(PetscFinalize());
199: return 0;
200: }
202: /*TEST
204: testset:
205: requires: !single !complex
206: args: -ts_trajectory_type memory
208: test:
209: suffix: implicit
210: args: -param_dependence implicit
212: test:
213: suffix: explicit
214: args: -param_dependence explicit
216: test:
217: suffix: both
218: args: -param_dependence both
220: # a matrix registered without a callback must contribute nothing, on the TSTHETA path
221: # where TSComputeIJacobianP() is called with imex = PETSC_FALSE
222: test:
223: suffix: null_ijacobianp
224: args: -param_dependence implicit -null_ijacobianp -ts_type beuler
226: TEST*/