xref: /petsc/src/snes/tutorials/ex28.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
1c4762a1bSJed Brown static const char help[] = "1D multiphysics prototype with analytic Jacobians to solve individual problems and a coupled problem.\n\n";
2c4762a1bSJed Brown 
3c4762a1bSJed Brown /* Solve a PDE coupled to an algebraic system in 1D
4c4762a1bSJed Brown  *
5c4762a1bSJed Brown  * PDE (U):
6c4762a1bSJed Brown  *     -(k u_x)_x = 1 on (0,1), subject to u(0) = 0, u(1) = 1
7c4762a1bSJed Brown  * Algebraic (K):
8c4762a1bSJed Brown  *     exp(k-1) + k = 1/(1/(1+u) + 1/(1+u_x^2))
9c4762a1bSJed Brown  *
10c4762a1bSJed Brown  * The discretization places k at staggered points, and a separate DMDA is used for each "physics".
11c4762a1bSJed Brown  *
12c4762a1bSJed Brown  * This example is a prototype for coupling in multi-physics problems, therefore residual evaluation and assembly for
13c4762a1bSJed Brown  * each problem (referred to as U and K) are written separately.  This permits the same "physics" code to be used for
14c4762a1bSJed Brown  * solving each uncoupled problem as well as the coupled system.  In particular, run with -problem_type 0 to solve only
15c4762a1bSJed Brown  * problem U (with K fixed), -problem_type 1 to solve only K (with U fixed), and -problem_type 2 to solve both at once.
16c4762a1bSJed Brown  *
17c4762a1bSJed Brown  * In all cases, a fully-assembled analytic Jacobian is available, so the systems can be solved with a direct solve or
18c4762a1bSJed Brown  * any other standard method.  Additionally, by running with
19c4762a1bSJed Brown  *
20c4762a1bSJed Brown  *   -pack_dm_mat_type nest
21c4762a1bSJed Brown  *
22c4762a1bSJed Brown  * The same code assembles a coupled matrix where each block is stored separately, which allows the use of PCFieldSplit
23c4762a1bSJed Brown  * without copying values to extract submatrices.
24c4762a1bSJed Brown  */
25c4762a1bSJed Brown 
26c4762a1bSJed Brown #include <petscsnes.h>
27c4762a1bSJed Brown #include <petscdm.h>
28c4762a1bSJed Brown #include <petscdmda.h>
29c4762a1bSJed Brown #include <petscdmcomposite.h>
30c4762a1bSJed Brown 
31c4762a1bSJed Brown typedef struct _UserCtx *User;
32c4762a1bSJed Brown struct _UserCtx {
33c4762a1bSJed Brown   PetscInt ptype;
34c4762a1bSJed Brown   DM       pack;
35c4762a1bSJed Brown   Vec      Uloc, Kloc;
36c4762a1bSJed Brown };
37c4762a1bSJed Brown 
38*9371c9d4SSatish Balay static PetscErrorCode FormFunctionLocal_U(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], PetscScalar f[]) {
39c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
40c4762a1bSJed Brown   PetscInt  i;
41c4762a1bSJed Brown 
42c4762a1bSJed Brown   PetscFunctionBeginUser;
43c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
44c4762a1bSJed Brown     if (i == 0) f[i] = 1. / hx * u[i];
45c4762a1bSJed Brown     else if (i == info->mx - 1) f[i] = 1. / hx * (u[i] - 1.0);
46c4762a1bSJed Brown     else f[i] = hx * ((k[i - 1] * (u[i] - u[i - 1]) - k[i] * (u[i + 1] - u[i])) / (hx * hx) - 1.0);
47c4762a1bSJed Brown   }
48c4762a1bSJed Brown   PetscFunctionReturn(0);
49c4762a1bSJed Brown }
50c4762a1bSJed Brown 
51*9371c9d4SSatish Balay static PetscErrorCode FormFunctionLocal_K(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], PetscScalar f[]) {
52c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
53c4762a1bSJed Brown   PetscInt  i;
54c4762a1bSJed Brown 
55c4762a1bSJed Brown   PetscFunctionBeginUser;
56c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
57*9371c9d4SSatish Balay     const PetscScalar ubar = 0.5 * (u[i + 1] + u[i]), gradu = (u[i + 1] - u[i]) / hx, g = 1. + gradu * gradu, w = 1. / (1. + ubar) + 1. / g;
58c4762a1bSJed Brown     f[i] = hx * (PetscExpScalar(k[i] - 1.0) + k[i] - 1. / w);
59c4762a1bSJed Brown   }
60c4762a1bSJed Brown   PetscFunctionReturn(0);
61c4762a1bSJed Brown }
62c4762a1bSJed Brown 
63*9371c9d4SSatish Balay static PetscErrorCode FormFunction_All(SNES snes, Vec X, Vec F, void *ctx) {
64c4762a1bSJed Brown   User          user = (User)ctx;
65c4762a1bSJed Brown   DM            dau, dak;
66c4762a1bSJed Brown   DMDALocalInfo infou, infok;
67c4762a1bSJed Brown   PetscScalar  *u, *k;
68c4762a1bSJed Brown   PetscScalar  *fu, *fk;
69c4762a1bSJed Brown   Vec           Uloc, Kloc, Fu, Fk;
70c4762a1bSJed Brown 
71c4762a1bSJed Brown   PetscFunctionBeginUser;
729566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
739566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
749566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
759566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack, &Uloc, &Kloc));
76c4762a1bSJed Brown   switch (user->ptype) {
77c4762a1bSJed Brown   case 0:
789566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau, X, INSERT_VALUES, Uloc));
799566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dau, X, INSERT_VALUES, Uloc));
809566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
819566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, user->Kloc, &k));
829566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, F, &fu));
839566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user, &infou, u, k, fu));
849566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, F, &fu));
859566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
869566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, user->Kloc, &k));
87c4762a1bSJed Brown     break;
88c4762a1bSJed Brown   case 1:
899566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak, X, INSERT_VALUES, Kloc));
909566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dak, X, INSERT_VALUES, Kloc));
919566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, user->Uloc, &u));
929566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
939566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, F, &fk));
949566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user, &infok, u, k, fk));
959566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, F, &fk));
969566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, user->Uloc, &u));
979566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
98c4762a1bSJed Brown     break;
99c4762a1bSJed Brown   case 2:
1009566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack, X, Uloc, Kloc));
1019566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
1029566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
1039566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(user->pack, F, &Fu, &Fk));
1049566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Fu, &fu));
1059566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Fk, &fk));
1069566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user, &infou, u, k, fu));
1079566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user, &infok, u, k, fk));
1089566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Fu, &fu));
1099566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Fk, &fk));
1109566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(user->pack, F, &Fu, &Fk));
1119566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
1129566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
113c4762a1bSJed Brown     break;
114c4762a1bSJed Brown   }
1159566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack, &Uloc, &Kloc));
116c4762a1bSJed Brown   PetscFunctionReturn(0);
117c4762a1bSJed Brown }
118c4762a1bSJed Brown 
119*9371c9d4SSatish Balay static PetscErrorCode FormJacobianLocal_U(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], Mat Buu) {
120c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
121c4762a1bSJed Brown   PetscInt  i;
122c4762a1bSJed Brown 
123c4762a1bSJed Brown   PetscFunctionBeginUser;
124c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
125c4762a1bSJed Brown     PetscInt    row = i - info->gxs, cols[] = {row - 1, row, row + 1};
126c4762a1bSJed Brown     PetscScalar val = 1. / hx;
127c4762a1bSJed Brown     if (i == 0) {
1289566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 1, &row, &val, INSERT_VALUES));
129c4762a1bSJed Brown     } else if (i == info->mx - 1) {
1309566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 1, &row, &val, INSERT_VALUES));
131c4762a1bSJed Brown     } else {
132c4762a1bSJed Brown       PetscScalar vals[] = {-k[i - 1] / hx, (k[i - 1] + k[i]) / hx, -k[i] / hx};
1339566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 3, cols, vals, INSERT_VALUES));
134c4762a1bSJed Brown     }
135c4762a1bSJed Brown   }
136c4762a1bSJed Brown   PetscFunctionReturn(0);
137c4762a1bSJed Brown }
138c4762a1bSJed Brown 
139*9371c9d4SSatish Balay static PetscErrorCode FormJacobianLocal_K(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], Mat Bkk) {
140c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
141c4762a1bSJed Brown   PetscInt  i;
142c4762a1bSJed Brown 
143c4762a1bSJed Brown   PetscFunctionBeginUser;
144c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
145c4762a1bSJed Brown     PetscInt    row    = i - info->gxs;
146c4762a1bSJed Brown     PetscScalar vals[] = {hx * (PetscExpScalar(k[i] - 1.) + (PetscScalar)1.)};
1479566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bkk, 1, &row, 1, &row, vals, INSERT_VALUES));
148c4762a1bSJed Brown   }
149c4762a1bSJed Brown   PetscFunctionReturn(0);
150c4762a1bSJed Brown }
151c4762a1bSJed Brown 
152*9371c9d4SSatish Balay static PetscErrorCode FormJacobianLocal_UK(User user, DMDALocalInfo *info, DMDALocalInfo *infok, const PetscScalar u[], const PetscScalar k[], Mat Buk) {
153c4762a1bSJed Brown   PetscReal   hx = 1. / info->mx;
154c4762a1bSJed Brown   PetscInt    i;
155c4762a1bSJed Brown   PetscInt    row, cols[2];
156c4762a1bSJed Brown   PetscScalar vals[2];
157c4762a1bSJed Brown 
158c4762a1bSJed Brown   PetscFunctionBeginUser;
159c4762a1bSJed Brown   if (!Buk) PetscFunctionReturn(0); /* Not assembling this block */
160c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
161c4762a1bSJed Brown     if (i == 0 || i == info->mx - 1) continue;
162c4762a1bSJed Brown     row     = i - info->gxs;
163*9371c9d4SSatish Balay     cols[0] = i - 1 - infok->gxs;
164*9371c9d4SSatish Balay     vals[0] = (u[i] - u[i - 1]) / hx;
165*9371c9d4SSatish Balay     cols[1] = i - infok->gxs;
166*9371c9d4SSatish Balay     vals[1] = (u[i] - u[i + 1]) / hx;
1679566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Buk, 1, &row, 2, cols, vals, INSERT_VALUES));
168c4762a1bSJed Brown   }
169c4762a1bSJed Brown   PetscFunctionReturn(0);
170c4762a1bSJed Brown }
171c4762a1bSJed Brown 
172*9371c9d4SSatish Balay static PetscErrorCode FormJacobianLocal_KU(User user, DMDALocalInfo *info, DMDALocalInfo *infok, const PetscScalar u[], const PetscScalar k[], Mat Bku) {
173c4762a1bSJed Brown   PetscInt  i;
174c4762a1bSJed Brown   PetscReal hx = 1. / (info->mx - 1);
175c4762a1bSJed Brown 
176c4762a1bSJed Brown   PetscFunctionBeginUser;
177c4762a1bSJed Brown   if (!Bku) PetscFunctionReturn(0); /* Not assembling this block */
178c4762a1bSJed Brown   for (i = infok->xs; i < infok->xs + infok->xm; i++) {
179c4762a1bSJed Brown     PetscInt          row = i - infok->gxs, cols[2];
180c4762a1bSJed Brown     PetscScalar       vals[2];
181*9371c9d4SSatish Balay     const PetscScalar ubar = 0.5 * (u[i] + u[i + 1]), ubar_L = 0.5, ubar_R = 0.5, gradu = (u[i + 1] - u[i]) / hx, gradu_L = -1. / hx, gradu_R = 1. / hx, g = 1. + PetscSqr(gradu), g_gradu = 2. * gradu, w = 1. / (1. + ubar) + 1. / g, w_ubar = -1. / PetscSqr(1. + ubar), w_gradu = -g_gradu / PetscSqr(g), iw = 1. / w, iw_ubar = -w_ubar * PetscSqr(iw), iw_gradu = -w_gradu * PetscSqr(iw);
182*9371c9d4SSatish Balay     cols[0] = i - info->gxs;
183*9371c9d4SSatish Balay     vals[0] = -hx * (iw_ubar * ubar_L + iw_gradu * gradu_L);
184*9371c9d4SSatish Balay     cols[1] = i + 1 - info->gxs;
185*9371c9d4SSatish Balay     vals[1] = -hx * (iw_ubar * ubar_R + iw_gradu * gradu_R);
1869566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bku, 1, &row, 2, cols, vals, INSERT_VALUES));
187c4762a1bSJed Brown   }
188c4762a1bSJed Brown   PetscFunctionReturn(0);
189c4762a1bSJed Brown }
190c4762a1bSJed Brown 
191*9371c9d4SSatish Balay static PetscErrorCode FormJacobian_All(SNES snes, Vec X, Mat J, Mat B, void *ctx) {
192c4762a1bSJed Brown   User          user = (User)ctx;
193c4762a1bSJed Brown   DM            dau, dak;
194c4762a1bSJed Brown   DMDALocalInfo infou, infok;
195c4762a1bSJed Brown   PetscScalar  *u, *k;
196c4762a1bSJed Brown   Vec           Uloc, Kloc;
197c4762a1bSJed Brown 
198c4762a1bSJed Brown   PetscFunctionBeginUser;
1999566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
2009566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
2019566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
2029566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack, &Uloc, &Kloc));
203c4762a1bSJed Brown   switch (user->ptype) {
204c4762a1bSJed Brown   case 0:
2059566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau, X, INSERT_VALUES, Uloc));
2069566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dau, X, INSERT_VALUES, Uloc));
2079566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
2089566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, user->Kloc, &k));
2099566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user, &infou, u, k, B));
2109566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
2119566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, user->Kloc, &k));
212c4762a1bSJed Brown     break;
213c4762a1bSJed Brown   case 1:
2149566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak, X, INSERT_VALUES, Kloc));
2159566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dak, X, INSERT_VALUES, Kloc));
2169566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, user->Uloc, &u));
2179566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
2189566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user, &infok, u, k, B));
2199566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, user->Uloc, &u));
2209566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
221c4762a1bSJed Brown     break;
222c4762a1bSJed Brown   case 2: {
223c4762a1bSJed Brown     Mat       Buu, Buk, Bku, Bkk;
224c4762a1bSJed Brown     PetscBool nest;
225c4762a1bSJed Brown     IS       *is;
2269566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack, X, Uloc, Kloc));
2279566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
2289566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
2299566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetLocalISs(user->pack, &is));
2309566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[0], is[0], &Buu));
2319566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[0], is[1], &Buk));
2329566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[1], is[0], &Bku));
2339566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[1], is[1], &Bkk));
2349566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user, &infou, u, k, Buu));
2359566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)B, MATNEST, &nest));
236c4762a1bSJed Brown     if (!nest) {
237c4762a1bSJed Brown       /*
238c4762a1bSJed Brown          DMCreateMatrix_Composite()  for a nested matrix does not generate off-block matrices that one can call MatSetValuesLocal() on, it just creates dummy
239c4762a1bSJed Brown          matrices with no entries; there cannot insert values into them. Commit b6480e041dd2293a65f96222772d68cdb4ed6306
240c4762a1bSJed Brown          changed Mat_Nest() from returning NULL pointers for these submatrices to dummy matrices because PCFIELDSPLIT could not
241c4762a1bSJed Brown          handle the returned null matrices.
242c4762a1bSJed Brown       */
2439566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_UK(user, &infou, &infok, u, k, Buk));
2449566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_KU(user, &infou, &infok, u, k, Bku));
245c4762a1bSJed Brown     }
2469566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user, &infok, u, k, Bkk));
2479566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[0], is[0], &Buu));
2489566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[0], is[1], &Buk));
2499566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[1], is[0], &Bku));
2509566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[1], is[1], &Bkk));
2519566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
2529566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
253c4762a1bSJed Brown 
2549566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[0]));
2559566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[1]));
2569566063dSJacob Faibussowitsch     PetscCall(PetscFree(is));
257c4762a1bSJed Brown   } break;
258c4762a1bSJed Brown   }
2599566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack, &Uloc, &Kloc));
2609566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2619566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
262c4762a1bSJed Brown   if (J != B) {
2639566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
2649566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
265c4762a1bSJed Brown   }
266c4762a1bSJed Brown   PetscFunctionReturn(0);
267c4762a1bSJed Brown }
268c4762a1bSJed Brown 
269*9371c9d4SSatish Balay static PetscErrorCode FormInitial_Coupled(User user, Vec X) {
270c4762a1bSJed Brown   DM            dau, dak;
271c4762a1bSJed Brown   DMDALocalInfo infou, infok;
272c4762a1bSJed Brown   Vec           Xu, Xk;
273c4762a1bSJed Brown   PetscScalar  *u, *k, hx;
274c4762a1bSJed Brown   PetscInt      i;
275c4762a1bSJed Brown 
276c4762a1bSJed Brown   PetscFunctionBeginUser;
2779566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
2789566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetAccess(user->pack, X, &Xu, &Xk));
2799566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dau, Xu, &u));
2809566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dak, Xk, &k));
2819566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
2829566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
283c4762a1bSJed Brown   hx = 1. / (infok.mx);
284c4762a1bSJed Brown   for (i = infou.xs; i < infou.xs + infou.xm; i++) u[i] = (PetscScalar)i * hx * (1. - (PetscScalar)i * hx);
285c4762a1bSJed Brown   for (i = infok.xs; i < infok.xs + infok.xm; i++) k[i] = 1.0 + 0.5 * PetscSinScalar(2 * PETSC_PI * i * hx);
2869566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dau, Xu, &u));
2879566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dak, Xk, &k));
2889566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreAccess(user->pack, X, &Xu, &Xk));
2899566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(user->pack, X, user->Uloc, user->Kloc));
290c4762a1bSJed Brown   PetscFunctionReturn(0);
291c4762a1bSJed Brown }
292c4762a1bSJed Brown 
293*9371c9d4SSatish Balay int main(int argc, char *argv[]) {
294c4762a1bSJed Brown   DM              dau, dak, pack;
295c4762a1bSJed Brown   const PetscInt *lxu;
296c4762a1bSJed Brown   PetscInt       *lxk, m, sizes;
297c4762a1bSJed Brown   User            user;
298c4762a1bSJed Brown   SNES            snes;
299c4762a1bSJed Brown   Vec             X, F, Xu, Xk, Fu, Fk;
300c4762a1bSJed Brown   Mat             B;
301c4762a1bSJed Brown   IS             *isg;
302c4762a1bSJed Brown   PetscBool       pass_dm;
303c4762a1bSJed Brown 
304327415f7SBarry Smith   PetscFunctionBeginUser;
3059566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc, &argv, 0, help));
3069566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 10, 1, 1, NULL, &dau));
3079566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dau, "u_"));
3089566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dau));
3099566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dau));
3109566063dSJacob Faibussowitsch   PetscCall(DMDAGetOwnershipRanges(dau, &lxu, 0, 0));
3119566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(dau, 0, &m, 0, 0, &sizes, 0, 0, 0, 0, 0, 0, 0, 0));
3129566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(sizes, &lxk));
3139566063dSJacob Faibussowitsch   PetscCall(PetscArraycpy(lxk, lxu, sizes));
314c4762a1bSJed Brown   lxk[0]--;
3159566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, m - 1, 1, 1, lxk, &dak));
3169566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dak, "k_"));
3179566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dak));
3189566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dak));
3199566063dSJacob Faibussowitsch   PetscCall(PetscFree(lxk));
320c4762a1bSJed Brown 
3219566063dSJacob Faibussowitsch   PetscCall(DMCompositeCreate(PETSC_COMM_WORLD, &pack));
3229566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(pack, "pack_"));
3239566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack, dau));
3249566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack, dak));
3259566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dau, 0, "u"));
3269566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dak, 0, "k"));
3279566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(pack));
328c4762a1bSJed Brown 
3299566063dSJacob Faibussowitsch   PetscCall(DMCreateGlobalVector(pack, &X));
3309566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(X, &F));
331c4762a1bSJed Brown 
3329566063dSJacob Faibussowitsch   PetscCall(PetscNew(&user));
333c4762a1bSJed Brown 
334c4762a1bSJed Brown   user->pack = pack;
335c4762a1bSJed Brown 
3369566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetGlobalISs(pack, &isg));
3379566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(pack, &user->Uloc, &user->Kloc));
3389566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(pack, X, user->Uloc, user->Kloc));
339c4762a1bSJed Brown 
340d0609cedSBarry Smith   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Coupled problem options", "SNES");
341c4762a1bSJed Brown   {
342*9371c9d4SSatish Balay     user->ptype = 0;
343*9371c9d4SSatish Balay     pass_dm     = PETSC_TRUE;
344c4762a1bSJed Brown 
3459566063dSJacob Faibussowitsch     PetscCall(PetscOptionsInt("-problem_type", "0: solve for u only, 1: solve for k only, 2: solve for both", 0, user->ptype, &user->ptype, NULL));
3469566063dSJacob Faibussowitsch     PetscCall(PetscOptionsBool("-pass_dm", "Pass the packed DM to SNES to use when determining splits and forward into splits", 0, pass_dm, &pass_dm, NULL));
347c4762a1bSJed Brown   }
348d0609cedSBarry Smith   PetscOptionsEnd();
349c4762a1bSJed Brown 
3509566063dSJacob Faibussowitsch   PetscCall(FormInitial_Coupled(user, X));
351c4762a1bSJed Brown 
3529566063dSJacob Faibussowitsch   PetscCall(SNESCreate(PETSC_COMM_WORLD, &snes));
353c4762a1bSJed Brown   switch (user->ptype) {
354c4762a1bSJed Brown   case 0:
3559566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, X, &Xu, 0));
3569566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, F, &Fu, 0));
3579566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dau, &B));
3589566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, Fu, FormFunction_All, user));
3599566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3609566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3619566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes, dau));
3629566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Xu));
3639566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, X, &Xu, 0));
3649566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, F, &Fu, 0));
365c4762a1bSJed Brown     break;
366c4762a1bSJed Brown   case 1:
3679566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, X, 0, &Xk));
3689566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, F, 0, &Fk));
3699566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dak, &B));
3709566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, Fk, FormFunction_All, user));
3719566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3729566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3739566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes, dak));
3749566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Xk));
3759566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, X, 0, &Xk));
3769566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, F, 0, &Fk));
377c4762a1bSJed Brown     break;
378c4762a1bSJed Brown   case 2:
3799566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(pack, &B));
380c4762a1bSJed Brown     /* This example does not correctly allocate off-diagonal blocks. These options allows new nonzeros (slow). */
3819566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_FALSE));
3829566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE));
3839566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, F, FormFunction_All, user));
3849566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3859566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
386c4762a1bSJed Brown     if (!pass_dm) { /* Manually provide index sets and names for the splits */
387c4762a1bSJed Brown       KSP ksp;
388c4762a1bSJed Brown       PC  pc;
3899566063dSJacob Faibussowitsch       PetscCall(SNESGetKSP(snes, &ksp));
3909566063dSJacob Faibussowitsch       PetscCall(KSPGetPC(ksp, &pc));
3919566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc, "u", isg[0]));
3929566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc, "k", isg[1]));
393c4762a1bSJed Brown     } else {
394c4762a1bSJed Brown       /* The same names come from the options prefix for dau and dak. This option can support geometric multigrid inside
395c4762a1bSJed Brown        * of splits, but it requires using a DM (perhaps your own implementation). */
3969566063dSJacob Faibussowitsch       PetscCall(SNESSetDM(snes, pack));
397c4762a1bSJed Brown     }
3989566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, X));
399c4762a1bSJed Brown     break;
400c4762a1bSJed Brown   }
4019566063dSJacob Faibussowitsch   PetscCall(VecViewFromOptions(X, NULL, "-view_sol"));
402c4762a1bSJed Brown 
403c4762a1bSJed Brown   if (0) {
404c4762a1bSJed Brown     PetscInt  col      = 0;
405c4762a1bSJed Brown     PetscBool mult_dup = PETSC_FALSE, view_dup = PETSC_FALSE;
406c4762a1bSJed Brown     Mat       D;
407c4762a1bSJed Brown     Vec       Y;
408c4762a1bSJed Brown 
4099566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetInt(NULL, 0, "-col", &col, 0));
4109566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL, 0, "-mult_dup", &mult_dup, 0));
4119566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL, 0, "-view_dup", &view_dup, 0));
412c4762a1bSJed Brown 
4139566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X, &Y));
4149566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); */
4159566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); */
4169566063dSJacob Faibussowitsch     PetscCall(MatConvert(B, MATAIJ, MAT_INITIAL_MATRIX, &D));
4179566063dSJacob Faibussowitsch     PetscCall(VecZeroEntries(X));
4189566063dSJacob Faibussowitsch     PetscCall(VecSetValue(X, col, 1.0, INSERT_VALUES));
4199566063dSJacob Faibussowitsch     PetscCall(VecAssemblyBegin(X));
4209566063dSJacob Faibussowitsch     PetscCall(VecAssemblyEnd(X));
4219566063dSJacob Faibussowitsch     PetscCall(MatMult(mult_dup ? D : B, X, Y));
4229566063dSJacob Faibussowitsch     PetscCall(MatView(view_dup ? D : B, PETSC_VIEWER_STDOUT_WORLD));
4239566063dSJacob Faibussowitsch     /* PetscCall(VecView(X,PETSC_VIEWER_STDOUT_WORLD)); */
4249566063dSJacob Faibussowitsch     PetscCall(VecView(Y, PETSC_VIEWER_STDOUT_WORLD));
4259566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&D));
4269566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&Y));
427c4762a1bSJed Brown   }
428c4762a1bSJed Brown 
4299566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(pack, &user->Uloc, &user->Kloc));
4309566063dSJacob Faibussowitsch   PetscCall(PetscFree(user));
431c4762a1bSJed Brown 
4329566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[0]));
4339566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[1]));
4349566063dSJacob Faibussowitsch   PetscCall(PetscFree(isg));
4359566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X));
4369566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&F));
4379566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
4389566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dau));
4399566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dak));
4409566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&pack));
4419566063dSJacob Faibussowitsch   PetscCall(SNESDestroy(&snes));
4429566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
443b122ec5aSJacob Faibussowitsch   return 0;
444c4762a1bSJed Brown }
445c4762a1bSJed Brown 
446c4762a1bSJed Brown /*TEST
447c4762a1bSJed Brown 
448c4762a1bSJed Brown    test:
449c4762a1bSJed Brown       suffix: 0
450c4762a1bSJed Brown       nsize: 3
451c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 0
452c4762a1bSJed Brown 
453c4762a1bSJed Brown    test:
454c4762a1bSJed Brown       suffix: 1
455c4762a1bSJed Brown       nsize: 3
456c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 1
457c4762a1bSJed Brown 
458c4762a1bSJed Brown    test:
459c4762a1bSJed Brown       suffix: 2
460c4762a1bSJed Brown       nsize: 3
461c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2
462c4762a1bSJed Brown 
463c4762a1bSJed Brown    test:
464c4762a1bSJed Brown       suffix: 3
465c4762a1bSJed Brown       nsize: 3
466c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -ksp_monitor_short -problem_type 2 -snes_mf_operator -pack_dm_mat_type {{aij nest}} -pc_type fieldsplit -pc_fieldsplit_dm_splits -pc_fieldsplit_type additive -fieldsplit_u_ksp_type gmres -fieldsplit_k_pc_type jacobi
467c4762a1bSJed Brown 
468c4762a1bSJed Brown    test:
469c4762a1bSJed Brown       suffix: 4
470c4762a1bSJed Brown       nsize: 6
471c4762a1bSJed Brown       args: -u_da_grid_x 257 -snes_converged_reason -snes_monitor_short -ksp_monitor_short -problem_type 2 -snes_mf_operator -pack_dm_mat_type aij -pc_type fieldsplit -pc_fieldsplit_type multiplicative -fieldsplit_u_ksp_type gmres -fieldsplit_u_ksp_pc_side right -fieldsplit_u_pc_type mg -fieldsplit_u_pc_mg_levels 4 -fieldsplit_u_mg_levels_ksp_type richardson -fieldsplit_u_mg_levels_ksp_max_it 1 -fieldsplit_u_mg_levels_pc_type sor -fieldsplit_u_pc_mg_galerkin pmat -fieldsplit_u_ksp_converged_reason -fieldsplit_k_pc_type jacobi
472c4762a1bSJed Brown       requires: !single
473c4762a1bSJed Brown 
474c4762a1bSJed Brown    test:
475c4762a1bSJed Brown       suffix: glvis_composite_da_1d
476c4762a1bSJed Brown       args: -u_da_grid_x 20 -problem_type 0 -snes_monitor_solution glvis:
477c4762a1bSJed Brown 
47805ec3129SRichard Tran Mills    test:
47905ec3129SRichard Tran Mills       suffix: cuda
48005ec3129SRichard Tran Mills       nsize: 1
48105ec3129SRichard Tran Mills       requires: cuda
48205ec3129SRichard Tran Mills       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2 -pack_dm_mat_type aijcusparse -pack_dm_vec_type cuda
48305ec3129SRichard Tran Mills 
48405ec3129SRichard Tran Mills    test:
48505ec3129SRichard Tran Mills       suffix: viennacl
48605ec3129SRichard Tran Mills       nsize: 1
48705ec3129SRichard Tran Mills       requires: viennacl
48805ec3129SRichard Tran Mills       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2 -pack_dm_mat_type aijviennacl -pack_dm_vec_type viennacl
48905ec3129SRichard Tran Mills 
490c4762a1bSJed Brown TEST*/
491