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