1 /* Routines to convert between a (subset of) DMStag and DMDA */
2 
3 #include <petscdmda.h>
4 #include <petsc/private/dmstagimpl.h>
5 #include <petscdmdatypes.h>
6 
DMStagCreateCompatibleDMDA(DM dm,DMStagStencilLocation loc,PetscInt c,DM * dmda)7 static PetscErrorCode DMStagCreateCompatibleDMDA(DM dm,DMStagStencilLocation loc,PetscInt c,DM *dmda)
8 {
9   PetscErrorCode  ierr;
10   DM_Stag * const stag = (DM_Stag*) dm->data;
11   PetscInt        dim,i,j,stencilWidth,dof,N[DMSTAG_MAX_DIM];
12   DMDAStencilType stencilType;
13   PetscInt        *l[DMSTAG_MAX_DIM];
14 
15   PetscFunctionBegin;
16   PetscValidHeaderSpecificType(dm,DM_CLASSID,1,DMSTAG);
17   ierr = DMGetDimension(dm,&dim);CHKERRQ(ierr);
18 
19   /* Create grid decomposition (to be adjusted later) */
20   for (i=0; i<dim; ++i) {
21     ierr = PetscMalloc1(stag->nRanks[i],&l[i]);CHKERRQ(ierr);
22     for (j=0; j<stag->nRanks[i]; ++j) l[i][j] = stag->l[i][j];
23     N[i] = stag->N[i];
24   }
25 
26   /* dof */
27   dof = c < 0 ? -c : 1;
28 
29   /* Determine/adjust sizes */
30   switch (loc) {
31     case DMSTAG_ELEMENT:
32       break;
33     case DMSTAG_LEFT:
34     case DMSTAG_RIGHT:
35       if (dim<1) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
36       l[0][stag->nRanks[0]-1] += 1; /* extra vertex in direction 0 on last rank in dimension 0 */
37       N[0] += 1;
38       break;
39     case DMSTAG_UP:
40     case DMSTAG_DOWN:
41       if (dim < 2) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
42       l[1][stag->nRanks[1]-1] += 1; /* extra vertex in direction 1 on last rank in dimension 1 */
43       N[1] += 1;
44       break;
45     case DMSTAG_BACK:
46     case DMSTAG_FRONT:
47       if (dim < 3) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
48       l[2][stag->nRanks[2]-1] += 1; /* extra vertex in direction 2 on last rank in dimension 2 */
49       N[2] += 1;
50       break;
51     case DMSTAG_DOWN_LEFT :
52     case DMSTAG_DOWN_RIGHT :
53     case DMSTAG_UP_LEFT :
54     case DMSTAG_UP_RIGHT :
55       if (dim < 2) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
56       for (i=0; i<2; ++i) { /* extra vertex in direction i on last rank in dimension i = 0,1 */
57         l[i][stag->nRanks[i]-1] += 1;
58         N[i] += 1;
59       }
60       break;
61     case DMSTAG_BACK_LEFT:
62     case DMSTAG_BACK_RIGHT:
63     case DMSTAG_FRONT_LEFT:
64     case DMSTAG_FRONT_RIGHT:
65       if (dim < 3) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
66       for (i=0; i<3; i+=2) { /* extra vertex in direction i on last rank in dimension i = 0,2 */
67         l[i][stag->nRanks[i]-1] += 1;
68         N[i] += 1;
69       }
70       break;
71     case DMSTAG_BACK_DOWN:
72     case DMSTAG_BACK_UP:
73     case DMSTAG_FRONT_DOWN:
74     case DMSTAG_FRONT_UP:
75       if (dim < 3) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
76       for (i=1; i<3; ++i) { /* extra vertex in direction i on last rank in dimension i = 1,2 */
77         l[i][stag->nRanks[i]-1] += 1;
78         N[i] += 1;
79       }
80       break;
81     case DMSTAG_BACK_DOWN_LEFT:
82     case DMSTAG_BACK_DOWN_RIGHT:
83     case DMSTAG_BACK_UP_LEFT:
84     case DMSTAG_BACK_UP_RIGHT:
85     case DMSTAG_FRONT_DOWN_LEFT:
86     case DMSTAG_FRONT_DOWN_RIGHT:
87     case DMSTAG_FRONT_UP_LEFT:
88     case DMSTAG_FRONT_UP_RIGHT:
89       if (dim < 3) SETERRQ2(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Incompatible dim (%d) and loc(%s) combination",dim,DMStagStencilLocations[loc]);
90       for (i=0; i<3; ++i) { /* extra vertex in direction i on last rank in dimension i = 0,1,2 */
91         l[i][stag->nRanks[i]-1] += 1;
92         N[i] += 1;
93       }
94       break;
95       break;
96     default : SETERRQ1(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Not implemented for location %s",DMStagStencilLocations[loc]);
97   }
98 
99   /* Use the same stencil type */
100   switch (stag->stencilType) {
101     case DMSTAG_STENCIL_STAR: stencilType = DMDA_STENCIL_STAR; stencilWidth = stag->stencilWidth; break;
102     case DMSTAG_STENCIL_BOX : stencilType = DMDA_STENCIL_BOX ; stencilWidth = stag->stencilWidth; break;
103     default: SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Unsupported Stencil Type %d",stag->stencilType);
104   }
105 
106   /* Create DMDA, using same boundary type */
107   switch (dim) {
108     case 1:
109       ierr = DMDACreate1d(PetscObjectComm((PetscObject)dm),stag->boundaryType[0],N[0],dof,stencilWidth,l[0],dmda);CHKERRQ(ierr);
110       break;
111     case 2:
112       ierr = DMDACreate2d(PetscObjectComm((PetscObject)dm),stag->boundaryType[0],stag->boundaryType[1],stencilType,N[0],N[1],stag->nRanks[0],stag->nRanks[1],dof,stencilWidth,l[0],l[1],dmda);CHKERRQ(ierr);
113       break;
114     case 3:
115       ierr = DMDACreate3d(PetscObjectComm((PetscObject)dm),stag->boundaryType[0],stag->boundaryType[1],stag->boundaryType[2],stencilType,N[0],N[1],N[2],stag->nRanks[0],stag->nRanks[1],stag->nRanks[2],dof,stencilWidth,l[0],l[1],l[2],dmda);CHKERRQ(ierr);
116       break;
117     default: SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not implemented for dim %d",dim);
118   }
119   for (i=0; i<dim; ++i) {
120     ierr = PetscFree(l[i]);CHKERRQ(ierr);
121   }
122   PetscFunctionReturn(0);
123 }
124 
125 /*
126 Helper function to get the number of extra points in a DMDA representation for a given canonical location.
127 */
DMStagDMDAGetExtraPoints(DM dm,DMStagStencilLocation locCanonical,PetscInt * extraPoint)128 static PetscErrorCode DMStagDMDAGetExtraPoints(DM dm,DMStagStencilLocation locCanonical,PetscInt *extraPoint)
129 {
130   PetscErrorCode ierr;
131   PetscInt       dim,d,nExtra[DMSTAG_MAX_DIM];
132 
133   PetscFunctionBegin;
134   PetscValidHeaderSpecificType(dm,DM_CLASSID,1,DMSTAG);
135   ierr = DMGetDimension(dm,&dim);CHKERRQ(ierr);
136   if (dim > DMSTAG_MAX_DIM) SETERRQ1(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Not implemented for %D dimensions",dim);
137   ierr = DMStagGetCorners(dm,NULL,NULL,NULL,NULL,NULL,NULL,&nExtra[0],&nExtra[1],&nExtra[2]);CHKERRQ(ierr);
138   for (d=0; d<dim; ++d) extraPoint[d] = 0;
139   switch (locCanonical) {
140     case DMSTAG_ELEMENT:
141       break; /* no extra points */
142     case DMSTAG_LEFT:
143       extraPoint[0] = nExtra[0]; break; /* only extra point in x */
144     case DMSTAG_DOWN:
145       extraPoint[1] = nExtra[1]; break; /* only extra point in y */
146     case DMSTAG_BACK:
147       extraPoint[2] = nExtra[2]; break; /* only extra point in z */
148     case DMSTAG_DOWN_LEFT:
149       extraPoint[0] = nExtra[0]; extraPoint[1] = nExtra[1]; break; /* extra point in both x and y  */
150     case DMSTAG_BACK_LEFT:
151       extraPoint[0] = nExtra[0]; extraPoint[2] = nExtra[2]; break; /* extra point in both x and z  */
152     case DMSTAG_BACK_DOWN:
153       extraPoint[1] = nExtra[1]; extraPoint[2] = nExtra[2]; break; /* extra point in both y and z  */
154     case DMSTAG_BACK_DOWN_LEFT:
155      extraPoint[0] = nExtra[0]; extraPoint[1] = nExtra[1]; extraPoint[2] = nExtra[2]; break; /* extra points in x,y,z */
156     default : SETERRQ1(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Not implemented for location (perhaps not canonical) %s",DMStagStencilLocations[locCanonical]);
157   }
158   PetscFunctionReturn(0);
159 }
160 
161 /*
162 Function much like DMStagMigrateVec(), but which accepts an additional position argument to disambiguate which
163 type of DMDA to migrate to.
164 */
165 
DMStagMigrateVecDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM dmTo,Vec vecTo)166 static PetscErrorCode DMStagMigrateVecDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM dmTo,Vec vecTo)
167 {
168   PetscErrorCode ierr;
169   PetscInt       i,j,k,d,dim,dof,dofToMax,start[DMSTAG_MAX_DIM],n[DMSTAG_MAX_DIM],extraPoint[DMSTAG_MAX_DIM];
170   Vec            vecLocal;
171 
172   PetscFunctionBegin;
173   PetscValidHeaderSpecificType(dm,DM_CLASSID,1,DMSTAG);
174   PetscValidHeaderSpecific(vec,VEC_CLASSID,2);
175   PetscValidHeaderSpecificType(dmTo,DM_CLASSID,4,DMDA);
176   PetscValidHeaderSpecific(vecTo,VEC_CLASSID,5);
177   ierr = DMGetDimension(dm,&dim);CHKERRQ(ierr);
178   ierr = DMDAGetDof(dmTo,&dofToMax);CHKERRQ(ierr);
179   if (-c > dofToMax) SETERRQ1(PetscObjectComm((PetscObject)dmTo),PETSC_ERR_ARG_OUTOFRANGE,"Invalid negative component value. Must be >= -%D",dofToMax);
180   ierr = DMStagGetCorners(dm,&start[0],&start[1],&start[2],&n[0],&n[1],&n[2],NULL,NULL,NULL);CHKERRQ(ierr);
181   ierr = DMStagDMDAGetExtraPoints(dm,loc,extraPoint);CHKERRQ(ierr);
182   ierr = DMStagGetLocationDOF(dm,loc,&dof);CHKERRQ(ierr);
183   ierr = DMGetLocalVector(dm,&vecLocal);CHKERRQ(ierr);
184   ierr = DMGlobalToLocalBegin(dm,vec,INSERT_VALUES,vecLocal);CHKERRQ(ierr);
185   ierr = DMGlobalToLocalEnd(dm,vec,INSERT_VALUES,vecLocal);CHKERRQ(ierr);
186   if (dim == 1) {
187     PetscScalar **arrTo;
188     ierr = DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
189     if (c < 0) {
190       const PetscInt dofTo = -c;
191       for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
192         for (d=0; d<PetscMin(dof,dofTo); ++d) {
193           DMStagStencil pos;
194           pos.i = i; pos.loc = loc; pos.c = d;
195           ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[i][d]);CHKERRQ(ierr);
196         }
197         for (;d<dofTo; ++d) {
198           arrTo[i][d] = 0.0; /* Pad extra dof with zeros */
199         }
200       }
201     } else {
202       for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
203         DMStagStencil pos;
204         pos.i = i; pos.loc = loc; pos.c = c;
205         ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[i][0]);CHKERRQ(ierr);
206       }
207     }
208     ierr = DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
209   } else if (dim == 2) {
210     PetscScalar ***arrTo;
211     ierr = DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
212     if (c < 0) {
213       const PetscInt dofTo = -c;
214       for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
215         for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
216           for (d=0; d<PetscMin(dof,dofTo); ++d) {
217             DMStagStencil pos;
218             pos.i = i; pos.j = j; pos.loc = loc; pos.c = d;
219             ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[j][i][d]);CHKERRQ(ierr);
220           }
221           for (;d<dofTo; ++d) {
222             arrTo[j][i][d] = 0.0; /* Pad extra dof with zeros */
223           }
224         }
225       }
226     } else {
227       for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
228         for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
229           DMStagStencil pos;
230           pos.i = i; pos.j = j; pos.loc = loc; pos.c = c;
231           ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[j][i][0]);CHKERRQ(ierr);
232         }
233       }
234     }
235     ierr = DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
236   } else if (dim == 3) {
237     PetscScalar ****arrTo;
238     ierr = DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
239     if (c < 0) {
240       const PetscInt dofTo = -c;
241       for (k=start[2]; k<start[2] + n[2] + extraPoint[2]; ++k) {
242         for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
243           for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
244             for (d=0; d<PetscMin(dof,dofTo); ++d) {
245               DMStagStencil pos;
246               pos.i = i; pos.j = j; pos.k = k; pos.loc = loc; pos.c = d;
247               ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[k][j][i][d]);CHKERRQ(ierr);
248             }
249             for (;d<dofTo; ++d) {
250               arrTo[k][j][i][d] = 0.0; /* Pad extra dof with zeros */
251             }
252           }
253         }
254       }
255     } else {
256       for (k=start[2]; k<start[2] + n[2] + extraPoint[2]; ++k) {
257         for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
258           for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
259             DMStagStencil pos;
260             pos.i = i; pos.j = j; pos.k = k; pos.loc = loc; pos.c = c;
261             ierr = DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[k][j][i][0]);CHKERRQ(ierr);
262           }
263         }
264       }
265     }
266     ierr = DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);CHKERRQ(ierr);
267   } else SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Unsupported dimension %d",dim);
268   ierr = DMRestoreLocalVector(dm,&vecLocal);CHKERRQ(ierr);
269   PetscFunctionReturn(0);
270 }
271 
272 /* Transfer coordinates from a DMStag to a DMDA, specifying which location */
DMStagTransferCoordinatesToDMDA(DM dmstag,DMStagStencilLocation loc,DM dmda)273 static PetscErrorCode DMStagTransferCoordinatesToDMDA(DM dmstag,DMStagStencilLocation loc,DM dmda)
274 {
275   PetscErrorCode ierr;
276   PetscInt       dim,start[DMSTAG_MAX_DIM],n[DMSTAG_MAX_DIM],extraPoint[DMSTAG_MAX_DIM],d;
277   DM             dmstagCoord,dmdaCoord;
278   DMType         dmstagCoordType;
279   Vec            stagCoord,daCoord;
280   PetscBool      daCoordIsStag,daCoordIsProduct;
281 
282   PetscFunctionBegin;
283   PetscValidHeaderSpecificType(dmstag,DM_CLASSID,1,DMSTAG);
284   PetscValidHeaderSpecificType(dmda,DM_CLASSID,3,DMDA);
285   ierr = DMGetDimension(dmstag,&dim);CHKERRQ(ierr);
286   ierr = DMGetCoordinateDM(dmstag,&dmstagCoord);CHKERRQ(ierr);
287   ierr = DMGetCoordinatesLocal(dmstag,&stagCoord);CHKERRQ(ierr); /* Note local */
288   ierr = DMGetCoordinateDM(dmda,&dmdaCoord);CHKERRQ(ierr);
289   daCoord = NULL;
290   ierr = DMGetCoordinates(dmda,&daCoord);CHKERRQ(ierr);
291   if (!daCoord) {
292     ierr = DMCreateGlobalVector(dmdaCoord,&daCoord);CHKERRQ(ierr);
293     ierr = DMSetCoordinates(dmda,daCoord);CHKERRQ(ierr);
294     ierr = VecDestroy(&daCoord);CHKERRQ(ierr);
295     ierr = DMGetCoordinates(dmda,&daCoord);CHKERRQ(ierr);
296   }
297   ierr = DMGetType(dmstagCoord,&dmstagCoordType);CHKERRQ(ierr);
298   ierr = PetscStrcmp(dmstagCoordType,DMSTAG,&daCoordIsStag);CHKERRQ(ierr);
299   ierr = PetscStrcmp(dmstagCoordType,DMPRODUCT,&daCoordIsProduct);CHKERRQ(ierr);
300   ierr = DMStagGetCorners(dmstag,&start[0],&start[1],&start[2],&n[0],&n[1],&n[2],NULL,NULL,NULL);CHKERRQ(ierr);
301   ierr = DMStagDMDAGetExtraPoints(dmstag,loc,extraPoint);CHKERRQ(ierr);
302   if (dim == 1) {
303     PetscInt ex;
304     PetscScalar **cArrDa;
305     ierr = DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
306     if (daCoordIsStag)  {
307       PetscInt slot;
308       PetscScalar **cArrStag;
309       ierr = DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);CHKERRQ(ierr);
310       ierr = DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
311       for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
312         cArrDa[ex][0] = cArrStag[ex][slot];
313       }
314       ierr = DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
315     } else if (daCoordIsProduct) {
316       PetscScalar **cArrX;
317       ierr = DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,NULL,NULL);CHKERRQ(ierr);
318       for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
319         cArrDa[ex][0] = cArrX[ex][0];
320       }
321       ierr = DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,NULL,NULL);CHKERRQ(ierr);
322     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
323     ierr = DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
324   } else if (dim == 2) {
325     PetscInt ex,ey;
326     PetscScalar ***cArrDa;
327     ierr = DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
328     if (daCoordIsStag)  {
329       PetscInt slot;
330       PetscScalar ***cArrStag;
331       ierr = DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);CHKERRQ(ierr);
332       ierr = DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
333       for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
334         for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
335           for (d=0; d<2; ++d) {
336             cArrDa[ey][ex][d] = cArrStag[ey][ex][slot+d];
337           }
338         }
339       }
340       ierr = DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
341     } else if (daCoordIsProduct) {
342       PetscScalar **cArrX,**cArrY;
343       ierr = DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,NULL);CHKERRQ(ierr);
344       for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
345         for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
346           cArrDa[ey][ex][0] = cArrX[ex][0];
347           cArrDa[ey][ex][1] = cArrY[ey][0];
348         }
349       }
350       ierr = DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,NULL);CHKERRQ(ierr);
351     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
352     ierr = DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
353   }  else if (dim == 3) {
354     PetscInt ex,ey,ez;
355     PetscScalar ****cArrDa;
356     ierr = DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
357     if (daCoordIsStag)  {
358       PetscInt slot;
359       PetscScalar ****cArrStag;
360       ierr = DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);CHKERRQ(ierr);
361       ierr = DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
362       for (ez=start[2]; ez<start[2] + n[2] + extraPoint[2]; ++ez) {
363         for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
364           for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
365             for (d=0; d<3; ++d) {
366               cArrDa[ez][ey][ex][d] = cArrStag[ez][ey][ex][slot+d];
367             }
368           }
369         }
370       }
371       ierr = DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);CHKERRQ(ierr);
372     } else if (daCoordIsProduct) {
373       PetscScalar **cArrX,**cArrY,**cArrZ;
374       ierr = DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,&cArrZ);CHKERRQ(ierr);
375       for (ez=start[2]; ez<start[2] + n[2] + extraPoint[2]; ++ez) {
376         for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
377           for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
378             cArrDa[ez][ey][ex][0] = cArrX[ex][0];
379             cArrDa[ez][ey][ex][1] = cArrY[ey][0];
380             cArrDa[ez][ey][ex][2] = cArrZ[ez][0];
381           }
382         }
383       }
384       ierr = DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,&cArrZ);CHKERRQ(ierr);
385     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
386     ierr = DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);CHKERRQ(ierr);
387   } else SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Unsupported dimension %d",dim);
388   PetscFunctionReturn(0);
389 }
390 
391 /*
392 Convert to a location value with only BACK, DOWN, LEFT, and ELEMENT involved (makes looping easier)
393 */
DMStagStencilLocationCanonicalize(DMStagStencilLocation loc,DMStagStencilLocation * locCanonical)394 static PetscErrorCode DMStagStencilLocationCanonicalize(DMStagStencilLocation loc,DMStagStencilLocation *locCanonical)
395 {
396   PetscFunctionBegin;
397   switch (loc) {
398     case DMSTAG_ELEMENT:
399       *locCanonical = DMSTAG_ELEMENT;
400       break;
401     case DMSTAG_LEFT:
402     case DMSTAG_RIGHT:
403       *locCanonical = DMSTAG_LEFT;
404       break;
405     case DMSTAG_DOWN:
406     case DMSTAG_UP:
407       *locCanonical = DMSTAG_DOWN;
408       break;
409     case DMSTAG_BACK:
410     case DMSTAG_FRONT:
411       *locCanonical = DMSTAG_BACK;
412       break;
413     case DMSTAG_DOWN_LEFT :
414     case DMSTAG_DOWN_RIGHT :
415     case DMSTAG_UP_LEFT :
416     case DMSTAG_UP_RIGHT :
417       *locCanonical = DMSTAG_DOWN_LEFT;
418       break;
419     case DMSTAG_BACK_LEFT:
420     case DMSTAG_BACK_RIGHT:
421     case DMSTAG_FRONT_LEFT:
422     case DMSTAG_FRONT_RIGHT:
423       *locCanonical = DMSTAG_BACK_LEFT;
424       break;
425     case DMSTAG_BACK_DOWN:
426     case DMSTAG_BACK_UP:
427     case DMSTAG_FRONT_DOWN:
428     case DMSTAG_FRONT_UP:
429       *locCanonical = DMSTAG_BACK_DOWN;
430       break;
431     case DMSTAG_BACK_DOWN_LEFT:
432     case DMSTAG_BACK_DOWN_RIGHT:
433     case DMSTAG_BACK_UP_LEFT:
434     case DMSTAG_BACK_UP_RIGHT:
435     case DMSTAG_FRONT_DOWN_LEFT:
436     case DMSTAG_FRONT_DOWN_RIGHT:
437     case DMSTAG_FRONT_UP_LEFT:
438     case DMSTAG_FRONT_UP_RIGHT:
439       *locCanonical = DMSTAG_BACK_DOWN_LEFT;
440       break;
441     default :
442       *locCanonical = DMSTAG_NULL_LOCATION;
443       break;
444   }
445   PetscFunctionReturn(0);
446 }
447 
448 /*@C
449   DMStagVecSplitToDMDA - create a DMDA and Vec from a DMStag and Vec
450 
451   Logically Collective
452 
453   High-level helper function which accepts a DMStag, a global vector, and location/dof,
454   and generates a corresponding DMDA and Vec.
455 
456   Input Parameters:
457 + dm - the DMStag object
458 . vec- Vec object associated with dm
459 . loc - which subgrid to extract (see DMStagStencilLocation)
460 - c - which component to extract (see note below)
461 
462   Output Parameters:
463 + pda - the new DMDA
464 - pdavec - the new Vec
465 
466   Notes:
467   If a c value of -k is provided, the first k dof for that position are extracted,
468   padding with zero values if needbe. If a non-negative value is provided, a single
469   dof is extracted.
470 
471   The caller is responsible for destroying the created DMDA and Vec.
472 
473   Level: advanced
474 
475 .seealso: DMSTAG, DMDA, DMStagMigrateVec(), DMStagCreateCompatibleDMStag()
476 @*/
DMStagVecSplitToDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM * pda,Vec * pdavec)477 PetscErrorCode DMStagVecSplitToDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM *pda,Vec *pdavec)
478 {
479   PetscErrorCode  ierr;
480   PetscInt        dim,locdof;
481   DM              da,coordDM;
482   Vec             davec;
483   DMStagStencilLocation locCanonical;
484 
485   PetscFunctionBegin;
486   PetscValidHeaderSpecificType(dm,DM_CLASSID,1,DMSTAG);
487   PetscValidHeaderSpecific(vec,VEC_CLASSID,2);
488   ierr = DMGetDimension(dm,&dim);CHKERRQ(ierr);
489   ierr = DMStagGetLocationDOF(dm,loc,&locdof);CHKERRQ(ierr);
490   if (c >= locdof) SETERRQ3(PetscObjectComm((PetscObject)dm),PETSC_ERR_ARG_OUTOFRANGE,"Location %s has %D dof, but component %D requested\n",DMStagStencilLocations[loc],locdof,c);
491   ierr = DMStagStencilLocationCanonicalize(loc,&locCanonical);CHKERRQ(ierr);
492   ierr = DMStagCreateCompatibleDMDA(dm,locCanonical,c,pda);CHKERRQ(ierr);
493   da = *pda;
494   ierr = DMSetUp(*pda);CHKERRQ(ierr);
495   ierr = DMGetCoordinateDM(dm,&coordDM);CHKERRQ(ierr);
496   if (coordDM) {
497     ierr = DMStagTransferCoordinatesToDMDA(dm,locCanonical,da);CHKERRQ(ierr);
498   }
499   ierr = DMCreateGlobalVector(da,pdavec);CHKERRQ(ierr);
500   davec = *pdavec;
501   ierr = DMStagMigrateVecDMDA(dm,vec,locCanonical,c,da,davec);CHKERRQ(ierr);
502   PetscFunctionReturn(0);
503 }
504