xref: /petsc/src/dm/tutorials/ex22.c (revision 030f984af8d8bb4c203755d35bded3c05b3d83ce)
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   ierr = PetscInitialize(&argc, &argv, (char*)0, help);if (ierr) return ierr;
56   ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRMPI(ierr);
57   ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRMPI(ierr);
58 
59   ierr = PetscOptionsBegin(PETSC_COMM_WORLD, "", "ex22 DMDA tutorial example options", "DMDA");CHKERRQ(ierr);
60   ierr = PetscOptionsRangeInt("-Mx", "dimension along x-axis", "ex22.c", Mx, &Mx, NULL, 0, PETSC_MAX_INT);CHKERRQ(ierr);
61   ierr = PetscOptionsRangeInt("-My", "dimension along y-axis", "ex22.c", My, &My, NULL, 0, PETSC_MAX_INT);CHKERRQ(ierr);
62   ierr = PetscOptionsRangeInt("-Mz", "dimension along z-axis", "ex22.c", Mz, &Mz, NULL, 0, PETSC_MAX_INT);CHKERRQ(ierr);
63   ierr = PetscOptionsEnum("-sliceaxis","axis along which 2D slice is extracted from : X, Y, Z","",sliceaxes,(PetscEnum)sliceaxis,(PetscEnum*)&sliceaxis,NULL);CHKERRQ(ierr);
64   ierr = PetscOptionsRangeInt("-gp", "index along sliceaxis at which 2D slice is extracted", "ex22.c", gp, &gp, NULL, 0, PETSC_MAX_INT);CHKERRQ(ierr);
65   ierr = PetscOptionsEnd();CHKERRQ(ierr);
66 
67   /* Ensure that the requested slice is not out of bounds for the selected axis */
68   if (sliceaxis==DM_X) {
69     if (gp>Mx) SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
70   } else if (sliceaxis==DM_Y) {
71     if (gp>My) SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
72   } else if (sliceaxis==DM_Z) {
73     if (gp>Mz) SETERRQ(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);CHKERRQ(ierr);
87   ierr = DMSetFromOptions(da3D);CHKERRQ(ierr);
88   ierr = DMSetUp(da3D);CHKERRQ(ierr);
89 
90   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
91      Create the parent vector
92      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
93   ierr = DMCreateGlobalVector(da3D, &vec_full);CHKERRQ(ierr);
94   ierr = PetscObjectSetName((PetscObject) vec_full, "full_vector");CHKERRQ(ierr);
95 
96   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
97      Populate the 3D vector
98      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
99   ierr = DMDAGetCorners(da3D, &ixs, &iys, &izs, &ixm, &iym, &izm);CHKERRQ(ierr);
100   ierr = DMDAVecGetArray(da3D, vec_full, &vecdata3d);CHKERRQ(ierr);
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   ierr = DMDAVecRestoreArray(da3D, vec_full, &vecdata3d);CHKERRQ(ierr);
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   ierr = DMDACreatePatchIS(da3D, &lower, &upper, &patchis_3d, patchis_offproc);CHKERRQ(ierr);
124   ierr = PetscPrintf(PETSC_COMM_WORLD, "\n IS to select slice from 3D DMDA vector : \n");CHKERRQ(ierr);
125   ierr = ISView(patchis_3d, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
126 
127   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
128      Use the obtained IS to extract the slice as a subvector
129      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
130   ierr = VecGetSubVector(vec_full, patchis_3d, &vec_extracted);CHKERRQ(ierr);
131 
132   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
133      View the extracted subvector
134      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
135   ierr = PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_DENSE);CHKERRQ(ierr);
136   ierr = PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in DMDA ordering : \n");CHKERRQ(ierr);
137   ierr = VecView(vec_extracted, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
138   ierr = PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
139 
140   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
141      Query 3D DMDA layout, get the subset MPI communicator
142      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
143   if (sliceaxis==DM_X) {
144     ierr = DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL);CHKERRQ(ierr);
145     ierr = DMDAGetOwnershipRanges(da3D, NULL, &l1, &l2);CHKERRQ(ierr);
146     M1 = My; M2 = Mz;
147     ierr = DMDAGetProcessorSubset(da3D, DM_X, gp, &subset_mpi_comm);CHKERRQ(ierr);
148   } else if (sliceaxis==DM_Y) {
149     ierr = DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, NULL, &m2, NULL, NULL, NULL, NULL, NULL, NULL);CHKERRQ(ierr);
150     ierr = DMDAGetOwnershipRanges(da3D, &l1, NULL, &l2);CHKERRQ(ierr);
151     M1 = Mx; M2 = Mz;
152     ierr = DMDAGetProcessorSubset(da3D, DM_Y, gp, &subset_mpi_comm);CHKERRQ(ierr);
153   } else if (sliceaxis==DM_Z) {
154     ierr = DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL, NULL);CHKERRQ(ierr);
155     ierr = DMDAGetOwnershipRanges(da3D, &l1, &l2, NULL);CHKERRQ(ierr);
156     M1 = Mx; M2 = My;
157     ierr = DMDAGetProcessorSubset(da3D, DM_Z, gp, &subset_mpi_comm);CHKERRQ(ierr);
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     ierr = MPI_Comm_size(subset_mpi_comm, &size);CHKERRMPI(ierr);
168     ierr = PetscSynchronizedPrintf(subset_mpi_comm, "subset MPI subcomm size is : %d, includes global rank : %d \n", size, rank);CHKERRQ(ierr);
169     ierr = PetscSynchronizedFlush(subset_mpi_comm, PETSC_STDOUT);CHKERRQ(ierr);
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);CHKERRQ(ierr);
178     ierr = DMSetFromOptions(da2D);CHKERRQ(ierr);
179     ierr = DMSetUp(da2D);CHKERRQ(ierr);
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     ierr = DMDACreatePatchIS(da2D, &lower, &upper, &patchis_2d, patchis_offproc);CHKERRQ(ierr);
185 
186     /* Convert the 2D patch IS to natural indexing (column major flattened) */
187     ierr = ISDuplicate(patchis_2d, &scatis_natural_slice);CHKERRQ(ierr);
188     ierr = DMDAGetAO(da2D, &da2D_ao);CHKERRQ(ierr);
189     ierr = AOPetscToApplicationIS(da2D_ao, scatis_natural_slice);CHKERRQ(ierr);
190     ierr = ISGetIndices(scatis_natural_slice, &is_array);CHKERRQ(ierr);
191     ierr = ISGetLocalSize(scatis_natural_slice, &in);CHKERRQ(ierr);
192 
193     /* Create an aliased IS on the 3D DMDA's communicator */
194     ierr = ISCreateGeneral(PETSC_COMM_WORLD, in, is_array, PETSC_USE_POINTER, &scatis_natural_slice_g);CHKERRQ(ierr);
195     ierr = ISRestoreIndices(scatis_natural_slice, &is_array);CHKERRQ(ierr);
196 
197     /* Create a 2D DMDA global vector */
198     ierr = DMCreateGlobalVector(da2D, &vec_slice);CHKERRQ(ierr);
199     ierr = PetscObjectSetName((PetscObject) vec_slice, "slice_vector_natural");CHKERRQ(ierr);
200     ierr = VecGetLocalSize(vec_slice ,&vn);CHKERRQ(ierr);
201     ierr = VecGetArrayRead(vec_slice, &array);CHKERRQ(ierr);
202 
203     /* Create an aliased version of the above on the 3D DMDA's communicator */
204     ierr = VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, vn, M1*M2, array, &vec_slice_g);CHKERRQ(ierr);
205     ierr = VecRestoreArrayRead(vec_slice, &array);CHKERRQ(ierr);
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     ierr = ISCreateGeneral(PETSC_COMM_WORLD, 0, NULL, PETSC_USE_POINTER, &scatis_natural_slice_g);CHKERRQ(ierr);
210     ierr = VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, 0, M1*M2, NULL, &vec_slice_g);CHKERRQ(ierr);
211   }
212   ierr = PetscBarrier(NULL);CHKERRQ(ierr);
213 
214   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
215      Create IS that maps from the extracted slice vector
216      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
217   ierr = VecGetOwnershipRange(vec_extracted, &low, &high);CHKERRQ(ierr);
218   ierr = ISCreateStride(PETSC_COMM_WORLD, high-low, low, 1, &scatis_extracted_slice);CHKERRQ(ierr);
219 
220   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
221      Scatter extracted subvector -> natural 2D slice vector
222      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
223   ierr = VecScatterCreate(vec_extracted, scatis_extracted_slice, vec_slice_g, scatis_natural_slice_g, &vscat);CHKERRQ(ierr);
224   ierr = VecScatterBegin(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD);CHKERRQ(ierr);
225   ierr = VecScatterEnd(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD);CHKERRQ(ierr);
226 
227   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
228      View the natural 2D slice vector
229      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
230   ierr = PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_DENSE);CHKERRQ(ierr);
231   ierr = PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in natural ordering : \n");CHKERRQ(ierr);
232   ierr = VecView(vec_slice_g, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
233   ierr = PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
234 
235   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
236      Restore subvector
237      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
238   ierr = VecRestoreSubVector(vec_full, patchis_3d, &vec_extracted);CHKERRQ(ierr);
239 
240   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
241      Destroy data structures and exit.
242      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
243   ierr = VecDestroy(&vec_full);CHKERRQ(ierr);
244   ierr = VecScatterDestroy(&vscat);CHKERRQ(ierr);
245   ierr = ISDestroy(&scatis_extracted_slice);CHKERRQ(ierr);
246   ierr = ISDestroy(&scatis_natural_slice_g);CHKERRQ(ierr);
247   ierr = VecDestroy(&vec_slice_g);CHKERRQ(ierr);
248   ierr = ISDestroy(&patchis_3d);CHKERRQ(ierr);
249   ierr = DMDestroy(&da3D);CHKERRQ(ierr);
250 
251   if (subset_mpi_comm != MPI_COMM_NULL) {
252     ierr = ISDestroy(&patchis_2d);CHKERRQ(ierr);
253     ierr = ISDestroy(&scatis_natural_slice);CHKERRQ(ierr);
254     ierr = VecDestroy(&vec_slice);CHKERRQ(ierr);
255     ierr = DMDestroy(&da2D);CHKERRQ(ierr);
256     ierr = MPI_Comm_free(&subset_mpi_comm);CHKERRMPI(ierr);
257   }
258 
259   ierr = PetscFinalize();
260   return ierr;
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