xref: /petsc/src/snes/tutorials/ex28.c (revision 9566063d113dddea24716c546802770db7481bc0)
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 
38c4762a1bSJed Brown static PetscErrorCode FormFunctionLocal_U(User user,DMDALocalInfo *info,const PetscScalar u[],const PetscScalar k[],PetscScalar f[])
39c4762a1bSJed Brown {
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 
52c4762a1bSJed Brown static PetscErrorCode FormFunctionLocal_K(User user,DMDALocalInfo *info,const PetscScalar u[],const PetscScalar k[],PetscScalar f[])
53c4762a1bSJed Brown {
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++) {
59c4762a1bSJed Brown     const PetscScalar
60c4762a1bSJed Brown       ubar  = 0.5*(u[i+1]+u[i]),
61c4762a1bSJed Brown       gradu = (u[i+1]-u[i])/hx,
62c4762a1bSJed Brown       g     = 1. + gradu*gradu,
63c4762a1bSJed Brown       w     = 1./(1.+ubar) + 1./g;
64c4762a1bSJed Brown     f[i] = hx*(PetscExpScalar(k[i]-1.0) + k[i] - 1./w);
65c4762a1bSJed Brown   }
66c4762a1bSJed Brown   PetscFunctionReturn(0);
67c4762a1bSJed Brown }
68c4762a1bSJed Brown 
69c4762a1bSJed Brown static PetscErrorCode FormFunction_All(SNES snes,Vec X,Vec F,void *ctx)
70c4762a1bSJed Brown {
71c4762a1bSJed Brown   User           user = (User)ctx;
72c4762a1bSJed Brown   DM             dau,dak;
73c4762a1bSJed Brown   DMDALocalInfo  infou,infok;
74c4762a1bSJed Brown   PetscScalar    *u,*k;
75c4762a1bSJed Brown   PetscScalar    *fu,*fk;
76c4762a1bSJed Brown   Vec            Uloc,Kloc,Fu,Fk;
77c4762a1bSJed Brown 
78c4762a1bSJed Brown   PetscFunctionBeginUser;
79*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
80*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
81*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak,&infok));
82*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
83c4762a1bSJed Brown   switch (user->ptype) {
84c4762a1bSJed Brown   case 0:
85*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
86*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
87*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
88*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,user->Kloc,&k));
89*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,F,&fu));
90*9566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user,&infou,u,k,fu));
91*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,F,&fu));
92*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
93*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,user->Kloc,&k));
94c4762a1bSJed Brown     break;
95c4762a1bSJed Brown   case 1:
96*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
97*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
98*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,user->Uloc,&u));
99*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
100*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,F,&fk));
101*9566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user,&infok,u,k,fk));
102*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,F,&fk));
103*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,user->Uloc,&u));
104*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
105c4762a1bSJed Brown     break;
106c4762a1bSJed Brown   case 2:
107*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack,X,Uloc,Kloc));
108*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
109*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
110*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(user->pack,F,&Fu,&Fk));
111*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Fu,&fu));
112*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Fk,&fk));
113*9566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user,&infou,u,k,fu));
114*9566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user,&infok,u,k,fk));
115*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Fu,&fu));
116*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Fk,&fk));
117*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(user->pack,F,&Fu,&Fk));
118*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
119*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
120c4762a1bSJed Brown     break;
121c4762a1bSJed Brown   }
122*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack,&Uloc,&Kloc));
123c4762a1bSJed Brown   PetscFunctionReturn(0);
124c4762a1bSJed Brown }
125c4762a1bSJed Brown 
126c4762a1bSJed Brown static PetscErrorCode FormJacobianLocal_U(User user,DMDALocalInfo *info,const PetscScalar u[],const PetscScalar k[],Mat Buu)
127c4762a1bSJed Brown {
128c4762a1bSJed Brown   PetscReal      hx = 1./info->mx;
129c4762a1bSJed Brown   PetscInt       i;
130c4762a1bSJed Brown 
131c4762a1bSJed Brown   PetscFunctionBeginUser;
132c4762a1bSJed Brown   for (i=info->xs; i<info->xs+info->xm; i++) {
133c4762a1bSJed Brown     PetscInt    row = i-info->gxs,cols[] = {row-1,row,row+1};
134c4762a1bSJed Brown     PetscScalar val = 1./hx;
135c4762a1bSJed Brown     if (i == 0) {
136*9566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu,1,&row,1,&row,&val,INSERT_VALUES));
137c4762a1bSJed Brown     } else if (i == info->mx-1) {
138*9566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu,1,&row,1,&row,&val,INSERT_VALUES));
139c4762a1bSJed Brown     } else {
140c4762a1bSJed Brown       PetscScalar vals[] = {-k[i-1]/hx,(k[i-1]+k[i])/hx,-k[i]/hx};
141*9566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu,1,&row,3,cols,vals,INSERT_VALUES));
142c4762a1bSJed Brown     }
143c4762a1bSJed Brown   }
144c4762a1bSJed Brown   PetscFunctionReturn(0);
145c4762a1bSJed Brown }
146c4762a1bSJed Brown 
147c4762a1bSJed Brown static PetscErrorCode FormJacobianLocal_K(User user,DMDALocalInfo *info,const PetscScalar u[],const PetscScalar k[],Mat Bkk)
148c4762a1bSJed Brown {
149c4762a1bSJed Brown   PetscReal      hx = 1./info->mx;
150c4762a1bSJed Brown   PetscInt       i;
151c4762a1bSJed Brown 
152c4762a1bSJed Brown   PetscFunctionBeginUser;
153c4762a1bSJed Brown   for (i=info->xs; i<info->xs+info->xm; i++) {
154c4762a1bSJed Brown     PetscInt    row    = i-info->gxs;
155c4762a1bSJed Brown     PetscScalar vals[] = {hx*(PetscExpScalar(k[i]-1.)+ (PetscScalar)1.)};
156*9566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bkk,1,&row,1,&row,vals,INSERT_VALUES));
157c4762a1bSJed Brown   }
158c4762a1bSJed Brown   PetscFunctionReturn(0);
159c4762a1bSJed Brown }
160c4762a1bSJed Brown 
161c4762a1bSJed Brown static PetscErrorCode FormJacobianLocal_UK(User user,DMDALocalInfo *info,DMDALocalInfo *infok,const PetscScalar u[],const PetscScalar k[],Mat Buk)
162c4762a1bSJed Brown {
163c4762a1bSJed Brown   PetscReal      hx = 1./info->mx;
164c4762a1bSJed Brown   PetscInt       i;
165c4762a1bSJed Brown   PetscInt       row,cols[2];
166c4762a1bSJed Brown   PetscScalar    vals[2];
167c4762a1bSJed Brown 
168c4762a1bSJed Brown   PetscFunctionBeginUser;
169c4762a1bSJed Brown   if (!Buk) PetscFunctionReturn(0); /* Not assembling this block */
170c4762a1bSJed Brown   for (i=info->xs; i<info->xs+info->xm; i++) {
171c4762a1bSJed Brown     if (i == 0 || i == info->mx-1) continue;
172c4762a1bSJed Brown     row     = i-info->gxs;
173c4762a1bSJed Brown     cols[0] = i-1-infok->gxs;  vals[0] = (u[i]-u[i-1])/hx;
174c4762a1bSJed Brown     cols[1] = i-infok->gxs;    vals[1] = (u[i]-u[i+1])/hx;
175*9566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Buk,1,&row,2,cols,vals,INSERT_VALUES));
176c4762a1bSJed Brown   }
177c4762a1bSJed Brown   PetscFunctionReturn(0);
178c4762a1bSJed Brown }
179c4762a1bSJed Brown 
180c4762a1bSJed Brown static PetscErrorCode FormJacobianLocal_KU(User user,DMDALocalInfo *info,DMDALocalInfo *infok,const PetscScalar u[],const PetscScalar k[],Mat Bku)
181c4762a1bSJed Brown {
182c4762a1bSJed Brown   PetscInt       i;
183c4762a1bSJed Brown   PetscReal      hx = 1./(info->mx-1);
184c4762a1bSJed Brown 
185c4762a1bSJed Brown   PetscFunctionBeginUser;
186c4762a1bSJed Brown   if (!Bku) PetscFunctionReturn(0); /* Not assembling this block */
187c4762a1bSJed Brown   for (i=infok->xs; i<infok->xs+infok->xm; i++) {
188c4762a1bSJed Brown     PetscInt    row = i-infok->gxs,cols[2];
189c4762a1bSJed Brown     PetscScalar vals[2];
190c4762a1bSJed Brown     const PetscScalar
191c4762a1bSJed Brown       ubar     = 0.5*(u[i]+u[i+1]),
192c4762a1bSJed Brown       ubar_L   = 0.5,
193c4762a1bSJed Brown       ubar_R   = 0.5,
194c4762a1bSJed Brown       gradu    = (u[i+1]-u[i])/hx,
195c4762a1bSJed Brown       gradu_L  = -1./hx,
196c4762a1bSJed Brown       gradu_R  = 1./hx,
197c4762a1bSJed Brown       g        = 1. + PetscSqr(gradu),
198c4762a1bSJed Brown       g_gradu  = 2.*gradu,
199c4762a1bSJed Brown       w        = 1./(1.+ubar) + 1./g,
200c4762a1bSJed Brown       w_ubar   = -1./PetscSqr(1.+ubar),
201c4762a1bSJed Brown       w_gradu  = -g_gradu/PetscSqr(g),
202c4762a1bSJed Brown       iw       = 1./w,
203c4762a1bSJed Brown       iw_ubar  = -w_ubar * PetscSqr(iw),
204c4762a1bSJed Brown       iw_gradu = -w_gradu * PetscSqr(iw);
205c4762a1bSJed Brown     cols[0] = i-info->gxs;         vals[0] = -hx*(iw_ubar*ubar_L + iw_gradu*gradu_L);
206c4762a1bSJed Brown     cols[1] = i+1-info->gxs;       vals[1] = -hx*(iw_ubar*ubar_R + iw_gradu*gradu_R);
207*9566063dSJacob Faibussowitsch     PetscCall(MatSetValuesLocal(Bku,1,&row,2,cols,vals,INSERT_VALUES));
208c4762a1bSJed Brown   }
209c4762a1bSJed Brown   PetscFunctionReturn(0);
210c4762a1bSJed Brown }
211c4762a1bSJed Brown 
212c4762a1bSJed Brown static PetscErrorCode FormJacobian_All(SNES snes,Vec X,Mat J,Mat B,void *ctx)
213c4762a1bSJed Brown {
214c4762a1bSJed Brown   User           user = (User)ctx;
215c4762a1bSJed Brown   DM             dau,dak;
216c4762a1bSJed Brown   DMDALocalInfo  infou,infok;
217c4762a1bSJed Brown   PetscScalar    *u,*k;
218c4762a1bSJed Brown   Vec            Uloc,Kloc;
219c4762a1bSJed Brown 
220c4762a1bSJed Brown   PetscFunctionBeginUser;
221*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
222*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
223*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak,&infok));
224*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
225c4762a1bSJed Brown   switch (user->ptype) {
226c4762a1bSJed Brown   case 0:
227*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
228*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
229*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
230*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,user->Kloc,&k));
231*9566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user,&infou,u,k,B));
232*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
233*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,user->Kloc,&k));
234c4762a1bSJed Brown     break;
235c4762a1bSJed Brown   case 1:
236*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
237*9566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
238*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,user->Uloc,&u));
239*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
240*9566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user,&infok,u,k,B));
241*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,user->Uloc,&u));
242*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
243c4762a1bSJed Brown     break;
244c4762a1bSJed Brown   case 2: {
245c4762a1bSJed Brown     Mat       Buu,Buk,Bku,Bkk;
246c4762a1bSJed Brown     PetscBool nest;
247c4762a1bSJed Brown     IS        *is;
248*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack,X,Uloc,Kloc));
249*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
250*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
251*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetLocalISs(user->pack,&is));
252*9566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[0],is[0],&Buu));
253*9566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[0],is[1],&Buk));
254*9566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[1],is[0],&Bku));
255*9566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[1],is[1],&Bkk));
256*9566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user,&infou,u,k,Buu));
257*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)B,MATNEST,&nest));
258c4762a1bSJed Brown     if (!nest) {
259c4762a1bSJed Brown       /*
260c4762a1bSJed Brown          DMCreateMatrix_Composite()  for a nested matrix does not generate off-block matrices that one can call MatSetValuesLocal() on, it just creates dummy
261c4762a1bSJed Brown          matrices with no entries; there cannot insert values into them. Commit b6480e041dd2293a65f96222772d68cdb4ed6306
262c4762a1bSJed Brown          changed Mat_Nest() from returning NULL pointers for these submatrices to dummy matrices because PCFIELDSPLIT could not
263c4762a1bSJed Brown          handle the returned null matrices.
264c4762a1bSJed Brown       */
265*9566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_UK(user,&infou,&infok,u,k,Buk));
266*9566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_KU(user,&infou,&infok,u,k,Bku));
267c4762a1bSJed Brown     }
268*9566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user,&infok,u,k,Bkk));
269*9566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[0],is[0],&Buu));
270*9566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[0],is[1],&Buk));
271*9566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[1],is[0],&Bku));
272*9566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[1],is[1],&Bkk));
273*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
274*9566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
275c4762a1bSJed Brown 
276*9566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[0]));
277*9566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[1]));
278*9566063dSJacob Faibussowitsch     PetscCall(PetscFree(is));
279c4762a1bSJed Brown   } break;
280c4762a1bSJed Brown   }
281*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack,&Uloc,&Kloc));
282*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
283*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd  (B,MAT_FINAL_ASSEMBLY));
284c4762a1bSJed Brown   if (J != B) {
285*9566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
286*9566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd  (J,MAT_FINAL_ASSEMBLY));
287c4762a1bSJed Brown   }
288c4762a1bSJed Brown   PetscFunctionReturn(0);
289c4762a1bSJed Brown }
290c4762a1bSJed Brown 
291c4762a1bSJed Brown static PetscErrorCode FormInitial_Coupled(User user,Vec X)
292c4762a1bSJed Brown {
293c4762a1bSJed Brown   DM             dau,dak;
294c4762a1bSJed Brown   DMDALocalInfo  infou,infok;
295c4762a1bSJed Brown   Vec            Xu,Xk;
296c4762a1bSJed Brown   PetscScalar    *u,*k,hx;
297c4762a1bSJed Brown   PetscInt       i;
298c4762a1bSJed Brown 
299c4762a1bSJed Brown   PetscFunctionBeginUser;
300*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
301*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetAccess(user->pack,X,&Xu,&Xk));
302*9566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dau,Xu,&u));
303*9566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dak,Xk,&k));
304*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
305*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak,&infok));
306c4762a1bSJed Brown   hx   = 1./(infok.mx);
307c4762a1bSJed Brown   for (i=infou.xs; i<infou.xs+infou.xm; i++) u[i] = (PetscScalar)i*hx * (1.-(PetscScalar)i*hx);
308c4762a1bSJed Brown   for (i=infok.xs; i<infok.xs+infok.xm; i++) k[i] = 1.0 + 0.5*PetscSinScalar(2*PETSC_PI*i*hx);
309*9566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dau,Xu,&u));
310*9566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dak,Xk,&k));
311*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreAccess(user->pack,X,&Xu,&Xk));
312*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(user->pack,X,user->Uloc,user->Kloc));
313c4762a1bSJed Brown   PetscFunctionReturn(0);
314c4762a1bSJed Brown }
315c4762a1bSJed Brown 
316c4762a1bSJed Brown int main(int argc, char *argv[])
317c4762a1bSJed Brown {
318c4762a1bSJed Brown   PetscErrorCode ierr;
319c4762a1bSJed Brown   DM             dau,dak,pack;
320c4762a1bSJed Brown   const PetscInt *lxu;
321c4762a1bSJed Brown   PetscInt       *lxk,m,sizes;
322c4762a1bSJed Brown   User           user;
323c4762a1bSJed Brown   SNES           snes;
324c4762a1bSJed Brown   Vec            X,F,Xu,Xk,Fu,Fk;
325c4762a1bSJed Brown   Mat            B;
326c4762a1bSJed Brown   IS             *isg;
327c4762a1bSJed Brown   PetscBool      pass_dm;
328c4762a1bSJed Brown 
329*9566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc,&argv,0,help));
330*9566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,10,1,1,NULL,&dau));
331*9566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dau,"u_"));
332*9566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dau));
333*9566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dau));
334*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetOwnershipRanges(dau,&lxu,0,0));
335*9566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(dau,0, &m,0,0, &sizes,0,0, 0,0,0,0,0,0));
336*9566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(sizes,&lxk));
337*9566063dSJacob Faibussowitsch   PetscCall(PetscArraycpy(lxk,lxu,sizes));
338c4762a1bSJed Brown   lxk[0]--;
339*9566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,m-1,1,1,lxk,&dak));
340*9566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dak,"k_"));
341*9566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dak));
342*9566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dak));
343*9566063dSJacob Faibussowitsch   PetscCall(PetscFree(lxk));
344c4762a1bSJed Brown 
345*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeCreate(PETSC_COMM_WORLD,&pack));
346*9566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(pack,"pack_"));
347*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack,dau));
348*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack,dak));
349*9566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dau,0,"u"));
350*9566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dak,0,"k"));
351*9566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(pack));
352c4762a1bSJed Brown 
353*9566063dSJacob Faibussowitsch   PetscCall(DMCreateGlobalVector(pack,&X));
354*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(X,&F));
355c4762a1bSJed Brown 
356*9566063dSJacob Faibussowitsch   PetscCall(PetscNew(&user));
357c4762a1bSJed Brown 
358c4762a1bSJed Brown   user->pack = pack;
359c4762a1bSJed Brown 
360*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetGlobalISs(pack,&isg));
361*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(pack,&user->Uloc,&user->Kloc));
362*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(pack,X,user->Uloc,user->Kloc));
363c4762a1bSJed Brown 
364*9566063dSJacob Faibussowitsch   ierr = PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"Coupled problem options","SNES");PetscCall(ierr);
365c4762a1bSJed Brown   {
366c4762a1bSJed Brown     user->ptype = 0; pass_dm = PETSC_TRUE;
367c4762a1bSJed Brown 
368*9566063dSJacob 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));
369*9566063dSJacob 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));
370c4762a1bSJed Brown   }
371*9566063dSJacob Faibussowitsch   ierr = PetscOptionsEnd();PetscCall(ierr);
372c4762a1bSJed Brown 
373*9566063dSJacob Faibussowitsch   PetscCall(FormInitial_Coupled(user,X));
374c4762a1bSJed Brown 
375*9566063dSJacob Faibussowitsch   PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes));
376c4762a1bSJed Brown   switch (user->ptype) {
377c4762a1bSJed Brown   case 0:
378*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,X,&Xu,0));
379*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,F,&Fu,0));
380*9566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dau,&B));
381*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,Fu,FormFunction_All,user));
382*9566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
383*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
384*9566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes,dau));
385*9566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,Xu));
386*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,X,&Xu,0));
387*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,F,&Fu,0));
388c4762a1bSJed Brown     break;
389c4762a1bSJed Brown   case 1:
390*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,X,0,&Xk));
391*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,F,0,&Fk));
392*9566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dak,&B));
393*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,Fk,FormFunction_All,user));
394*9566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
395*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
396*9566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes,dak));
397*9566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,Xk));
398*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,X,0,&Xk));
399*9566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,F,0,&Fk));
400c4762a1bSJed Brown     break;
401c4762a1bSJed Brown   case 2:
402*9566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(pack,&B));
403c4762a1bSJed Brown     /* This example does not correctly allocate off-diagonal blocks. These options allows new nonzeros (slow). */
404*9566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_FALSE));
405*9566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_FALSE));
406*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,F,FormFunction_All,user));
407*9566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
408*9566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
409c4762a1bSJed Brown     if (!pass_dm) {             /* Manually provide index sets and names for the splits */
410c4762a1bSJed Brown       KSP ksp;
411c4762a1bSJed Brown       PC  pc;
412*9566063dSJacob Faibussowitsch       PetscCall(SNESGetKSP(snes,&ksp));
413*9566063dSJacob Faibussowitsch       PetscCall(KSPGetPC(ksp,&pc));
414*9566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc,"u",isg[0]));
415*9566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc,"k",isg[1]));
416c4762a1bSJed Brown     } else {
417c4762a1bSJed Brown       /* The same names come from the options prefix for dau and dak. This option can support geometric multigrid inside
418c4762a1bSJed Brown        * of splits, but it requires using a DM (perhaps your own implementation). */
419*9566063dSJacob Faibussowitsch       PetscCall(SNESSetDM(snes,pack));
420c4762a1bSJed Brown     }
421*9566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,X));
422c4762a1bSJed Brown     break;
423c4762a1bSJed Brown   }
424*9566063dSJacob Faibussowitsch   PetscCall(VecViewFromOptions(X,NULL,"-view_sol"));
425c4762a1bSJed Brown 
426c4762a1bSJed Brown   if (0) {
427c4762a1bSJed Brown     PetscInt  col      = 0;
428c4762a1bSJed Brown     PetscBool mult_dup = PETSC_FALSE,view_dup = PETSC_FALSE;
429c4762a1bSJed Brown     Mat       D;
430c4762a1bSJed Brown     Vec       Y;
431c4762a1bSJed Brown 
432*9566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetInt(NULL,0,"-col",&col,0));
433*9566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL,0,"-mult_dup",&mult_dup,0));
434*9566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL,0,"-view_dup",&view_dup,0));
435c4762a1bSJed Brown 
436*9566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X,&Y));
437*9566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); */
438*9566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); */
439*9566063dSJacob Faibussowitsch     PetscCall(MatConvert(B,MATAIJ,MAT_INITIAL_MATRIX,&D));
440*9566063dSJacob Faibussowitsch     PetscCall(VecZeroEntries(X));
441*9566063dSJacob Faibussowitsch     PetscCall(VecSetValue(X,col,1.0,INSERT_VALUES));
442*9566063dSJacob Faibussowitsch     PetscCall(VecAssemblyBegin(X));
443*9566063dSJacob Faibussowitsch     PetscCall(VecAssemblyEnd(X));
444*9566063dSJacob Faibussowitsch     PetscCall(MatMult(mult_dup ? D : B,X,Y));
445*9566063dSJacob Faibussowitsch     PetscCall(MatView(view_dup ? D : B,PETSC_VIEWER_STDOUT_WORLD));
446*9566063dSJacob Faibussowitsch     /* PetscCall(VecView(X,PETSC_VIEWER_STDOUT_WORLD)); */
447*9566063dSJacob Faibussowitsch     PetscCall(VecView(Y,PETSC_VIEWER_STDOUT_WORLD));
448*9566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&D));
449*9566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&Y));
450c4762a1bSJed Brown   }
451c4762a1bSJed Brown 
452*9566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(pack,&user->Uloc,&user->Kloc));
453*9566063dSJacob Faibussowitsch   PetscCall(PetscFree(user));
454c4762a1bSJed Brown 
455*9566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[0]));
456*9566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[1]));
457*9566063dSJacob Faibussowitsch   PetscCall(PetscFree(isg));
458*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X));
459*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&F));
460*9566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
461*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dau));
462*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dak));
463*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&pack));
464*9566063dSJacob Faibussowitsch   PetscCall(SNESDestroy(&snes));
465*9566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
466b122ec5aSJacob Faibussowitsch   return 0;
467c4762a1bSJed Brown }
468c4762a1bSJed Brown 
469c4762a1bSJed Brown /*TEST
470c4762a1bSJed Brown 
471c4762a1bSJed Brown    test:
472c4762a1bSJed Brown       suffix: 0
473c4762a1bSJed Brown       nsize: 3
474c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 0
475c4762a1bSJed Brown 
476c4762a1bSJed Brown    test:
477c4762a1bSJed Brown       suffix: 1
478c4762a1bSJed Brown       nsize: 3
479c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 1
480c4762a1bSJed Brown 
481c4762a1bSJed Brown    test:
482c4762a1bSJed Brown       suffix: 2
483c4762a1bSJed Brown       nsize: 3
484c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2
485c4762a1bSJed Brown 
486c4762a1bSJed Brown    test:
487c4762a1bSJed Brown       suffix: 3
488c4762a1bSJed Brown       nsize: 3
489c4762a1bSJed 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
490c4762a1bSJed Brown 
491c4762a1bSJed Brown    test:
492c4762a1bSJed Brown       suffix: 4
493c4762a1bSJed Brown       nsize: 6
494c4762a1bSJed 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
495c4762a1bSJed Brown       requires: !single
496c4762a1bSJed Brown 
497c4762a1bSJed Brown    test:
498c4762a1bSJed Brown       suffix: glvis_composite_da_1d
499c4762a1bSJed Brown       args: -u_da_grid_x 20 -problem_type 0 -snes_monitor_solution glvis:
500c4762a1bSJed Brown 
50105ec3129SRichard Tran Mills    test:
50205ec3129SRichard Tran Mills       suffix: cuda
50305ec3129SRichard Tran Mills       nsize: 1
50405ec3129SRichard Tran Mills       requires: cuda
50505ec3129SRichard 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
50605ec3129SRichard Tran Mills 
50705ec3129SRichard Tran Mills    test:
50805ec3129SRichard Tran Mills       suffix: viennacl
50905ec3129SRichard Tran Mills       nsize: 1
51005ec3129SRichard Tran Mills       requires: viennacl
51105ec3129SRichard 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
51205ec3129SRichard Tran Mills 
513c4762a1bSJed Brown TEST*/
514