xref: /petsc/src/ts/tutorials/advection-diffusion-reaction/ex5adj_mf.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
160f0b76eSHong Zhang static char help[] = "Demonstrates adjoint sensitivity analysis for Reaction-Diffusion Equations.\n";
260f0b76eSHong Zhang 
360f0b76eSHong Zhang /*
460f0b76eSHong Zhang   See ex5.c for details on the equation.
560f0b76eSHong Zhang   This code demonestrates the TSAdjoint interface to a system of time-dependent partial differential equations.
660f0b76eSHong Zhang   It computes the sensitivity of a component in the final solution, which locates in the center of the 2D domain, w.r.t. the initial conditions.
760f0b76eSHong Zhang   The user does not need to provide any additional functions. The required functions in the original simulation are reused in the adjoint run.
860f0b76eSHong Zhang 
960f0b76eSHong Zhang   Runtime options:
1060f0b76eSHong Zhang     -forwardonly  - run the forward simulation without adjoint
1160f0b76eSHong Zhang     -implicitform - provide IFunction and IJacobian to TS, if not set, RHSFunction and RHSJacobian will be used
1260f0b76eSHong Zhang     -aijpc        - set the preconditioner matrix to be aij (the Jacobian matrix can be of a different type such as ELL)
1360f0b76eSHong Zhang */
1460f0b76eSHong Zhang 
1560f0b76eSHong Zhang #include "reaction_diffusion.h"
1660f0b76eSHong Zhang #include <petscdm.h>
1760f0b76eSHong Zhang #include <petscdmda.h>
1860f0b76eSHong Zhang 
1960f0b76eSHong Zhang /* Matrix free support */
2060f0b76eSHong Zhang typedef struct {
2160f0b76eSHong Zhang   PetscReal time;
2260f0b76eSHong Zhang   Vec       U;
2360f0b76eSHong Zhang   Vec       Udot;
2460f0b76eSHong Zhang   PetscReal shift;
2560f0b76eSHong Zhang   AppCtx   *appctx;
2660f0b76eSHong Zhang   TS        ts;
2760f0b76eSHong Zhang } MCtx;
2860f0b76eSHong Zhang 
2960f0b76eSHong Zhang /*
3060f0b76eSHong Zhang    User-defined routines
3160f0b76eSHong Zhang */
3260f0b76eSHong Zhang PetscErrorCode InitialConditions(DM, Vec);
3360f0b76eSHong Zhang PetscErrorCode RHSJacobianShell(TS, PetscReal, Vec, Mat, Mat, void *);
3460f0b76eSHong Zhang PetscErrorCode IJacobianShell(TS, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *);
3560f0b76eSHong Zhang 
36*9371c9d4SSatish Balay PetscErrorCode InitializeLambda(DM da, Vec lambda, PetscReal x, PetscReal y) {
3760f0b76eSHong Zhang   PetscInt i, j, Mx, My, xs, ys, xm, ym;
3860f0b76eSHong Zhang   Field  **l;
3960f0b76eSHong Zhang   PetscFunctionBegin;
4060f0b76eSHong Zhang 
419566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
4260f0b76eSHong Zhang   /* locate the global i index for x and j index for y */
4360f0b76eSHong Zhang   i = (PetscInt)(x * (Mx - 1));
4460f0b76eSHong Zhang   j = (PetscInt)(y * (My - 1));
459566063dSJacob Faibussowitsch   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
4660f0b76eSHong Zhang 
4760f0b76eSHong Zhang   if (xs <= i && i < xs + xm && ys <= j && j < ys + ym) {
4860f0b76eSHong Zhang     /* the i,j vertex is on this process */
499566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(da, lambda, &l));
5060f0b76eSHong Zhang     l[j][i].u = 1.0;
5160f0b76eSHong Zhang     l[j][i].v = 1.0;
529566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(da, lambda, &l));
5360f0b76eSHong Zhang   }
5460f0b76eSHong Zhang   PetscFunctionReturn(0);
5560f0b76eSHong Zhang }
5660f0b76eSHong Zhang 
57*9371c9d4SSatish Balay static PetscErrorCode MyRHSMatMultTranspose(Mat A_shell, Vec X, Vec Y) {
5860f0b76eSHong Zhang   MCtx       *mctx;
5960f0b76eSHong Zhang   AppCtx     *appctx;
6060f0b76eSHong Zhang   DM          da;
6160f0b76eSHong Zhang   PetscInt    i, j, Mx, My, xs, ys, xm, ym;
6260f0b76eSHong Zhang   PetscReal   hx, hy, sx, sy;
6360f0b76eSHong Zhang   PetscScalar uc, uxx, uyy, vc, vxx, vyy, ucb, vcb;
6460f0b76eSHong Zhang   Field     **u, **x, **y;
6560f0b76eSHong Zhang   Vec         localX;
6660f0b76eSHong Zhang 
6760f0b76eSHong Zhang   PetscFunctionBeginUser;
689566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(A_shell, &mctx));
6960f0b76eSHong Zhang   appctx = mctx->appctx;
709566063dSJacob Faibussowitsch   PetscCall(TSGetDM(mctx->ts, &da));
719566063dSJacob Faibussowitsch   PetscCall(DMGetLocalVector(da, &localX));
729566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
73*9371c9d4SSatish Balay   hx = 2.50 / (PetscReal)Mx;
74*9371c9d4SSatish Balay   sx = 1.0 / (hx * hx);
75*9371c9d4SSatish Balay   hy = 2.50 / (PetscReal)My;
76*9371c9d4SSatish Balay   sy = 1.0 / (hy * hy);
7760f0b76eSHong Zhang 
7860f0b76eSHong Zhang   /* Scatter ghost points to local vector,using the 2-step process
7960f0b76eSHong Zhang      DMGlobalToLocalBegin(),DMGlobalToLocalEnd().
8060f0b76eSHong Zhang      By placing code between these two statements, computations can be
8160f0b76eSHong Zhang      done while messages are in transition. */
829566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX));
839566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX));
8460f0b76eSHong Zhang 
8560f0b76eSHong Zhang   /* Get pointers to vector data */
869566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, localX, &x));
879566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, mctx->U, &u));
889566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(da, Y, &y));
8960f0b76eSHong Zhang 
9060f0b76eSHong Zhang   /* Get local grid boundaries */
919566063dSJacob Faibussowitsch   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
9260f0b76eSHong Zhang 
9360f0b76eSHong Zhang   /* Compute function over the locally owned part of the grid */
9460f0b76eSHong Zhang   for (j = ys; j < ys + ym; j++) {
9560f0b76eSHong Zhang     for (i = xs; i < xs + xm; i++) {
9660f0b76eSHong Zhang       uc        = u[j][i].u;
9760f0b76eSHong Zhang       ucb       = x[j][i].u;
9860f0b76eSHong Zhang       uxx       = (-2.0 * ucb + x[j][i - 1].u + x[j][i + 1].u) * sx;
9960f0b76eSHong Zhang       uyy       = (-2.0 * ucb + x[j - 1][i].u + x[j + 1][i].u) * sy;
10060f0b76eSHong Zhang       vc        = u[j][i].v;
10160f0b76eSHong Zhang       vcb       = x[j][i].v;
10260f0b76eSHong Zhang       vxx       = (-2.0 * vcb + x[j][i - 1].v + x[j][i + 1].v) * sx;
10360f0b76eSHong Zhang       vyy       = (-2.0 * vcb + x[j - 1][i].v + x[j + 1][i].v) * sy;
10460f0b76eSHong Zhang       y[j][i].u = appctx->D1 * (uxx + uyy) - ucb * (vc * vc + appctx->gamma) + vc * vc * vcb;
10560f0b76eSHong Zhang       y[j][i].v = appctx->D2 * (vxx + vyy) - 2.0 * uc * vc * ucb + (2.0 * uc * vc - appctx->gamma - appctx->kappa) * vcb;
10660f0b76eSHong Zhang     }
10760f0b76eSHong Zhang   }
1089566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, localX, &x));
1099566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, mctx->U, &u));
1109566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(da, Y, &y));
1119566063dSJacob Faibussowitsch   PetscCall(DMRestoreLocalVector(da, &localX));
11260f0b76eSHong Zhang   PetscFunctionReturn(0);
11360f0b76eSHong Zhang }
11460f0b76eSHong Zhang 
115*9371c9d4SSatish Balay static PetscErrorCode MyIMatMultTranspose(Mat A_shell, Vec X, Vec Y) {
11660f0b76eSHong Zhang   MCtx       *mctx;
11760f0b76eSHong Zhang   AppCtx     *appctx;
11860f0b76eSHong Zhang   DM          da;
11960f0b76eSHong Zhang   PetscInt    i, j, Mx, My, xs, ys, xm, ym;
12060f0b76eSHong Zhang   PetscReal   hx, hy, sx, sy;
12160f0b76eSHong Zhang   PetscScalar uc, uxx, uyy, vc, vxx, vyy, ucb, vcb;
12260f0b76eSHong Zhang   Field     **u, **x, **y;
12360f0b76eSHong Zhang   Vec         localX;
12460f0b76eSHong Zhang 
12560f0b76eSHong Zhang   PetscFunctionBeginUser;
1269566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(A_shell, &mctx));
12760f0b76eSHong Zhang   appctx = mctx->appctx;
1289566063dSJacob Faibussowitsch   PetscCall(TSGetDM(mctx->ts, &da));
1299566063dSJacob Faibussowitsch   PetscCall(DMGetLocalVector(da, &localX));
1309566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
131*9371c9d4SSatish Balay   hx = 2.50 / (PetscReal)Mx;
132*9371c9d4SSatish Balay   sx = 1.0 / (hx * hx);
133*9371c9d4SSatish Balay   hy = 2.50 / (PetscReal)My;
134*9371c9d4SSatish Balay   sy = 1.0 / (hy * hy);
13560f0b76eSHong Zhang 
13660f0b76eSHong Zhang   /* Scatter ghost points to local vector,using the 2-step process
13760f0b76eSHong Zhang      DMGlobalToLocalBegin(),DMGlobalToLocalEnd().
13860f0b76eSHong Zhang      By placing code between these two statements, computations can be
13960f0b76eSHong Zhang      done while messages are in transition. */
1409566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX));
1419566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX));
14260f0b76eSHong Zhang 
14360f0b76eSHong Zhang   /* Get pointers to vector data */
1449566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, localX, &x));
1459566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, mctx->U, &u));
1469566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(da, Y, &y));
14760f0b76eSHong Zhang 
14860f0b76eSHong Zhang   /* Get local grid boundaries */
1499566063dSJacob Faibussowitsch   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
15060f0b76eSHong Zhang 
15160f0b76eSHong Zhang   /* Compute function over the locally owned part of the grid */
15260f0b76eSHong Zhang   for (j = ys; j < ys + ym; j++) {
15360f0b76eSHong Zhang     for (i = xs; i < xs + xm; i++) {
15460f0b76eSHong Zhang       uc        = u[j][i].u;
15560f0b76eSHong Zhang       ucb       = x[j][i].u;
15660f0b76eSHong Zhang       uxx       = (-2.0 * ucb + x[j][i - 1].u + x[j][i + 1].u) * sx;
15760f0b76eSHong Zhang       uyy       = (-2.0 * ucb + x[j - 1][i].u + x[j + 1][i].u) * sy;
15860f0b76eSHong Zhang       vc        = u[j][i].v;
15960f0b76eSHong Zhang       vcb       = x[j][i].v;
16060f0b76eSHong Zhang       vxx       = (-2.0 * vcb + x[j][i - 1].v + x[j][i + 1].v) * sx;
16160f0b76eSHong Zhang       vyy       = (-2.0 * vcb + x[j - 1][i].v + x[j + 1][i].v) * sy;
16260f0b76eSHong Zhang       y[j][i].u = appctx->D1 * (uxx + uyy) - ucb * (vc * vc + appctx->gamma) + vc * vc * vcb;
16360f0b76eSHong Zhang       y[j][i].v = appctx->D2 * (vxx + vyy) - 2.0 * uc * vc * ucb + (2.0 * uc * vc - appctx->gamma - appctx->kappa) * vcb;
16460f0b76eSHong Zhang       y[j][i].u = mctx->shift * ucb - y[j][i].u;
16560f0b76eSHong Zhang       y[j][i].v = mctx->shift * vcb - y[j][i].v;
16660f0b76eSHong Zhang     }
16760f0b76eSHong Zhang   }
1689566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, localX, &x));
1699566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, mctx->U, &u));
1709566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(da, Y, &y));
1719566063dSJacob Faibussowitsch   PetscCall(DMRestoreLocalVector(da, &localX));
17260f0b76eSHong Zhang   PetscFunctionReturn(0);
17360f0b76eSHong Zhang }
17460f0b76eSHong Zhang 
175*9371c9d4SSatish Balay static PetscErrorCode MyIMatMult(Mat A_shell, Vec X, Vec Y) {
17660f0b76eSHong Zhang   MCtx       *mctx;
17760f0b76eSHong Zhang   AppCtx     *appctx;
17860f0b76eSHong Zhang   DM          da;
17960f0b76eSHong Zhang   PetscInt    i, j, Mx, My, xs, ys, xm, ym;
18060f0b76eSHong Zhang   PetscReal   hx, hy, sx, sy;
18160f0b76eSHong Zhang   PetscScalar uc, uxx, uyy, vc, vxx, vyy, ucb, vcb;
18260f0b76eSHong Zhang   Field     **u, **x, **y;
18360f0b76eSHong Zhang   Vec         localX;
18460f0b76eSHong Zhang 
18560f0b76eSHong Zhang   PetscFunctionBeginUser;
1869566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(A_shell, &mctx));
18760f0b76eSHong Zhang   appctx = mctx->appctx;
1889566063dSJacob Faibussowitsch   PetscCall(TSGetDM(mctx->ts, &da));
1899566063dSJacob Faibussowitsch   PetscCall(DMGetLocalVector(da, &localX));
1909566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
191*9371c9d4SSatish Balay   hx = 2.50 / (PetscReal)Mx;
192*9371c9d4SSatish Balay   sx = 1.0 / (hx * hx);
193*9371c9d4SSatish Balay   hy = 2.50 / (PetscReal)My;
194*9371c9d4SSatish Balay   sy = 1.0 / (hy * hy);
19560f0b76eSHong Zhang 
19660f0b76eSHong Zhang   /* Scatter ghost points to local vector,using the 2-step process
19760f0b76eSHong Zhang      DMGlobalToLocalBegin(),DMGlobalToLocalEnd().
19860f0b76eSHong Zhang      By placing code between these two statements, computations can be
19960f0b76eSHong Zhang      done while messages are in transition. */
2009566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX));
2019566063dSJacob Faibussowitsch   PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX));
20260f0b76eSHong Zhang 
20360f0b76eSHong Zhang   /* Get pointers to vector data */
2049566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, localX, &x));
2059566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArrayRead(da, mctx->U, &u));
2069566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(da, Y, &y));
20760f0b76eSHong Zhang 
20860f0b76eSHong Zhang   /* Get local grid boundaries */
2099566063dSJacob Faibussowitsch   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
21060f0b76eSHong Zhang 
21160f0b76eSHong Zhang   /* Compute function over the locally owned part of the grid */
21260f0b76eSHong Zhang   for (j = ys; j < ys + ym; j++) {
21360f0b76eSHong Zhang     for (i = xs; i < xs + xm; i++) {
21460f0b76eSHong Zhang       uc        = u[j][i].u;
21560f0b76eSHong Zhang       ucb       = x[j][i].u;
21660f0b76eSHong Zhang       uxx       = (-2.0 * ucb + x[j][i - 1].u + x[j][i + 1].u) * sx;
21760f0b76eSHong Zhang       uyy       = (-2.0 * ucb + x[j - 1][i].u + x[j + 1][i].u) * sy;
21860f0b76eSHong Zhang       vc        = u[j][i].v;
21960f0b76eSHong Zhang       vcb       = x[j][i].v;
22060f0b76eSHong Zhang       vxx       = (-2.0 * vcb + x[j][i - 1].v + x[j][i + 1].v) * sx;
22160f0b76eSHong Zhang       vyy       = (-2.0 * vcb + x[j - 1][i].v + x[j + 1][i].v) * sy;
22260f0b76eSHong Zhang       y[j][i].u = appctx->D1 * (uxx + uyy) - (vc * vc + appctx->gamma) * ucb - 2.0 * uc * vc * vcb;
22360f0b76eSHong Zhang       y[j][i].v = appctx->D2 * (vxx + vyy) + vc * vc * ucb + (2.0 * uc * vc - appctx->gamma - appctx->kappa) * vcb;
22460f0b76eSHong Zhang       y[j][i].u = mctx->shift * ucb - y[j][i].u;
22560f0b76eSHong Zhang       y[j][i].v = mctx->shift * vcb - y[j][i].v;
22660f0b76eSHong Zhang     }
22760f0b76eSHong Zhang   }
2289566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, localX, &x));
2299566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArrayRead(da, mctx->U, &u));
2309566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(da, Y, &y));
2319566063dSJacob Faibussowitsch   PetscCall(DMRestoreLocalVector(da, &localX));
23260f0b76eSHong Zhang   PetscFunctionReturn(0);
23360f0b76eSHong Zhang }
23460f0b76eSHong Zhang 
235*9371c9d4SSatish Balay int main(int argc, char **argv) {
23660f0b76eSHong Zhang   TS             ts; /* ODE integrator */
23760f0b76eSHong Zhang   Vec            x;  /* solution */
23860f0b76eSHong Zhang   DM             da;
23960f0b76eSHong Zhang   AppCtx         appctx;
24060f0b76eSHong Zhang   MCtx           mctx;
24160f0b76eSHong Zhang   Vec            lambda[1];
24260f0b76eSHong Zhang   PetscBool      forwardonly = PETSC_FALSE, implicitform = PETSC_TRUE, mf = PETSC_FALSE;
24360f0b76eSHong Zhang   PetscLogDouble v1, v2;
24460f0b76eSHong Zhang #if defined(PETSC_USE_LOG)
24560f0b76eSHong Zhang   PetscLogStage stage;
24660f0b76eSHong Zhang #endif
24760f0b76eSHong Zhang 
248327415f7SBarry Smith   PetscFunctionBeginUser;
2499566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
2509566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL, NULL, "-forwardonly", &forwardonly, NULL));
2519566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL, NULL, "-implicitform", &implicitform, NULL));
25260f0b76eSHong Zhang   appctx.aijpc = PETSC_FALSE;
2539566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL, NULL, "-aijpc", &appctx.aijpc, NULL));
2549566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL, NULL, "-mf", &mf, NULL));
25560f0b76eSHong Zhang 
25660f0b76eSHong Zhang   appctx.D1    = 8.0e-5;
25760f0b76eSHong Zhang   appctx.D2    = 4.0e-5;
25860f0b76eSHong Zhang   appctx.gamma = .024;
25960f0b76eSHong Zhang   appctx.kappa = .06;
26060f0b76eSHong Zhang 
26160f0b76eSHong Zhang   PetscLogStageRegister("MyAdjoint", &stage);
26260f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
26360f0b76eSHong Zhang      Create distributed array (DMDA) to manage parallel grid and vectors
26460f0b76eSHong Zhang   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
2659566063dSJacob Faibussowitsch   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR, 64, 64, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, NULL, &da));
2669566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(da));
2679566063dSJacob Faibussowitsch   PetscCall(DMSetUp(da));
2689566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(da, 0, "u"));
2699566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(da, 1, "v"));
27060f0b76eSHong Zhang 
27160f0b76eSHong Zhang   /*  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
27260f0b76eSHong Zhang      Extract global vectors from DMDA; then duplicate for remaining
27360f0b76eSHong Zhang      vectors that are the same types
27460f0b76eSHong Zhang    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
2759566063dSJacob Faibussowitsch   PetscCall(DMCreateGlobalVector(da, &x));
27660f0b76eSHong Zhang 
27760f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
27860f0b76eSHong Zhang      Create timestepping solver context
27960f0b76eSHong Zhang      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
2809566063dSJacob Faibussowitsch   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2819566063dSJacob Faibussowitsch   PetscCall(TSSetDM(ts, da));
2829566063dSJacob Faibussowitsch   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2839566063dSJacob Faibussowitsch   PetscCall(TSSetEquationType(ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
28460f0b76eSHong Zhang   if (!implicitform) {
2859566063dSJacob Faibussowitsch     PetscCall(TSSetType(ts, TSRK));
2869566063dSJacob Faibussowitsch     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, &appctx));
2879566063dSJacob Faibussowitsch     PetscCall(TSSetRHSJacobian(ts, NULL, NULL, RHSJacobian, &appctx));
28860f0b76eSHong Zhang   } else {
2899566063dSJacob Faibussowitsch     PetscCall(TSSetType(ts, TSCN));
2909566063dSJacob Faibussowitsch     PetscCall(TSSetIFunction(ts, NULL, IFunction, &appctx));
29160f0b76eSHong Zhang     if (appctx.aijpc) {
29260f0b76eSHong Zhang       Mat A, B;
29360f0b76eSHong Zhang 
2949566063dSJacob Faibussowitsch       PetscCall(DMSetMatType(da, MATSELL));
2959566063dSJacob Faibussowitsch       PetscCall(DMCreateMatrix(da, &A));
2969566063dSJacob Faibussowitsch       PetscCall(MatConvert(A, MATAIJ, MAT_INITIAL_MATRIX, &B));
29760f0b76eSHong Zhang       /* FIXME do we need to change viewer to display matrix in natural ordering as DMCreateMatrix_DA does? */
2989566063dSJacob Faibussowitsch       PetscCall(TSSetIJacobian(ts, A, B, IJacobian, &appctx));
2999566063dSJacob Faibussowitsch       PetscCall(MatDestroy(&A));
3009566063dSJacob Faibussowitsch       PetscCall(MatDestroy(&B));
30160f0b76eSHong Zhang     } else {
3029566063dSJacob Faibussowitsch       PetscCall(TSSetIJacobian(ts, NULL, NULL, IJacobian, &appctx));
30360f0b76eSHong Zhang     }
30460f0b76eSHong Zhang   }
30560f0b76eSHong Zhang 
30660f0b76eSHong Zhang   if (mf) {
30760f0b76eSHong Zhang     PetscInt xm, ym, Mx, My, dof;
30860f0b76eSHong Zhang     mctx.ts     = ts;
30960f0b76eSHong Zhang     mctx.appctx = &appctx;
3109566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(x, &mctx.U));
31160f0b76eSHong Zhang     /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
31260f0b76eSHong Zhang       Create matrix free context
31360f0b76eSHong Zhang       - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3149566063dSJacob Faibussowitsch     PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, &dof, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
3159566063dSJacob Faibussowitsch     PetscCall(DMDAGetCorners(da, NULL, NULL, NULL, &xm, &ym, NULL));
3169566063dSJacob Faibussowitsch     PetscCall(MatCreateShell(PETSC_COMM_WORLD, dof * xm * ym, PETSC_DETERMINE, dof * Mx * My, dof * Mx * My, &mctx, &appctx.A));
3179566063dSJacob Faibussowitsch     PetscCall(MatShellSetOperation(appctx.A, MATOP_MULT_TRANSPOSE, (void (*)(void))MyRHSMatMultTranspose));
31860f0b76eSHong Zhang     if (!implicitform) { /* for explicit methods only */
3199566063dSJacob Faibussowitsch       PetscCall(TSSetRHSJacobian(ts, appctx.A, appctx.A, RHSJacobianShell, &appctx));
32060f0b76eSHong Zhang     } else {
3219566063dSJacob Faibussowitsch       /* PetscCall(VecDuplicate(appctx.U,&mctx.Udot)); */
3229566063dSJacob Faibussowitsch       PetscCall(MatShellSetOperation(appctx.A, MATOP_MULT, (void (*)(void))MyIMatMult));
3239566063dSJacob Faibussowitsch       PetscCall(MatShellSetOperation(appctx.A, MATOP_MULT_TRANSPOSE, (void (*)(void))MyIMatMultTranspose));
3249566063dSJacob Faibussowitsch       PetscCall(TSSetIJacobian(ts, appctx.A, appctx.A, IJacobianShell, &appctx));
32560f0b76eSHong Zhang     }
32660f0b76eSHong Zhang   }
32760f0b76eSHong Zhang 
32860f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
32960f0b76eSHong Zhang      Set initial conditions
33060f0b76eSHong Zhang    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3319566063dSJacob Faibussowitsch   PetscCall(InitialConditions(da, x));
3329566063dSJacob Faibussowitsch   PetscCall(TSSetSolution(ts, x));
33360f0b76eSHong Zhang 
33460f0b76eSHong Zhang   /*
33560f0b76eSHong Zhang     Have the TS save its trajectory so that TSAdjointSolve() may be used
33660f0b76eSHong Zhang   */
3379566063dSJacob Faibussowitsch   if (!forwardonly) PetscCall(TSSetSaveTrajectory(ts));
33860f0b76eSHong Zhang 
33960f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
34060f0b76eSHong Zhang      Set solver options
34160f0b76eSHong Zhang    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3429566063dSJacob Faibussowitsch   PetscCall(TSSetMaxTime(ts, 200.0));
3439566063dSJacob Faibussowitsch   PetscCall(TSSetTimeStep(ts, 0.5));
3449566063dSJacob Faibussowitsch   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
3459566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(ts));
34660f0b76eSHong Zhang 
3479566063dSJacob Faibussowitsch   PetscCall(PetscTime(&v1));
34860f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
34960f0b76eSHong Zhang      Solve ODE system
35060f0b76eSHong Zhang      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3519566063dSJacob Faibussowitsch   PetscCall(TSSolve(ts, x));
35260f0b76eSHong Zhang   if (!forwardonly) {
35360f0b76eSHong Zhang     /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
35460f0b76eSHong Zhang        Start the Adjoint model
35560f0b76eSHong Zhang        - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3569566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(x, &lambda[0]));
35760f0b76eSHong Zhang     /*   Reset initial conditions for the adjoint integration */
3589566063dSJacob Faibussowitsch     PetscCall(InitializeLambda(da, lambda[0], 0.5, 0.5));
3599566063dSJacob Faibussowitsch     PetscCall(TSSetCostGradients(ts, 1, lambda, NULL));
36060f0b76eSHong Zhang     PetscLogStagePush(stage);
3619566063dSJacob Faibussowitsch     PetscCall(TSAdjointSolve(ts));
36260f0b76eSHong Zhang     PetscLogStagePop();
3639566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&lambda[0]));
36460f0b76eSHong Zhang   }
3659566063dSJacob Faibussowitsch   PetscCall(PetscTime(&v2));
3669566063dSJacob Faibussowitsch   PetscCall(PetscPrintf(PETSC_COMM_WORLD, " %.3lf ", v2 - v1));
36760f0b76eSHong Zhang 
36860f0b76eSHong Zhang   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
36960f0b76eSHong Zhang      Free work space.  All PETSc objects should be destroyed when they
37060f0b76eSHong Zhang      are no longer needed.
37160f0b76eSHong Zhang    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
3729566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&x));
3739566063dSJacob Faibussowitsch   PetscCall(TSDestroy(&ts));
3749566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&da));
37560f0b76eSHong Zhang   if (mf) {
3769566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&mctx.U));
3779566063dSJacob Faibussowitsch     /* PetscCall(VecDestroy(&mctx.Udot));*/
3789566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&appctx.A));
37960f0b76eSHong Zhang   }
3809566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
381b122ec5aSJacob Faibussowitsch   return 0;
38260f0b76eSHong Zhang }
38360f0b76eSHong Zhang 
38460f0b76eSHong Zhang /* ------------------------------------------------------------------- */
385*9371c9d4SSatish Balay PetscErrorCode RHSJacobianShell(TS ts, PetscReal t, Vec U, Mat A, Mat BB, void *ctx) {
38660f0b76eSHong Zhang   MCtx *mctx;
38760f0b76eSHong Zhang 
38860f0b76eSHong Zhang   PetscFunctionBegin;
3899566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(A, &mctx));
3909566063dSJacob Faibussowitsch   PetscCall(VecCopy(U, mctx->U));
39160f0b76eSHong Zhang   PetscFunctionReturn(0);
39260f0b76eSHong Zhang }
39360f0b76eSHong Zhang 
394*9371c9d4SSatish Balay PetscErrorCode IJacobianShell(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal a, Mat A, Mat BB, void *ctx) {
39560f0b76eSHong Zhang   MCtx *mctx;
39660f0b76eSHong Zhang 
39760f0b76eSHong Zhang   PetscFunctionBegin;
3989566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(A, &mctx));
3999566063dSJacob Faibussowitsch   PetscCall(VecCopy(U, mctx->U));
4009566063dSJacob Faibussowitsch   /* PetscCall(VecCopy(Udot,mctx->Udot)); */
40160f0b76eSHong Zhang   mctx->shift = a;
40260f0b76eSHong Zhang   PetscFunctionReturn(0);
40360f0b76eSHong Zhang }
40460f0b76eSHong Zhang 
40560f0b76eSHong Zhang /* ------------------------------------------------------------------- */
406*9371c9d4SSatish Balay PetscErrorCode InitialConditions(DM da, Vec U) {
40760f0b76eSHong Zhang   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
40860f0b76eSHong Zhang   Field   **u;
40960f0b76eSHong Zhang   PetscReal hx, hy, x, y;
41060f0b76eSHong Zhang 
41160f0b76eSHong Zhang   PetscFunctionBegin;
4129566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
41360f0b76eSHong Zhang 
41460f0b76eSHong Zhang   hx = 2.5 / (PetscReal)Mx;
41560f0b76eSHong Zhang   hy = 2.5 / (PetscReal)My;
41660f0b76eSHong Zhang 
41760f0b76eSHong Zhang   /*
41860f0b76eSHong Zhang      Get pointers to vector data
41960f0b76eSHong Zhang   */
4209566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(da, U, &u));
42160f0b76eSHong Zhang 
42260f0b76eSHong Zhang   /*
42360f0b76eSHong Zhang      Get local grid boundaries
42460f0b76eSHong Zhang   */
4259566063dSJacob Faibussowitsch   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
42660f0b76eSHong Zhang 
42760f0b76eSHong Zhang   /*
42860f0b76eSHong Zhang      Compute function over the locally owned part of the grid
42960f0b76eSHong Zhang   */
43060f0b76eSHong Zhang   for (j = ys; j < ys + ym; j++) {
43160f0b76eSHong Zhang     y = j * hy;
43260f0b76eSHong Zhang     for (i = xs; i < xs + xm; i++) {
43360f0b76eSHong Zhang       x = i * hx;
43460f0b76eSHong Zhang       if ((1.0 <= x) && (x <= 1.5) && (1.0 <= y) && (y <= 1.5)) u[j][i].v = .25 * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0);
43560f0b76eSHong Zhang       else u[j][i].v = 0.0;
43660f0b76eSHong Zhang 
43760f0b76eSHong Zhang       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
43860f0b76eSHong Zhang     }
43960f0b76eSHong Zhang   }
44060f0b76eSHong Zhang 
44160f0b76eSHong Zhang   /*
44260f0b76eSHong Zhang      Restore vectors
44360f0b76eSHong Zhang   */
4449566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(da, U, &u));
44560f0b76eSHong Zhang   PetscFunctionReturn(0);
44660f0b76eSHong Zhang }
44760f0b76eSHong Zhang 
44860f0b76eSHong Zhang /*TEST
44960f0b76eSHong Zhang 
45060f0b76eSHong Zhang    build:
45160f0b76eSHong Zhang       depends: reaction_diffusion.c
45260f0b76eSHong Zhang       requires: !complex !single
45360f0b76eSHong Zhang 
45460f0b76eSHong Zhang    test:
45560f0b76eSHong Zhang       requires: cams
456533251d3SHong Zhang       args: -mf -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor -ts_trajectory_type memory -ts_trajectory_solution_only 0 -ts_trajectory_max_units_ram 6 -ts_trajectory_memory_type cams
45760f0b76eSHong Zhang TEST*/
458