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*/