xref: /petsc/src/snes/tutorials/ex28.c (revision d0609ced746bc51b019815ca91d747429db24893)
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;
799566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
809566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
819566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak,&infok));
829566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
83c4762a1bSJed Brown   switch (user->ptype) {
84c4762a1bSJed Brown   case 0:
859566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
869566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
879566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
889566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,user->Kloc,&k));
899566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,F,&fu));
909566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user,&infou,u,k,fu));
919566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,F,&fu));
929566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
939566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,user->Kloc,&k));
94c4762a1bSJed Brown     break;
95c4762a1bSJed Brown   case 1:
969566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
979566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
989566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,user->Uloc,&u));
999566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
1009566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,F,&fk));
1019566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user,&infok,u,k,fk));
1029566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,F,&fk));
1039566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,user->Uloc,&u));
1049566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
105c4762a1bSJed Brown     break;
106c4762a1bSJed Brown   case 2:
1079566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack,X,Uloc,Kloc));
1089566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
1099566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
1109566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(user->pack,F,&Fu,&Fk));
1119566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Fu,&fu));
1129566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Fk,&fk));
1139566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_U(user,&infou,u,k,fu));
1149566063dSJacob Faibussowitsch     PetscCall(FormFunctionLocal_K(user,&infok,u,k,fk));
1159566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Fu,&fu));
1169566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Fk,&fk));
1179566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(user->pack,F,&Fu,&Fk));
1189566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
1199566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
120c4762a1bSJed Brown     break;
121c4762a1bSJed Brown   }
1229566063dSJacob 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) {
1369566063dSJacob Faibussowitsch       PetscCall(MatSetValuesLocal(Buu,1,&row,1,&row,&val,INSERT_VALUES));
137c4762a1bSJed Brown     } else if (i == info->mx-1) {
1389566063dSJacob 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};
1419566063dSJacob 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.)};
1569566063dSJacob 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;
1759566063dSJacob 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);
2079566063dSJacob 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;
2219566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
2229566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
2239566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dak,&infok));
2249566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
225c4762a1bSJed Brown   switch (user->ptype) {
226c4762a1bSJed Brown   case 0:
2279566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
2289566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
2299566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
2309566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,user->Kloc,&k));
2319566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user,&infou,u,k,B));
2329566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
2339566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,user->Kloc,&k));
234c4762a1bSJed Brown     break;
235c4762a1bSJed Brown   case 1:
2369566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
2379566063dSJacob Faibussowitsch     PetscCall(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
2389566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,user->Uloc,&u));
2399566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
2409566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user,&infok,u,k,B));
2419566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,user->Uloc,&u));
2429566063dSJacob 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;
2489566063dSJacob Faibussowitsch     PetscCall(DMCompositeScatter(user->pack,X,Uloc,Kloc));
2499566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dau,Uloc,&u));
2509566063dSJacob Faibussowitsch     PetscCall(DMDAVecGetArray(dak,Kloc,&k));
2519566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetLocalISs(user->pack,&is));
2529566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[0],is[0],&Buu));
2539566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[0],is[1],&Buk));
2549566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[1],is[0],&Bku));
2559566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSubMatrix(B,is[1],is[1],&Bkk));
2569566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_U(user,&infou,u,k,Buu));
2579566063dSJacob 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       */
2659566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_UK(user,&infou,&infok,u,k,Buk));
2669566063dSJacob Faibussowitsch       PetscCall(FormJacobianLocal_KU(user,&infou,&infok,u,k,Bku));
267c4762a1bSJed Brown     }
2689566063dSJacob Faibussowitsch     PetscCall(FormJacobianLocal_K(user,&infok,u,k,Bkk));
2699566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[0],is[0],&Buu));
2709566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[0],is[1],&Buk));
2719566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[1],is[0],&Bku));
2729566063dSJacob Faibussowitsch     PetscCall(MatRestoreLocalSubMatrix(B,is[1],is[1],&Bkk));
2739566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dau,Uloc,&u));
2749566063dSJacob Faibussowitsch     PetscCall(DMDAVecRestoreArray(dak,Kloc,&k));
275c4762a1bSJed Brown 
2769566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[0]));
2779566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&is[1]));
2789566063dSJacob Faibussowitsch     PetscCall(PetscFree(is));
279c4762a1bSJed Brown   } break;
280c4762a1bSJed Brown   }
2819566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(user->pack,&Uloc,&Kloc));
2829566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
2839566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd  (B,MAT_FINAL_ASSEMBLY));
284c4762a1bSJed Brown   if (J != B) {
2859566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
2869566063dSJacob 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;
3009566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetEntries(user->pack,&dau,&dak));
3019566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetAccess(user->pack,X,&Xu,&Xk));
3029566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dau,Xu,&u));
3039566063dSJacob Faibussowitsch   PetscCall(DMDAVecGetArray(dak,Xk,&k));
3049566063dSJacob Faibussowitsch   PetscCall(DMDAGetLocalInfo(dau,&infou));
3059566063dSJacob 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);
3099566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dau,Xu,&u));
3109566063dSJacob Faibussowitsch   PetscCall(DMDAVecRestoreArray(dak,Xk,&k));
3119566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreAccess(user->pack,X,&Xu,&Xk));
3129566063dSJacob 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   DM             dau,dak,pack;
319c4762a1bSJed Brown   const PetscInt *lxu;
320c4762a1bSJed Brown   PetscInt       *lxk,m,sizes;
321c4762a1bSJed Brown   User           user;
322c4762a1bSJed Brown   SNES           snes;
323c4762a1bSJed Brown   Vec            X,F,Xu,Xk,Fu,Fk;
324c4762a1bSJed Brown   Mat            B;
325c4762a1bSJed Brown   IS             *isg;
326c4762a1bSJed Brown   PetscBool      pass_dm;
327c4762a1bSJed Brown 
3289566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc,&argv,0,help));
3299566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,10,1,1,NULL,&dau));
3309566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dau,"u_"));
3319566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dau));
3329566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dau));
3339566063dSJacob Faibussowitsch   PetscCall(DMDAGetOwnershipRanges(dau,&lxu,0,0));
3349566063dSJacob Faibussowitsch   PetscCall(DMDAGetInfo(dau,0, &m,0,0, &sizes,0,0, 0,0,0,0,0,0));
3359566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(sizes,&lxk));
3369566063dSJacob Faibussowitsch   PetscCall(PetscArraycpy(lxk,lxu,sizes));
337c4762a1bSJed Brown   lxk[0]--;
3389566063dSJacob Faibussowitsch   PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,m-1,1,1,lxk,&dak));
3399566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(dak,"k_"));
3409566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(dak));
3419566063dSJacob Faibussowitsch   PetscCall(DMSetUp(dak));
3429566063dSJacob Faibussowitsch   PetscCall(PetscFree(lxk));
343c4762a1bSJed Brown 
3449566063dSJacob Faibussowitsch   PetscCall(DMCompositeCreate(PETSC_COMM_WORLD,&pack));
3459566063dSJacob Faibussowitsch   PetscCall(DMSetOptionsPrefix(pack,"pack_"));
3469566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack,dau));
3479566063dSJacob Faibussowitsch   PetscCall(DMCompositeAddDM(pack,dak));
3489566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dau,0,"u"));
3499566063dSJacob Faibussowitsch   PetscCall(DMDASetFieldName(dak,0,"k"));
3509566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(pack));
351c4762a1bSJed Brown 
3529566063dSJacob Faibussowitsch   PetscCall(DMCreateGlobalVector(pack,&X));
3539566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(X,&F));
354c4762a1bSJed Brown 
3559566063dSJacob Faibussowitsch   PetscCall(PetscNew(&user));
356c4762a1bSJed Brown 
357c4762a1bSJed Brown   user->pack = pack;
358c4762a1bSJed Brown 
3599566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetGlobalISs(pack,&isg));
3609566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetLocalVectors(pack,&user->Uloc,&user->Kloc));
3619566063dSJacob Faibussowitsch   PetscCall(DMCompositeScatter(pack,X,user->Uloc,user->Kloc));
362c4762a1bSJed Brown 
363*d0609cedSBarry Smith   PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"Coupled problem options","SNES");
364c4762a1bSJed Brown   {
365c4762a1bSJed Brown     user->ptype = 0; pass_dm = PETSC_TRUE;
366c4762a1bSJed Brown 
3679566063dSJacob 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));
3689566063dSJacob 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));
369c4762a1bSJed Brown   }
370*d0609cedSBarry Smith   PetscOptionsEnd();
371c4762a1bSJed Brown 
3729566063dSJacob Faibussowitsch   PetscCall(FormInitial_Coupled(user,X));
373c4762a1bSJed Brown 
3749566063dSJacob Faibussowitsch   PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes));
375c4762a1bSJed Brown   switch (user->ptype) {
376c4762a1bSJed Brown   case 0:
3779566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,X,&Xu,0));
3789566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,F,&Fu,0));
3799566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dau,&B));
3809566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,Fu,FormFunction_All,user));
3819566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
3829566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3839566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes,dau));
3849566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,Xu));
3859566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,X,&Xu,0));
3869566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,F,&Fu,0));
387c4762a1bSJed Brown     break;
388c4762a1bSJed Brown   case 1:
3899566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,X,0,&Xk));
3909566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccess(pack,F,0,&Fk));
3919566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(dak,&B));
3929566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,Fk,FormFunction_All,user));
3939566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
3949566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
3959566063dSJacob Faibussowitsch     PetscCall(SNESSetDM(snes,dak));
3969566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,Xk));
3979566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,X,0,&Xk));
3989566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccess(pack,F,0,&Fk));
399c4762a1bSJed Brown     break;
400c4762a1bSJed Brown   case 2:
4019566063dSJacob Faibussowitsch     PetscCall(DMCreateMatrix(pack,&B));
402c4762a1bSJed Brown     /* This example does not correctly allocate off-diagonal blocks. These options allows new nonzeros (slow). */
4039566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_FALSE));
4049566063dSJacob Faibussowitsch     PetscCall(MatSetOption(B,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_FALSE));
4059566063dSJacob Faibussowitsch     PetscCall(SNESSetFunction(snes,F,FormFunction_All,user));
4069566063dSJacob Faibussowitsch     PetscCall(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
4079566063dSJacob Faibussowitsch     PetscCall(SNESSetFromOptions(snes));
408c4762a1bSJed Brown     if (!pass_dm) {             /* Manually provide index sets and names for the splits */
409c4762a1bSJed Brown       KSP ksp;
410c4762a1bSJed Brown       PC  pc;
4119566063dSJacob Faibussowitsch       PetscCall(SNESGetKSP(snes,&ksp));
4129566063dSJacob Faibussowitsch       PetscCall(KSPGetPC(ksp,&pc));
4139566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc,"u",isg[0]));
4149566063dSJacob Faibussowitsch       PetscCall(PCFieldSplitSetIS(pc,"k",isg[1]));
415c4762a1bSJed Brown     } else {
416c4762a1bSJed Brown       /* The same names come from the options prefix for dau and dak. This option can support geometric multigrid inside
417c4762a1bSJed Brown        * of splits, but it requires using a DM (perhaps your own implementation). */
4189566063dSJacob Faibussowitsch       PetscCall(SNESSetDM(snes,pack));
419c4762a1bSJed Brown     }
4209566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes,NULL,X));
421c4762a1bSJed Brown     break;
422c4762a1bSJed Brown   }
4239566063dSJacob Faibussowitsch   PetscCall(VecViewFromOptions(X,NULL,"-view_sol"));
424c4762a1bSJed Brown 
425c4762a1bSJed Brown   if (0) {
426c4762a1bSJed Brown     PetscInt  col      = 0;
427c4762a1bSJed Brown     PetscBool mult_dup = PETSC_FALSE,view_dup = PETSC_FALSE;
428c4762a1bSJed Brown     Mat       D;
429c4762a1bSJed Brown     Vec       Y;
430c4762a1bSJed Brown 
4319566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetInt(NULL,0,"-col",&col,0));
4329566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL,0,"-mult_dup",&mult_dup,0));
4339566063dSJacob Faibussowitsch     PetscCall(PetscOptionsGetBool(NULL,0,"-view_dup",&view_dup,0));
434c4762a1bSJed Brown 
4359566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X,&Y));
4369566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); */
4379566063dSJacob Faibussowitsch     /* PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); */
4389566063dSJacob Faibussowitsch     PetscCall(MatConvert(B,MATAIJ,MAT_INITIAL_MATRIX,&D));
4399566063dSJacob Faibussowitsch     PetscCall(VecZeroEntries(X));
4409566063dSJacob Faibussowitsch     PetscCall(VecSetValue(X,col,1.0,INSERT_VALUES));
4419566063dSJacob Faibussowitsch     PetscCall(VecAssemblyBegin(X));
4429566063dSJacob Faibussowitsch     PetscCall(VecAssemblyEnd(X));
4439566063dSJacob Faibussowitsch     PetscCall(MatMult(mult_dup ? D : B,X,Y));
4449566063dSJacob Faibussowitsch     PetscCall(MatView(view_dup ? D : B,PETSC_VIEWER_STDOUT_WORLD));
4459566063dSJacob Faibussowitsch     /* PetscCall(VecView(X,PETSC_VIEWER_STDOUT_WORLD)); */
4469566063dSJacob Faibussowitsch     PetscCall(VecView(Y,PETSC_VIEWER_STDOUT_WORLD));
4479566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&D));
4489566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&Y));
449c4762a1bSJed Brown   }
450c4762a1bSJed Brown 
4519566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreLocalVectors(pack,&user->Uloc,&user->Kloc));
4529566063dSJacob Faibussowitsch   PetscCall(PetscFree(user));
453c4762a1bSJed Brown 
4549566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[0]));
4559566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&isg[1]));
4569566063dSJacob Faibussowitsch   PetscCall(PetscFree(isg));
4579566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X));
4589566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&F));
4599566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
4609566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dau));
4619566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dak));
4629566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&pack));
4639566063dSJacob Faibussowitsch   PetscCall(SNESDestroy(&snes));
4649566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
465b122ec5aSJacob Faibussowitsch   return 0;
466c4762a1bSJed Brown }
467c4762a1bSJed Brown 
468c4762a1bSJed Brown /*TEST
469c4762a1bSJed Brown 
470c4762a1bSJed Brown    test:
471c4762a1bSJed Brown       suffix: 0
472c4762a1bSJed Brown       nsize: 3
473c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 0
474c4762a1bSJed Brown 
475c4762a1bSJed Brown    test:
476c4762a1bSJed Brown       suffix: 1
477c4762a1bSJed Brown       nsize: 3
478c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 1
479c4762a1bSJed Brown 
480c4762a1bSJed Brown    test:
481c4762a1bSJed Brown       suffix: 2
482c4762a1bSJed Brown       nsize: 3
483c4762a1bSJed Brown       args: -u_da_grid_x 20 -snes_converged_reason -snes_monitor_short -problem_type 2
484c4762a1bSJed Brown 
485c4762a1bSJed Brown    test:
486c4762a1bSJed Brown       suffix: 3
487c4762a1bSJed Brown       nsize: 3
488c4762a1bSJed 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
489c4762a1bSJed Brown 
490c4762a1bSJed Brown    test:
491c4762a1bSJed Brown       suffix: 4
492c4762a1bSJed Brown       nsize: 6
493c4762a1bSJed 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
494c4762a1bSJed Brown       requires: !single
495c4762a1bSJed Brown 
496c4762a1bSJed Brown    test:
497c4762a1bSJed Brown       suffix: glvis_composite_da_1d
498c4762a1bSJed Brown       args: -u_da_grid_x 20 -problem_type 0 -snes_monitor_solution glvis:
499c4762a1bSJed Brown 
50005ec3129SRichard Tran Mills    test:
50105ec3129SRichard Tran Mills       suffix: cuda
50205ec3129SRichard Tran Mills       nsize: 1
50305ec3129SRichard Tran Mills       requires: cuda
50405ec3129SRichard 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
50505ec3129SRichard Tran Mills 
50605ec3129SRichard Tran Mills    test:
50705ec3129SRichard Tran Mills       suffix: viennacl
50805ec3129SRichard Tran Mills       nsize: 1
50905ec3129SRichard Tran Mills       requires: viennacl
51005ec3129SRichard 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
51105ec3129SRichard Tran Mills 
512c4762a1bSJed Brown TEST*/
513