xref: /petsc/src/snes/tutorials/ex28.c (revision fb3c7e2dd4f43a268b95c335e0bddf05b42b9816)
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 
38d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormFunctionLocal_U(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], PetscScalar f[])
39d71ae5a4SJacob Faibussowitsch {
40c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
41c4762a1bSJed Brown   PetscInt  i;
42c4762a1bSJed Brown 
43c4762a1bSJed Brown   PetscFunctionBeginUser;
44c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
45c4762a1bSJed Brown     if (i == 0) f[i] = 1. / hx * u[i];
46c4762a1bSJed Brown     else if (i == info->mx - 1) f[i] = 1. / hx * (u[i] - 1.0);
47c4762a1bSJed Brown     else f[i] = hx * ((k[i - 1] * (u[i] - u[i - 1]) - k[i] * (u[i + 1] - u[i])) / (hx * hx) - 1.0);
48c4762a1bSJed Brown   }
49c4762a1bSJed Brown   PetscFunctionReturn(0);
50c4762a1bSJed Brown }
51c4762a1bSJed Brown 
52d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormFunctionLocal_K(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], PetscScalar f[])
53d71ae5a4SJacob Faibussowitsch {
54c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
55c4762a1bSJed Brown   PetscInt  i;
56c4762a1bSJed Brown 
57c4762a1bSJed Brown   PetscFunctionBeginUser;
58c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
599371c9d4SSatish 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;
60c4762a1bSJed Brown     f[i] = hx * (PetscExpScalar(k[i] - 1.0) + k[i] - 1. / w);
61c4762a1bSJed Brown   }
62c4762a1bSJed Brown   PetscFunctionReturn(0);
63c4762a1bSJed Brown }
64c4762a1bSJed Brown 
65d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormFunction_All(SNES snes, Vec X, Vec F, void *ctx)
66d71ae5a4SJacob Faibussowitsch {
67c4762a1bSJed Brown   User          user = (User)ctx;
68c4762a1bSJed Brown   DM            dau, dak;
69c4762a1bSJed Brown   DMDALocalInfo infou, infok;
70c4762a1bSJed Brown   PetscScalar  *u, *k;
71c4762a1bSJed Brown   PetscScalar  *fu, *fk;
72c4762a1bSJed Brown   Vec           Uloc, Kloc, Fu, Fk;
73c4762a1bSJed Brown 
74c4762a1bSJed Brown   PetscFunctionBeginUser;
759566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
769566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
779566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
789566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack, &Uloc, &Kloc));
79c4762a1bSJed Brown   switch (user->ptype) {
80c4762a1bSJed Brown   case 0:
819566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau, X, INSERT_VALUES, Uloc));
829566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dau, X, INSERT_VALUES, Uloc));
839566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
849566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, user->Kloc, &k));
859566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, F, &fu));
869566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user, &infou, u, k, fu));
879566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, F, &fu));
889566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
899566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, user->Kloc, &k));
90c4762a1bSJed Brown     break;
91c4762a1bSJed Brown   case 1:
929566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak, X, INSERT_VALUES, Kloc));
939566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dak, X, INSERT_VALUES, Kloc));
949566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, user->Uloc, &u));
959566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
969566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, F, &fk));
979566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user, &infok, u, k, fk));
989566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, F, &fk));
999566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, user->Uloc, &u));
1009566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
101c4762a1bSJed Brown     break;
102c4762a1bSJed Brown   case 2:
1039566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack, X, Uloc, Kloc));
1049566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
1059566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
1069566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(user->pack, F, &Fu, &Fk));
1079566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Fu, &fu));
1089566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Fk, &fk));
1099566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user, &infou, u, k, fu));
1109566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user, &infok, u, k, fk));
1119566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Fu, &fu));
1129566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Fk, &fk));
1139566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(user->pack, F, &Fu, &Fk));
1149566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
1159566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
116c4762a1bSJed Brown     break;
117c4762a1bSJed Brown   }
1189566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack, &Uloc, &Kloc));
119c4762a1bSJed Brown   PetscFunctionReturn(0);
120c4762a1bSJed Brown }
121c4762a1bSJed Brown 
122d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormJacobianLocal_U(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], Mat Buu)
123d71ae5a4SJacob Faibussowitsch {
124c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
125c4762a1bSJed Brown   PetscInt  i;
126c4762a1bSJed Brown 
127c4762a1bSJed Brown   PetscFunctionBeginUser;
128c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
129c4762a1bSJed Brown     PetscInt    row = i - info->gxs, cols[] = {row - 1, row, row + 1};
130c4762a1bSJed Brown     PetscScalar val = 1. / hx;
131c4762a1bSJed Brown     if (i == 0) {
1329566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 1, &row, &val, INSERT_VALUES));
133c4762a1bSJed Brown     } else if (i == info->mx - 1) {
1349566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 1, &row, &val, INSERT_VALUES));
135c4762a1bSJed Brown     } else {
136c4762a1bSJed Brown       PetscScalar vals[] = {-k[i - 1] / hx, (k[i - 1] + k[i]) / hx, -k[i] / hx};
1379566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu, 1, &row, 3, cols, vals, INSERT_VALUES));
138c4762a1bSJed Brown     }
139c4762a1bSJed Brown   }
140c4762a1bSJed Brown   PetscFunctionReturn(0);
141c4762a1bSJed Brown }
142c4762a1bSJed Brown 
143d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormJacobianLocal_K(User user, DMDALocalInfo *info, const PetscScalar u[], const PetscScalar k[], Mat Bkk)
144d71ae5a4SJacob Faibussowitsch {
145c4762a1bSJed Brown   PetscReal hx = 1. / info->mx;
146c4762a1bSJed Brown   PetscInt  i;
147c4762a1bSJed Brown 
148c4762a1bSJed Brown   PetscFunctionBeginUser;
149c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
150c4762a1bSJed Brown     PetscInt    row    = i - info->gxs;
151c4762a1bSJed Brown     PetscScalar vals[] = {hx * (PetscExpScalar(k[i] - 1.) + (PetscScalar)1.)};
1529566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bkk, 1, &row, 1, &row, vals, INSERT_VALUES));
153c4762a1bSJed Brown   }
154c4762a1bSJed Brown   PetscFunctionReturn(0);
155c4762a1bSJed Brown }
156c4762a1bSJed Brown 
157d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormJacobianLocal_UK(User user, DMDALocalInfo *info, DMDALocalInfo *infok, const PetscScalar u[], const PetscScalar k[], Mat Buk)
158d71ae5a4SJacob Faibussowitsch {
159c4762a1bSJed Brown   PetscReal   hx = 1. / info->mx;
160c4762a1bSJed Brown   PetscInt    i;
161c4762a1bSJed Brown   PetscInt    row, cols[2];
162c4762a1bSJed Brown   PetscScalar vals[2];
163c4762a1bSJed Brown 
164c4762a1bSJed Brown   PetscFunctionBeginUser;
165c4762a1bSJed Brown   if (!Buk) PetscFunctionReturn(0); /* Not assembling this block */
166c4762a1bSJed Brown   for (i = info->xs; i < info->xs + info->xm; i++) {
167c4762a1bSJed Brown     if (i == 0 || i == info->mx - 1) continue;
168c4762a1bSJed Brown     row     = i - info->gxs;
1699371c9d4SSatish Balay     cols[0] = i - 1 - infok->gxs;
1709371c9d4SSatish Balay     vals[0] = (u[i] - u[i - 1]) / hx;
1719371c9d4SSatish Balay     cols[1] = i - infok->gxs;
1729371c9d4SSatish Balay     vals[1] = (u[i] - u[i + 1]) / hx;
1739566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Buk, 1, &row, 2, cols, vals, INSERT_VALUES));
174c4762a1bSJed Brown   }
175c4762a1bSJed Brown   PetscFunctionReturn(0);
176c4762a1bSJed Brown }
177c4762a1bSJed Brown 
178d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormJacobianLocal_KU(User user, DMDALocalInfo *info, DMDALocalInfo *infok, const PetscScalar u[], const PetscScalar k[], Mat Bku)
179d71ae5a4SJacob Faibussowitsch {
180c4762a1bSJed Brown   PetscInt  i;
181c4762a1bSJed Brown   PetscReal hx = 1. / (info->mx - 1);
182c4762a1bSJed Brown 
183c4762a1bSJed Brown   PetscFunctionBeginUser;
184c4762a1bSJed Brown   if (!Bku) PetscFunctionReturn(0); /* Not assembling this block */
185c4762a1bSJed Brown   for (i = infok->xs; i < infok->xs + infok->xm; i++) {
186c4762a1bSJed Brown     PetscInt    row = i - infok->gxs, cols[2];
187c4762a1bSJed Brown     PetscScalar vals[2];
1889371c9d4SSatish 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);
1899371c9d4SSatish Balay     cols[0] = i - info->gxs;
1909371c9d4SSatish Balay     vals[0] = -hx * (iw_ubar * ubar_L + iw_gradu * gradu_L);
1919371c9d4SSatish Balay     cols[1] = i + 1 - info->gxs;
1929371c9d4SSatish Balay     vals[1] = -hx * (iw_ubar * ubar_R + iw_gradu * gradu_R);
1939566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bku, 1, &row, 2, cols, vals, INSERT_VALUES));
194c4762a1bSJed Brown   }
195c4762a1bSJed Brown   PetscFunctionReturn(0);
196c4762a1bSJed Brown }
197c4762a1bSJed Brown 
198d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormJacobian_All(SNES snes, Vec X, Mat J, Mat B, void *ctx)
199d71ae5a4SJacob Faibussowitsch {
200c4762a1bSJed Brown   User          user = (User)ctx;
201c4762a1bSJed Brown   DM            dau, dak;
202c4762a1bSJed Brown   DMDALocalInfo infou, infok;
203c4762a1bSJed Brown   PetscScalar  *u, *k;
204c4762a1bSJed Brown   Vec           Uloc, Kloc;
205c4762a1bSJed Brown 
206c4762a1bSJed Brown   PetscFunctionBeginUser;
2079566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
2089566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
2099566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
2109566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack, &Uloc, &Kloc));
211c4762a1bSJed Brown   switch (user->ptype) {
212c4762a1bSJed Brown   case 0:
2139566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau, X, INSERT_VALUES, Uloc));
2149566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dau, X, INSERT_VALUES, Uloc));
2159566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
2169566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, user->Kloc, &k));
2179566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user, &infou, u, k, B));
2189566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
2199566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, user->Kloc, &k));
220c4762a1bSJed Brown     break;
221c4762a1bSJed Brown   case 1:
2229566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak, X, INSERT_VALUES, Kloc));
2239566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd(dak, X, INSERT_VALUES, Kloc));
2249566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, user->Uloc, &u));
2259566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
2269566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user, &infok, u, k, B));
2279566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, user->Uloc, &u));
2289566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
229c4762a1bSJed Brown     break;
230c4762a1bSJed Brown   case 2: {
231c4762a1bSJed Brown     Mat       Buu, Buk, Bku, Bkk;
232c4762a1bSJed Brown     PetscBool nest;
233c4762a1bSJed Brown     IS       *is;
2349566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack, X, Uloc, Kloc));
2359566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau, Uloc, &u));
2369566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak, Kloc, &k));
2379566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetLocalISs(user->pack, &is));
2389566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[0], is[0], &Buu));
2399566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[0], is[1], &Buk));
2409566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[1], is[0], &Bku));
2419566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B, is[1], is[1], &Bkk));
2429566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user, &infou, u, k, Buu));
2439566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)B, MATNEST, &nest));
244c4762a1bSJed Brown     if (!nest) {
245c4762a1bSJed Brown       /*
246c4762a1bSJed Brown          DMCreateMatrix_Composite()  for a nested matrix does not generate off-block matrices that one can call MatSetValuesLocal() on, it just creates dummy
247c4762a1bSJed Brown          matrices with no entries; there cannot insert values into them. Commit b6480e041dd2293a65f96222772d68cdb4ed6306
248c4762a1bSJed Brown          changed Mat_Nest() from returning NULL pointers for these submatrices to dummy matrices because PCFIELDSPLIT could not
249c4762a1bSJed Brown          handle the returned null matrices.
250c4762a1bSJed Brown       */
2519566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_UK(user, &infou, &infok, u, k, Buk));
2529566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_KU(user, &infou, &infok, u, k, Bku));
253c4762a1bSJed Brown     }
2549566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user, &infok, u, k, Bkk));
2559566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[0], is[0], &Buu));
2569566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[0], is[1], &Buk));
2579566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[1], is[0], &Bku));
2589566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B, is[1], is[1], &Bkk));
2599566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau, Uloc, &u));
2609566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak, Kloc, &k));
261c4762a1bSJed Brown 
2629566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[0]));
2639566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[1]));
2649566063dSJacob Faibussowitsch     PetscCall(PetscFree(is));
265c4762a1bSJed Brown   } break;
266c4762a1bSJed Brown   }
2679566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack, &Uloc, &Kloc));
2689566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2699566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
270c4762a1bSJed Brown   if (J != B) {
2719566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
2729566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
273c4762a1bSJed Brown   }
274c4762a1bSJed Brown   PetscFunctionReturn(0);
275c4762a1bSJed Brown }
276c4762a1bSJed Brown 
277d71ae5a4SJacob Faibussowitsch static PetscErrorCode FormInitial_Coupled(User user, Vec X)
278d71ae5a4SJacob Faibussowitsch {
279c4762a1bSJed Brown   DM            dau, dak;
280c4762a1bSJed Brown   DMDALocalInfo infou, infok;
281c4762a1bSJed Brown   Vec           Xu, Xk;
282c4762a1bSJed Brown   PetscScalar  *u, *k, hx;
283c4762a1bSJed Brown   PetscInt      i;
284c4762a1bSJed Brown 
285c4762a1bSJed Brown   PetscFunctionBeginUser;
2869566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack, &dau, &dak));
2879566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetAccess(user->pack, X, &Xu, &Xk));
2889566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dau, Xu, &u));
2899566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dak, Xk, &k));
2909566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau, &infou));
2919566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak, &infok));
292c4762a1bSJed Brown   hx = 1. / (infok.mx);
293c4762a1bSJed Brown   for (i = infou.xs; i < infou.xs + infou.xm; i++) u[i] = (PetscScalar)i * hx * (1. - (PetscScalar)i * hx);
294c4762a1bSJed Brown   for (i = infok.xs; i < infok.xs + infok.xm; i++) k[i] = 1.0 + 0.5 * PetscSinScalar(2 * PETSC_PI * i * hx);
2959566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dau, Xu, &u));
2969566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dak, Xk, &k));
2979566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreAccess(user->pack, X, &Xu, &Xk));
2989566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(user->pack, X, user->Uloc, user->Kloc));
299c4762a1bSJed Brown   PetscFunctionReturn(0);
300c4762a1bSJed Brown }
301c4762a1bSJed Brown 
302d71ae5a4SJacob Faibussowitsch int main(int argc, char *argv[])
303d71ae5a4SJacob Faibussowitsch {
304c4762a1bSJed Brown   DM              dau, dak, pack;
305c4762a1bSJed Brown   const PetscInt *lxu;
306c4762a1bSJed Brown   PetscInt       *lxk, m, sizes;
307c4762a1bSJed Brown   User            user;
308c4762a1bSJed Brown   SNES            snes;
309c4762a1bSJed Brown   Vec             X, F, Xu, Xk, Fu, Fk;
310c4762a1bSJed Brown   Mat             B;
311c4762a1bSJed Brown   IS             *isg;
312c4762a1bSJed Brown   PetscBool       pass_dm;
313c4762a1bSJed Brown 
314327415f7SBarry Smith   PetscFunctionBeginUser;
3159566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc, &argv, 0, help));
3169566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 10, 1, 1, NULL, &dau));
3179566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dau, "u_"));
3189566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dau));
3199566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dau));
3209566063dSJacob Faibussowitsch   PetscCall(DMDAGetOwnershipRanges(dau, &lxu, 0, 0));
3219566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(dau, 0, &m, 0, 0, &sizes, 0, 0, 0, 0, 0, 0, 0, 0));
3229566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(sizes, &lxk));
3239566063dSJacob Faibussowitsch   PetscCall(PetscArraycpy(lxk, lxu, sizes));
324c4762a1bSJed Brown   lxk[0]--;
3259566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, m - 1, 1, 1, lxk, &dak));
3269566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dak, "k_"));
3279566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dak));
3289566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dak));
3299566063dSJacob Faibussowitsch   PetscCall(PetscFree(lxk));
330c4762a1bSJed Brown 
3319566063dSJacob Faibussowitsch   PetscCall(DMCompositeCreate(PETSC_COMM_WORLD, &pack));
3329566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(pack, "pack_"));
3339566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack, dau));
3349566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack, dak));
3359566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dau, 0, "u"));
3369566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dak, 0, "k"));
3379566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(pack));
338c4762a1bSJed Brown 
3399566063dSJacob Faibussowitsch   PetscCall(DMCreateGlobalVector(pack, &X));
3409566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(X, &F));
341c4762a1bSJed Brown 
3429566063dSJacob Faibussowitsch   PetscCall(PetscNew(&user));
343c4762a1bSJed Brown 
344c4762a1bSJed Brown   user->pack = pack;
345c4762a1bSJed Brown 
3469566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetGlobalISs(pack, &isg));
3479566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(pack, &user->Uloc, &user->Kloc));
3489566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(pack, X, user->Uloc, user->Kloc));
349c4762a1bSJed Brown 
350d0609cedSBarry Smith   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Coupled problem options", "SNES");
351c4762a1bSJed Brown   {
3529371c9d4SSatish Balay     user->ptype = 0;
3539371c9d4SSatish Balay     pass_dm     = PETSC_TRUE;
354c4762a1bSJed Brown 
3559566063dSJacob 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));
3569566063dSJacob 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));
357c4762a1bSJed Brown   }
358d0609cedSBarry Smith   PetscOptionsEnd();
359c4762a1bSJed Brown 
3609566063dSJacob Faibussowitsch   PetscCall(FormInitial_Coupled(user, X));
361c4762a1bSJed Brown 
3629566063dSJacob Faibussowitsch   PetscCall(SNESCreate(PETSC_COMM_WORLD, &snes));
363c4762a1bSJed Brown   switch (user->ptype) {
364c4762a1bSJed Brown   case 0:
3659566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, X, &Xu, 0));
3669566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, F, &Fu, 0));
3679566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dau, &B));
3689566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, Fu, FormFunction_All, user));
3699566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3709566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3719566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes, dau));
3729566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Xu));
3739566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, X, &Xu, 0));
3749566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, F, &Fu, 0));
375c4762a1bSJed Brown     break;
376c4762a1bSJed Brown   case 1:
3779566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, X, 0, &Xk));
3789566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack, F, 0, &Fk));
3799566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dak, &B));
3809566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, Fk, FormFunction_All, user));
3819566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3829566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3839566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes, dak));
3849566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Xk));
3859566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, X, 0, &Xk));
3869566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack, F, 0, &Fk));
387c4762a1bSJed Brown     break;
388c4762a1bSJed Brown   case 2:
3899566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(pack, &B));
390c4762a1bSJed Brown     /* This example does not correctly allocate off-diagonal blocks. These options allows new nonzeros (slow). */
3919566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_FALSE));
3929566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE));
3939566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes, F, FormFunction_All, user));
3949566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes, B, B, FormJacobian_All, user));
3959566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
396c4762a1bSJed Brown     if (!pass_dm) { /* Manually provide index sets and names for the splits */
397c4762a1bSJed Brown       KSP ksp;
398c4762a1bSJed Brown       PC  pc;
3999566063dSJacob Faibussowitsch       PetscCall(SNESGetKSP(snes, &ksp));
4009566063dSJacob Faibussowitsch       PetscCall(KSPGetPC(ksp, &pc));
4019566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc, "u", isg[0]));
4029566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc, "k", isg[1]));
403c4762a1bSJed Brown     } else {
404c4762a1bSJed Brown       /* The same names come from the options prefix for dau and dak. This option can support geometric multigrid inside
405c4762a1bSJed Brown        * of splits, but it requires using a DM (perhaps your own implementation). */
4069566063dSJacob Faibussowitsch       PetscCall(SNESSetDM(snes, pack));
407c4762a1bSJed Brown     }
4089566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, X));
409c4762a1bSJed Brown     break;
410c4762a1bSJed Brown   }
4119566063dSJacob Faibussowitsch   PetscCall(VecViewFromOptions(X, NULL, "-view_sol"));
412c4762a1bSJed Brown 
413c4762a1bSJed Brown   if (0) {
414c4762a1bSJed Brown     PetscInt  col      = 0;
415c4762a1bSJed Brown     PetscBool mult_dup = PETSC_FALSE, view_dup = PETSC_FALSE;
416c4762a1bSJed Brown     Mat       D;
417c4762a1bSJed Brown     Vec       Y;
418c4762a1bSJed Brown 
4199566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetInt(NULL, 0, "-col", &col, 0));
4209566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL, 0, "-mult_dup", &mult_dup, 0));
4219566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL, 0, "-view_dup", &view_dup, 0));
422c4762a1bSJed Brown 
4239566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X, &Y));
4249566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); */
4259566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); */
4269566063dSJacob Faibussowitsch     PetscCall(MatConvert(B, MATAIJ, MAT_INITIAL_MATRIX, &D));
4279566063dSJacob Faibussowitsch     PetscCall(VecZeroEntries(X));
4289566063dSJacob Faibussowitsch     PetscCall(VecSetValue(X, col, 1.0, INSERT_VALUES));
4299566063dSJacob Faibussowitsch     PetscCall(VecAssemblyBegin(X));
4309566063dSJacob Faibussowitsch     PetscCall(VecAssemblyEnd(X));
4319566063dSJacob Faibussowitsch     PetscCall(MatMult(mult_dup ? D : B, X, Y));
4329566063dSJacob Faibussowitsch     PetscCall(MatView(view_dup ? D : B, PETSC_VIEWER_STDOUT_WORLD));
4339566063dSJacob Faibussowitsch     /* PetscCall(VecView(X,PETSC_VIEWER_STDOUT_WORLD)); */
4349566063dSJacob Faibussowitsch     PetscCall(VecView(Y, PETSC_VIEWER_STDOUT_WORLD));
4359566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&D));
4369566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&Y));
437c4762a1bSJed Brown   }
438c4762a1bSJed Brown 
4399566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(pack, &user->Uloc, &user->Kloc));
4409566063dSJacob Faibussowitsch   PetscCall(PetscFree(user));
441c4762a1bSJed Brown 
4429566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[0]));
4439566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[1]));
4449566063dSJacob Faibussowitsch   PetscCall(PetscFree(isg));
4459566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X));
4469566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&F));
4479566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
4489566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dau));
4499566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dak));
4509566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&pack));
4519566063dSJacob Faibussowitsch   PetscCall(SNESDestroy(&snes));
4529566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
453b122ec5aSJacob Faibussowitsch   return 0;
454c4762a1bSJed Brown }
455c4762a1bSJed Brown 
456c4762a1bSJed Brown /*TEST
457c4762a1bSJed Brown 
458c4762a1bSJed Brown    test:
459c4762a1bSJed Brown       suffix: 0
460c4762a1bSJed Brown       nsize: 3
461c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 0
462c4762a1bSJed Brown 
463c4762a1bSJed Brown    test:
464c4762a1bSJed Brown       suffix: 1
465c4762a1bSJed Brown       nsize: 3
466c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 1
467c4762a1bSJed Brown 
468c4762a1bSJed Brown    test:
469c4762a1bSJed Brown       suffix: 2
470c4762a1bSJed Brown       nsize: 3
471c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2
472c4762a1bSJed Brown 
473*fb3c7e2dSJunchao Zhang    testset:
474c4762a1bSJed Brown       suffix: 3
475c4762a1bSJed Brown       nsize: 3
476*fb3c7e2dSJunchao Zhang       output_file: output/ex28_3.out
477*fb3c7e2dSJunchao Zhang       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -ksp_monitor_short -problem_type 2 -snes_mf_operator  -pc_type fieldsplit -pc_fieldsplit_dm_splits -pc_fieldsplit_type additive -fieldsplit_u_ksp_type gmres -fieldsplit_k_pc_type jacobi
478*fb3c7e2dSJunchao Zhang 
479*fb3c7e2dSJunchao Zhang       test:
480*fb3c7e2dSJunchao Zhang         suffix: std
481*fb3c7e2dSJunchao Zhang         args: -pack_dm_mat_type {{aij nest}}
482*fb3c7e2dSJunchao Zhang 
483*fb3c7e2dSJunchao Zhang       test:
484*fb3c7e2dSJunchao Zhang         suffix: cuda
485*fb3c7e2dSJunchao Zhang         requires: cuda
486*fb3c7e2dSJunchao Zhang         args: -pack_dm_mat_type aijcusparse -pack_dm_vec_type cuda
487*fb3c7e2dSJunchao Zhang 
488*fb3c7e2dSJunchao Zhang       test:
489*fb3c7e2dSJunchao Zhang         suffix: hip
490*fb3c7e2dSJunchao Zhang         requires: hip
491*fb3c7e2dSJunchao Zhang         args: -pack_dm_mat_type aijhipsparse -pack_dm_vec_type hip
492*fb3c7e2dSJunchao Zhang 
493*fb3c7e2dSJunchao Zhang       test:
494*fb3c7e2dSJunchao Zhang         suffix: kok
495*fb3c7e2dSJunchao Zhang         requires: kokkos_kernels
496*fb3c7e2dSJunchao Zhang         args: -pack_dm_mat_type aijkokkos -pack_dm_vec_type kokkos
497c4762a1bSJed Brown 
498c4762a1bSJed Brown    test:
499c4762a1bSJed Brown       suffix: 4
500c4762a1bSJed Brown       nsize: 6
501c4762a1bSJed 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
502c4762a1bSJed Brown       requires: !single
503c4762a1bSJed Brown 
504c4762a1bSJed Brown    test:
505c4762a1bSJed Brown       suffix: glvis_composite_da_1d
506c4762a1bSJed Brown       args: -u_da_grid_x 20 -problem_type 0 -snes_monitor_solution glvis:
507c4762a1bSJed Brown 
50805ec3129SRichard Tran Mills    test:
50905ec3129SRichard Tran Mills       suffix: cuda
51005ec3129SRichard Tran Mills       nsize: 1
51105ec3129SRichard Tran Mills       requires: cuda
51205ec3129SRichard 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
51305ec3129SRichard Tran Mills 
51405ec3129SRichard Tran Mills    test:
51505ec3129SRichard Tran Mills       suffix: viennacl
51605ec3129SRichard Tran Mills       nsize: 1
51705ec3129SRichard Tran Mills       requires: viennacl
51805ec3129SRichard 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
51905ec3129SRichard Tran Mills 
520c4762a1bSJed Brown TEST*/
521