xref: /petsc/src/dm/tutorials/ex22.c (revision b33f4bec9907f62d08679bf5e7ff704a731f8c0f)
1 static char help[] ="Extract a 2D slice in natural ordering from a 3D vector, Command line options : \n\
2 Mx/My/Mz - set the dimensions of the parent vector \n\
3 sliceaxis - string describing the axis along which the sice will be selected : X, Y, Z \n\
4 gp - global grid point number along the sliceaxis direction where the slice will be extracted from the parent vector \n";
5 
6 /*
7   This example shows to extract a 2D slice in natural ordering
8   from a 3D DMDA vector (first by extracting the slice and then
9   by converting it to natural ordering)
10 */
11 
12 #include <petscdmda.h>
13 
14 const char *const sliceaxes[] = {"X","Y","Z","sliceaxis","DM_",NULL};
15 
16 int main(int argc,char **argv)
17 {
18   DM                da3D;                          /* 3D DMDA object */
19   DM                da2D;                          /* 2D DMDA object */
20   Vec               vec_full;                      /* Parent vector */
21   Vec               vec_extracted;                 /* Extracted slice vector (in DMDA ordering) */
22   Vec               vec_slice;                     /* vector in natural ordering */
23   Vec               vec_slice_g;                   /* aliased vector in natural ordering */
24   IS                patchis_3d;                    /* IS to select slice and extract subvector */
25   IS                patchis_2d;                    /* Patch IS for 2D vector, will be converted to application ordering */
26   IS                scatis_extracted_slice;        /* PETSc indexed IS for extracted slice */
27   IS                scatis_natural_slice;          /* natural/application ordered IS for slice*/
28   IS                scatis_natural_slice_g;        /* aliased natural/application ordered IS  for slice */
29   VecScatter        vscat;                         /* scatter slice in DMDA ordering <-> slice in column major ordering */
30   AO                da2D_ao;                       /* AO associated with 2D DMDA */
31   MPI_Comm          subset_mpi_comm=MPI_COMM_NULL; /* MPI communicator where the slice lives */
32   PetscScalar       ***vecdata3d;                  /* Pointer to access 3d parent vector */
33   const PetscScalar *array;                        /* pointer to create aliased Vec */
34   PetscInt          Mx=4,My=4,Mz=4;                /* Dimensions for 3D DMDA */
35   const PetscInt    *l1,*l2;                       /* 3D DMDA layout */
36   PetscInt          M1=-1,M2=-1;                   /* Dimensions for 2D DMDA */
37   PetscInt          m1=-1,m2=-1;                   /* Layouts for 2D DMDA */
38   PetscInt          gp=2;                          /* grid point along sliceaxis to pick the slice */
39   DMDirection       sliceaxis=DM_X;                /* Select axis along which the slice will be extracted */
40   PetscInt          i,j,k;                         /* Iteration indices */
41   PetscInt          ixs,iys,izs;                   /* Corner indices for 3D vector */
42   PetscInt          ixm,iym,izm;                   /* Widths of parent vector */
43   PetscInt          low, high;                     /* ownership range indices */
44   PetscInt          in;                            /* local size index for IS*/
45   PetscInt          vn;                            /* local size index */
46   const PetscInt    *is_array;                     /* pointer to create aliased IS */
47   MatStencil        lower, upper;                  /* Stencils to select slice for Vec */
48   PetscBool         patchis_offproc = PETSC_FALSE; /* flag to DMDACreatePatchIS indicating that off-proc values are to be ignored */
49   PetscMPIInt       rank,size;                     /* MPI rank and size */
50   PetscErrorCode    ierr;                          /* error checking */
51 
52   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
53      Initialize program and set problem parameters
54      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
55   PetscCall(PetscInitialize(&argc, &argv, (char*)0, help));
56   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
57   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&rank));
58 
59   ierr = PetscOptionsBegin(PETSC_COMM_WORLD, "", "ex22 DMDA tutorial example options", "DMDA");PetscCall(ierr);
60   PetscCall(PetscOptionsRangeInt("-Mx", "dimension along x-axis", "ex22.c", Mx, &Mx, NULL, 0, PETSC_MAX_INT));
61   PetscCall(PetscOptionsRangeInt("-My", "dimension along y-axis", "ex22.c", My, &My, NULL, 0, PETSC_MAX_INT));
62   PetscCall(PetscOptionsRangeInt("-Mz", "dimension along z-axis", "ex22.c", Mz, &Mz, NULL, 0, PETSC_MAX_INT));
63   PetscCall(PetscOptionsEnum("-sliceaxis","axis along which 2D slice is extracted from : X, Y, Z","",sliceaxes,(PetscEnum)sliceaxis,(PetscEnum*)&sliceaxis,NULL));
64   PetscCall(PetscOptionsRangeInt("-gp", "index along sliceaxis at which 2D slice is extracted", "ex22.c", gp, &gp, NULL, 0, PETSC_MAX_INT));
65   ierr = PetscOptionsEnd();PetscCall(ierr);
66 
67   /* Ensure that the requested slice is not out of bounds for the selected axis */
68   if (sliceaxis==DM_X) {
69     PetscCheck(gp<=Mx,PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
70   } else if (sliceaxis==DM_Y) {
71     PetscCheck(gp<=My,PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
72   } else if (sliceaxis==DM_Z) {
73     PetscCheck(gp<=Mz,PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
74   }
75 
76   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
77      Create 3D DMDA object.
78      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
79   ierr = DMDACreate3d(PETSC_COMM_WORLD,
80                       DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE,
81                       DMDA_STENCIL_STAR,
82                       Mx, My, Mz,
83                       PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE,
84                       1, 1,
85                       NULL, NULL, NULL,
86                       &da3D);PetscCall(ierr);
87   PetscCall(DMSetFromOptions(da3D));
88   PetscCall(DMSetUp(da3D));
89 
90   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
91      Create the parent vector
92      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
93   PetscCall(DMCreateGlobalVector(da3D, &vec_full));
94   PetscCall(PetscObjectSetName((PetscObject) vec_full, "full_vector"));
95 
96   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
97      Populate the 3D vector
98      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
99   PetscCall(DMDAGetCorners(da3D, &ixs, &iys, &izs, &ixm, &iym, &izm));
100   PetscCall(DMDAVecGetArray(da3D, vec_full, &vecdata3d));
101   for (k=izs; k<izs+izm; k++) {
102     for (j=iys; j<iys+iym; j++) {
103       for (i=ixs; i<ixs+ixm; i++) {
104         vecdata3d[k][j][i] = ((i-Mx/2.0)*(j+Mx/2.0))+k*100;
105       }
106     }
107   }
108   PetscCall(DMDAVecRestoreArray(da3D, vec_full, &vecdata3d));
109 
110   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
111      Get an IS corresponding to a 2D slice
112      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
113   if (sliceaxis==DM_X) {
114     lower.i = gp; lower.j = 0;  lower.k = 0;
115     upper.i = gp; upper.j = My; upper.k = Mz;
116   } else if (sliceaxis==DM_Y) {
117     lower.i = 0;  lower.j = gp; lower.k = 0;
118     upper.i = Mx; upper.j = gp; upper.k = Mz;
119   } else if (sliceaxis==DM_Z) {
120     lower.i = 0;  lower.j = 0;  lower.k = gp;
121     upper.i = Mx; upper.j = My; upper.k = gp;
122   }
123   PetscCall(DMDACreatePatchIS(da3D, &lower, &upper, &patchis_3d, patchis_offproc));
124   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n IS to select slice from 3D DMDA vector : \n"));
125   PetscCall(ISView(patchis_3d, PETSC_VIEWER_STDOUT_WORLD));
126 
127   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
128      Use the obtained IS to extract the slice as a subvector
129      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
130   PetscCall(VecGetSubVector(vec_full, patchis_3d, &vec_extracted));
131 
132   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
133      View the extracted subvector
134      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
135   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_DENSE));
136   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in DMDA ordering : \n"));
137   PetscCall(VecView(vec_extracted, PETSC_VIEWER_STDOUT_WORLD));
138   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
139 
140   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
141      Query 3D DMDA layout, get the subset MPI communicator
142      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
143   if (sliceaxis==DM_X) {
144     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
145     PetscCall(DMDAGetOwnershipRanges(da3D, NULL, &l1, &l2));
146     M1 = My; M2 = Mz;
147     PetscCall(DMDAGetProcessorSubset(da3D, DM_X, gp, &subset_mpi_comm));
148   } else if (sliceaxis==DM_Y) {
149     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, NULL, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
150     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, NULL, &l2));
151     M1 = Mx; M2 = Mz;
152     PetscCall(DMDAGetProcessorSubset(da3D, DM_Y, gp, &subset_mpi_comm));
153   } else if (sliceaxis==DM_Z) {
154     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
155     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, &l2, NULL));
156     M1 = Mx; M2 = My;
157     PetscCall(DMDAGetProcessorSubset(da3D, DM_Z, gp, &subset_mpi_comm));
158   }
159 
160   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
161      Create 2D DMDA object,
162      vector (that will hold the slice as a column major flattened array) &
163      index set (that will be used for scattering to the column major
164      indexed slice vector)
165      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
166   if (subset_mpi_comm != MPI_COMM_NULL) {
167     PetscCallMPI(MPI_Comm_size(subset_mpi_comm, &size));
168     PetscCall(PetscSynchronizedPrintf(subset_mpi_comm, "subset MPI subcomm size is : %d, includes global rank : %d \n", size, rank));
169     PetscCall(PetscSynchronizedFlush(subset_mpi_comm, PETSC_STDOUT));
170     ierr = DMDACreate2d(subset_mpi_comm,
171                         DM_BOUNDARY_NONE, DM_BOUNDARY_NONE,
172                         DMDA_STENCIL_STAR,
173                         M1, M2,
174                         m1, m2,
175                         1, 1,
176                         l1, l2,
177                         &da2D);PetscCall(ierr);
178     PetscCall(DMSetFromOptions(da2D));
179     PetscCall(DMSetUp(da2D));
180 
181     /* Create a 2D patch IS for the slice */
182     lower.i = 0;  lower.j = 0;
183     upper.i = M1; upper.j = M2;
184     PetscCall(DMDACreatePatchIS(da2D, &lower, &upper, &patchis_2d, patchis_offproc));
185 
186     /* Convert the 2D patch IS to natural indexing (column major flattened) */
187     PetscCall(ISDuplicate(patchis_2d, &scatis_natural_slice));
188     PetscCall(DMDAGetAO(da2D, &da2D_ao));
189     PetscCall(AOPetscToApplicationIS(da2D_ao, scatis_natural_slice));
190     PetscCall(ISGetIndices(scatis_natural_slice, &is_array));
191     PetscCall(ISGetLocalSize(scatis_natural_slice, &in));
192 
193     /* Create an aliased IS on the 3D DMDA's communicator */
194     PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, in, is_array, PETSC_USE_POINTER, &scatis_natural_slice_g));
195     PetscCall(ISRestoreIndices(scatis_natural_slice, &is_array));
196 
197     /* Create a 2D DMDA global vector */
198     PetscCall(DMCreateGlobalVector(da2D, &vec_slice));
199     PetscCall(PetscObjectSetName((PetscObject) vec_slice, "slice_vector_natural"));
200     PetscCall(VecGetLocalSize(vec_slice ,&vn));
201     PetscCall(VecGetArrayRead(vec_slice, &array));
202 
203     /* Create an aliased version of the above on the 3D DMDA's communicator */
204     PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, vn, M1*M2, array, &vec_slice_g));
205     PetscCall(VecRestoreArrayRead(vec_slice, &array));
206   } else {
207     /* Ranks not part of the subset MPI communicator provide no entries, but the routines for creating
208        the IS and Vec on the 3D DMDA's communicator still need to called, since they are collective routines */
209     PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 0, NULL, PETSC_USE_POINTER, &scatis_natural_slice_g));
210     PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, 0, M1*M2, NULL, &vec_slice_g));
211   }
212   PetscCall(PetscBarrier(NULL));
213 
214   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
215      Create IS that maps from the extracted slice vector
216      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
217   PetscCall(VecGetOwnershipRange(vec_extracted, &low, &high));
218   PetscCall(ISCreateStride(PETSC_COMM_WORLD, high-low, low, 1, &scatis_extracted_slice));
219 
220   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
221      Scatter extracted subvector -> natural 2D slice vector
222      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
223   PetscCall(VecScatterCreate(vec_extracted, scatis_extracted_slice, vec_slice_g, scatis_natural_slice_g, &vscat));
224   PetscCall(VecScatterBegin(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));
225   PetscCall(VecScatterEnd(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));
226 
227   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
228      View the natural 2D slice vector
229      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
230   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_DENSE));
231   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in natural ordering : \n"));
232   PetscCall(VecView(vec_slice_g, PETSC_VIEWER_STDOUT_WORLD));
233   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
234 
235   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
236      Restore subvector
237      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
238   PetscCall(VecRestoreSubVector(vec_full, patchis_3d, &vec_extracted));
239 
240   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
241      Destroy data structures and exit.
242      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
243   PetscCall(VecDestroy(&vec_full));
244   PetscCall(VecScatterDestroy(&vscat));
245   PetscCall(ISDestroy(&scatis_extracted_slice));
246   PetscCall(ISDestroy(&scatis_natural_slice_g));
247   PetscCall(VecDestroy(&vec_slice_g));
248   PetscCall(ISDestroy(&patchis_3d));
249   PetscCall(DMDestroy(&da3D));
250 
251   if (subset_mpi_comm != MPI_COMM_NULL) {
252     PetscCall(ISDestroy(&patchis_2d));
253     PetscCall(ISDestroy(&scatis_natural_slice));
254     PetscCall(VecDestroy(&vec_slice));
255     PetscCall(DMDestroy(&da2D));
256     PetscCallMPI(MPI_Comm_free(&subset_mpi_comm));
257   }
258 
259   PetscCall(PetscFinalize());
260   return 0;
261 }
262 
263 /*TEST
264 
265     test:
266       nsize: 1
267       args: -sliceaxis X -gp 0
268 
269     test:
270       suffix: 2
271       nsize:  2
272       args: -sliceaxis Y -gp 1
273       filter: grep -v "subset MPI subcomm"
274 
275     test:
276       suffix: 3
277       nsize:  3
278       args:  -sliceaxis Z -gp 2
279       filter: grep -v "subset MPI subcomm"
280 
281     test:
282       suffix: 4
283       nsize:  4
284       args: -sliceaxis X -gp 2
285       filter: grep -v "subset MPI subcomm"
286 
287     test:
288       suffix: 5
289       nsize:  4
290       args: -sliceaxis Z -gp 1
291       filter: grep -v "subset MPI subcomm"
292 
293 TEST*/
294