xref: /petsc/src/snes/tutorials/ex28.c (revision b122ec5aa1bd4469eb4e0673542fb7de3f411254)
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;
795f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetEntries(user->pack,&dau,&dak));
805f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetLocalInfo(dau,&infou));
815f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetLocalInfo(dak,&infok));
825f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
83c4762a1bSJed Brown   switch (user->ptype) {
84c4762a1bSJed Brown   case 0:
855f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
865f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
875f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,Uloc,&u));
885f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,user->Kloc,&k));
895f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,F,&fu));
905f80ce2aSJacob Faibussowitsch     CHKERRQ(FormFunctionLocal_U(user,&infou,u,k,fu));
915f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,F,&fu));
925f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,Uloc,&u));
935f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,user->Kloc,&k));
94c4762a1bSJed Brown     break;
95c4762a1bSJed Brown   case 1:
965f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
975f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
985f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,user->Uloc,&u));
995f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,Kloc,&k));
1005f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,F,&fk));
1015f80ce2aSJacob Faibussowitsch     CHKERRQ(FormFunctionLocal_K(user,&infok,u,k,fk));
1025f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,F,&fk));
1035f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,user->Uloc,&u));
1045f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,Kloc,&k));
105c4762a1bSJed Brown     break;
106c4762a1bSJed Brown   case 2:
1075f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeScatter(user->pack,X,Uloc,Kloc));
1085f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,Uloc,&u));
1095f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,Kloc,&k));
1105f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetAccess(user->pack,F,&Fu,&Fk));
1115f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,Fu,&fu));
1125f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,Fk,&fk));
1135f80ce2aSJacob Faibussowitsch     CHKERRQ(FormFunctionLocal_U(user,&infou,u,k,fu));
1145f80ce2aSJacob Faibussowitsch     CHKERRQ(FormFunctionLocal_K(user,&infok,u,k,fk));
1155f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,Fu,&fu));
1165f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,Fk,&fk));
1175f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeRestoreAccess(user->pack,F,&Fu,&Fk));
1185f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,Uloc,&u));
1195f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,Kloc,&k));
120c4762a1bSJed Brown     break;
121c4762a1bSJed Brown   }
1225f80ce2aSJacob Faibussowitsch   CHKERRQ(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) {
1365f80ce2aSJacob Faibussowitsch       CHKERRQ(MatSetValuesLocal(Buu,1,&row,1,&row,&val,INSERT_VALUES));
137c4762a1bSJed Brown     } else if (i == info->mx-1) {
1385f80ce2aSJacob Faibussowitsch       CHKERRQ(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};
1415f80ce2aSJacob Faibussowitsch       CHKERRQ(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.)};
1565f80ce2aSJacob Faibussowitsch     CHKERRQ(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;
1755f80ce2aSJacob Faibussowitsch     CHKERRQ(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);
2075f80ce2aSJacob Faibussowitsch     CHKERRQ(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;
2215f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetEntries(user->pack,&dau,&dak));
2225f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetLocalInfo(dau,&infou));
2235f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetLocalInfo(dak,&infok));
2245f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetLocalVectors(user->pack,&Uloc,&Kloc));
225c4762a1bSJed Brown   switch (user->ptype) {
226c4762a1bSJed Brown   case 0:
2275f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalBegin(dau,X,INSERT_VALUES,Uloc));
2285f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalEnd  (dau,X,INSERT_VALUES,Uloc));
2295f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,Uloc,&u));
2305f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,user->Kloc,&k));
2315f80ce2aSJacob Faibussowitsch     CHKERRQ(FormJacobianLocal_U(user,&infou,u,k,B));
2325f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,Uloc,&u));
2335f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,user->Kloc,&k));
234c4762a1bSJed Brown     break;
235c4762a1bSJed Brown   case 1:
2365f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalBegin(dak,X,INSERT_VALUES,Kloc));
2375f80ce2aSJacob Faibussowitsch     CHKERRQ(DMGlobalToLocalEnd  (dak,X,INSERT_VALUES,Kloc));
2385f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,user->Uloc,&u));
2395f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,Kloc,&k));
2405f80ce2aSJacob Faibussowitsch     CHKERRQ(FormJacobianLocal_K(user,&infok,u,k,B));
2415f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,user->Uloc,&u));
2425f80ce2aSJacob Faibussowitsch     CHKERRQ(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;
2485f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeScatter(user->pack,X,Uloc,Kloc));
2495f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dau,Uloc,&u));
2505f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecGetArray(dak,Kloc,&k));
2515f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetLocalISs(user->pack,&is));
2525f80ce2aSJacob Faibussowitsch     CHKERRQ(MatGetLocalSubMatrix(B,is[0],is[0],&Buu));
2535f80ce2aSJacob Faibussowitsch     CHKERRQ(MatGetLocalSubMatrix(B,is[0],is[1],&Buk));
2545f80ce2aSJacob Faibussowitsch     CHKERRQ(MatGetLocalSubMatrix(B,is[1],is[0],&Bku));
2555f80ce2aSJacob Faibussowitsch     CHKERRQ(MatGetLocalSubMatrix(B,is[1],is[1],&Bkk));
2565f80ce2aSJacob Faibussowitsch     CHKERRQ(FormJacobianLocal_U(user,&infou,u,k,Buu));
2575f80ce2aSJacob Faibussowitsch     CHKERRQ(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       */
2655f80ce2aSJacob Faibussowitsch       CHKERRQ(FormJacobianLocal_UK(user,&infou,&infok,u,k,Buk));
2665f80ce2aSJacob Faibussowitsch       CHKERRQ(FormJacobianLocal_KU(user,&infou,&infok,u,k,Bku));
267c4762a1bSJed Brown     }
2685f80ce2aSJacob Faibussowitsch     CHKERRQ(FormJacobianLocal_K(user,&infok,u,k,Bkk));
2695f80ce2aSJacob Faibussowitsch     CHKERRQ(MatRestoreLocalSubMatrix(B,is[0],is[0],&Buu));
2705f80ce2aSJacob Faibussowitsch     CHKERRQ(MatRestoreLocalSubMatrix(B,is[0],is[1],&Buk));
2715f80ce2aSJacob Faibussowitsch     CHKERRQ(MatRestoreLocalSubMatrix(B,is[1],is[0],&Bku));
2725f80ce2aSJacob Faibussowitsch     CHKERRQ(MatRestoreLocalSubMatrix(B,is[1],is[1],&Bkk));
2735f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dau,Uloc,&u));
2745f80ce2aSJacob Faibussowitsch     CHKERRQ(DMDAVecRestoreArray(dak,Kloc,&k));
275c4762a1bSJed Brown 
2765f80ce2aSJacob Faibussowitsch     CHKERRQ(ISDestroy(&is[0]));
2775f80ce2aSJacob Faibussowitsch     CHKERRQ(ISDestroy(&is[1]));
2785f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscFree(is));
279c4762a1bSJed Brown   } break;
280c4762a1bSJed Brown   }
2815f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeRestoreLocalVectors(user->pack,&Uloc,&Kloc));
2825f80ce2aSJacob Faibussowitsch   CHKERRQ(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
2835f80ce2aSJacob Faibussowitsch   CHKERRQ(MatAssemblyEnd  (B,MAT_FINAL_ASSEMBLY));
284c4762a1bSJed Brown   if (J != B) {
2855f80ce2aSJacob Faibussowitsch     CHKERRQ(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
2865f80ce2aSJacob Faibussowitsch     CHKERRQ(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;
3005f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetEntries(user->pack,&dau,&dak));
3015f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetAccess(user->pack,X,&Xu,&Xk));
3025f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAVecGetArray(dau,Xu,&u));
3035f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAVecGetArray(dak,Xk,&k));
3045f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetLocalInfo(dau,&infou));
3055f80ce2aSJacob Faibussowitsch   CHKERRQ(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);
3095f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAVecRestoreArray(dau,Xu,&u));
3105f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAVecRestoreArray(dak,Xk,&k));
3115f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeRestoreAccess(user->pack,X,&Xu,&Xk));
3125f80ce2aSJacob Faibussowitsch   CHKERRQ(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*b122ec5aSJacob Faibussowitsch   CHKERRQ(PetscInitialize(&argc,&argv,0,help));
3305f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,10,1,1,NULL,&dau));
3315f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetOptionsPrefix(dau,"u_"));
3325f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetFromOptions(dau));
3335f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetUp(dau));
3345f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetOwnershipRanges(dau,&lxu,0,0));
3355f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDAGetInfo(dau,0, &m,0,0, &sizes,0,0, 0,0,0,0,0,0));
3365f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscMalloc1(sizes,&lxk));
3375f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscArraycpy(lxk,lxu,sizes));
338c4762a1bSJed Brown   lxk[0]--;
3395f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,m-1,1,1,lxk,&dak));
3405f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetOptionsPrefix(dak,"k_"));
3415f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetFromOptions(dak));
3425f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetUp(dak));
3435f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscFree(lxk));
344c4762a1bSJed Brown 
3455f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeCreate(PETSC_COMM_WORLD,&pack));
3465f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetOptionsPrefix(pack,"pack_"));
3475f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeAddDM(pack,dau));
3485f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeAddDM(pack,dak));
3495f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDASetFieldName(dau,0,"u"));
3505f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDASetFieldName(dak,0,"k"));
3515f80ce2aSJacob Faibussowitsch   CHKERRQ(DMSetFromOptions(pack));
352c4762a1bSJed Brown 
3535f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCreateGlobalVector(pack,&X));
3545f80ce2aSJacob Faibussowitsch   CHKERRQ(VecDuplicate(X,&F));
355c4762a1bSJed Brown 
3565f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscNew(&user));
357c4762a1bSJed Brown 
358c4762a1bSJed Brown   user->pack = pack;
359c4762a1bSJed Brown 
3605f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetGlobalISs(pack,&isg));
3615f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeGetLocalVectors(pack,&user->Uloc,&user->Kloc));
3625f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeScatter(pack,X,user->Uloc,user->Kloc));
363c4762a1bSJed Brown 
364c4762a1bSJed Brown   ierr = PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"Coupled problem options","SNES");CHKERRQ(ierr);
365c4762a1bSJed Brown   {
366c4762a1bSJed Brown     user->ptype = 0; pass_dm = PETSC_TRUE;
367c4762a1bSJed Brown 
3685f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscOptionsInt("-problem_type","0: solve for u only, 1: solve for k only, 2: solve for both",0,user->ptype,&user->ptype,NULL));
3695f80ce2aSJacob Faibussowitsch     CHKERRQ(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   }
371c4762a1bSJed Brown   ierr = PetscOptionsEnd();CHKERRQ(ierr);
372c4762a1bSJed Brown 
3735f80ce2aSJacob Faibussowitsch   CHKERRQ(FormInitial_Coupled(user,X));
374c4762a1bSJed Brown 
3755f80ce2aSJacob Faibussowitsch   CHKERRQ(SNESCreate(PETSC_COMM_WORLD,&snes));
376c4762a1bSJed Brown   switch (user->ptype) {
377c4762a1bSJed Brown   case 0:
3785f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetAccess(pack,X,&Xu,0));
3795f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetAccess(pack,F,&Fu,0));
3805f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCreateMatrix(dau,&B));
3815f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFunction(snes,Fu,FormFunction_All,user));
3825f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
3835f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFromOptions(snes));
3845f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetDM(snes,dau));
3855f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSolve(snes,NULL,Xu));
3865f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeRestoreAccess(pack,X,&Xu,0));
3875f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeRestoreAccess(pack,F,&Fu,0));
388c4762a1bSJed Brown     break;
389c4762a1bSJed Brown   case 1:
3905f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetAccess(pack,X,0,&Xk));
3915f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeGetAccess(pack,F,0,&Fk));
3925f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCreateMatrix(dak,&B));
3935f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFunction(snes,Fk,FormFunction_All,user));
3945f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
3955f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFromOptions(snes));
3965f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetDM(snes,dak));
3975f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSolve(snes,NULL,Xk));
3985f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeRestoreAccess(pack,X,0,&Xk));
3995f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCompositeRestoreAccess(pack,F,0,&Fk));
400c4762a1bSJed Brown     break;
401c4762a1bSJed Brown   case 2:
4025f80ce2aSJacob Faibussowitsch     CHKERRQ(DMCreateMatrix(pack,&B));
403c4762a1bSJed Brown     /* This example does not correctly allocate off-diagonal blocks. These options allows new nonzeros (slow). */
4045f80ce2aSJacob Faibussowitsch     CHKERRQ(MatSetOption(B,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_FALSE));
4055f80ce2aSJacob Faibussowitsch     CHKERRQ(MatSetOption(B,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_FALSE));
4065f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFunction(snes,F,FormFunction_All,user));
4075f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetJacobian(snes,B,B,FormJacobian_All,user));
4085f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSetFromOptions(snes));
409c4762a1bSJed Brown     if (!pass_dm) {             /* Manually provide index sets and names for the splits */
410c4762a1bSJed Brown       KSP ksp;
411c4762a1bSJed Brown       PC  pc;
4125f80ce2aSJacob Faibussowitsch       CHKERRQ(SNESGetKSP(snes,&ksp));
4135f80ce2aSJacob Faibussowitsch       CHKERRQ(KSPGetPC(ksp,&pc));
4145f80ce2aSJacob Faibussowitsch       CHKERRQ(PCFieldSplitSetIS(pc,"u",isg[0]));
4155f80ce2aSJacob Faibussowitsch       CHKERRQ(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). */
4195f80ce2aSJacob Faibussowitsch       CHKERRQ(SNESSetDM(snes,pack));
420c4762a1bSJed Brown     }
4215f80ce2aSJacob Faibussowitsch     CHKERRQ(SNESSolve(snes,NULL,X));
422c4762a1bSJed Brown     break;
423c4762a1bSJed Brown   }
4245f80ce2aSJacob Faibussowitsch   CHKERRQ(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 
4325f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscOptionsGetInt(NULL,0,"-col",&col,0));
4335f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscOptionsGetBool(NULL,0,"-mult_dup",&mult_dup,0));
4345f80ce2aSJacob Faibussowitsch     CHKERRQ(PetscOptionsGetBool(NULL,0,"-view_dup",&view_dup,0));
435c4762a1bSJed Brown 
4365f80ce2aSJacob Faibussowitsch     CHKERRQ(VecDuplicate(X,&Y));
4375f80ce2aSJacob Faibussowitsch     /* CHKERRQ(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); */
4385f80ce2aSJacob Faibussowitsch     /* CHKERRQ(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); */
4395f80ce2aSJacob Faibussowitsch     CHKERRQ(MatConvert(B,MATAIJ,MAT_INITIAL_MATRIX,&D));
4405f80ce2aSJacob Faibussowitsch     CHKERRQ(VecZeroEntries(X));
4415f80ce2aSJacob Faibussowitsch     CHKERRQ(VecSetValue(X,col,1.0,INSERT_VALUES));
4425f80ce2aSJacob Faibussowitsch     CHKERRQ(VecAssemblyBegin(X));
4435f80ce2aSJacob Faibussowitsch     CHKERRQ(VecAssemblyEnd(X));
4445f80ce2aSJacob Faibussowitsch     CHKERRQ(MatMult(mult_dup ? D : B,X,Y));
4455f80ce2aSJacob Faibussowitsch     CHKERRQ(MatView(view_dup ? D : B,PETSC_VIEWER_STDOUT_WORLD));
4465f80ce2aSJacob Faibussowitsch     /* CHKERRQ(VecView(X,PETSC_VIEWER_STDOUT_WORLD)); */
4475f80ce2aSJacob Faibussowitsch     CHKERRQ(VecView(Y,PETSC_VIEWER_STDOUT_WORLD));
4485f80ce2aSJacob Faibussowitsch     CHKERRQ(MatDestroy(&D));
4495f80ce2aSJacob Faibussowitsch     CHKERRQ(VecDestroy(&Y));
450c4762a1bSJed Brown   }
451c4762a1bSJed Brown 
4525f80ce2aSJacob Faibussowitsch   CHKERRQ(DMCompositeRestoreLocalVectors(pack,&user->Uloc,&user->Kloc));
4535f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscFree(user));
454c4762a1bSJed Brown 
4555f80ce2aSJacob Faibussowitsch   CHKERRQ(ISDestroy(&isg[0]));
4565f80ce2aSJacob Faibussowitsch   CHKERRQ(ISDestroy(&isg[1]));
4575f80ce2aSJacob Faibussowitsch   CHKERRQ(PetscFree(isg));
4585f80ce2aSJacob Faibussowitsch   CHKERRQ(VecDestroy(&X));
4595f80ce2aSJacob Faibussowitsch   CHKERRQ(VecDestroy(&F));
4605f80ce2aSJacob Faibussowitsch   CHKERRQ(MatDestroy(&B));
4615f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDestroy(&dau));
4625f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDestroy(&dak));
4635f80ce2aSJacob Faibussowitsch   CHKERRQ(DMDestroy(&pack));
4645f80ce2aSJacob Faibussowitsch   CHKERRQ(SNESDestroy(&snes));
465*b122ec5aSJacob Faibussowitsch   CHKERRQ(PetscFinalize());
466*b122ec5aSJacob 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